From 8f5e8d8fe7de5b0c6c5b461981e8dfec17dfaa50 Mon Sep 17 00:00:00 2001 From: MathisWellmann Date: Tue, 28 Jul 2026 00:01:40 +0200 Subject: [PATCH] add example `cuda-softmax` to show GPU kernel evolution --- Cargo.lock | 19 + Cargo.toml | 1 + README.md | 26 +- examples/cuda-softmax/Cargo.toml | 31 + examples/cuda-softmax/README.md | 174 +++++ examples/cuda-softmax/src/gpu.rs | 858 +++++++++++++++++++++++ examples/cuda-softmax/src/isolate.rs | 383 ++++++++++ examples/cuda-softmax/src/lib.rs | 177 +++++ examples/cuda-softmax/src/main.rs | 529 ++++++++++++++ examples/cuda-softmax/tests/isolation.rs | 75 ++ 10 files changed, 2267 insertions(+), 6 deletions(-) create mode 100644 examples/cuda-softmax/Cargo.toml create mode 100644 examples/cuda-softmax/README.md create mode 100644 examples/cuda-softmax/src/gpu.rs create mode 100644 examples/cuda-softmax/src/isolate.rs create mode 100644 examples/cuda-softmax/src/lib.rs create mode 100644 examples/cuda-softmax/src/main.rs create mode 100644 examples/cuda-softmax/tests/isolation.rs diff --git a/Cargo.lock b/Cargo.lock index 6971a8a..f903f96 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1509,6 +1509,25 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "cuda-softmax-example" +version = "0.22.1" +dependencies = [ + "cudarc", + "symbiont", + "tokio", + "tracing", +] + +[[package]] +name = "cudarc" +version = "0.19.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42310153e06cf4cd532901f7096beb27504d681736a29ee90728ae4e2d93b2a8" +dependencies = [ + "libloading 0.9.0", +] + [[package]] name = "cursor-icon" version = "1.2.0" diff --git a/Cargo.toml b/Cargo.toml index 91a2a27..7e200be 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ exclude = ["assets", "website"] members = [ "examples/batched-evolution", "examples/counter", + "examples/cuda-softmax", "examples/evolving-trader", "examples/fizzbuzz", "examples/fractal-studio", diff --git a/README.md b/README.md index d26e848..d0f5a45 100644 --- a/README.md +++ b/README.md @@ -120,6 +120,26 @@ Type a prompt — *"an animated Julia set, c orbiting the main cardioid, with a cargo run -p fractal-studio-example --release ``` +## Showcase: evolving a CUDA kernel + +The [cuda-softmax-example](examples/cuda-softmax/README.md) puts the loop on a GPU: the agent implements +`fn plan(rows, cols) -> KernelPlan`, returning **CUDA C source plus the launch geometry**, and the host +compiles it with NVRTC, checks it against a CPU oracle, and times it. + +- The fitness function is not a matter of taste: a row-wise softmax must move exactly one read and one write + of the matrix, and the ceiling is *measured* — a copy kernel over the same buffers on the same card — so + every candidate scores as a percentage of something real. +- Measured on an RTX PRO 6000 Blackwell, one round of two lanes against a local Qwen3.6-35B-A3B took the + naive starting kernel from 523 us to **12.7 us (41x, 47% of the copy ceiling)**. +- A GPU kernel cannot be `catch_unwind`-ed: an illegal access is a *sticky* CUDA error, after which no + context can be created in that process at all. So candidates are compiled and timed in a **child process** — + the process boundary is the GPU's panic handler — and a kernel that faults the device costs one child and + one table row instead of the run. + +```bash +MODEL=... cargo run -p cuda-softmax-example --release +``` + ## Core highlights - **Type-safe agentic code**: @@ -266,12 +286,6 @@ These constraints arise from the binary/dylib interaction boundary. The harness The LLM can only rewrite function *bodies* — the signature declared in `evolvable!` is fixed at compile time and enforced on every evolution. This is by design (it's what makes constrained generation possible), but it means the agent cannot add parameters, change return types, or introduce new functions at runtime. It would be UB to hot-swap a different function signature in, when the main binary expects a certain memory layout. -- **Sequential feedback loop**: - All evolvable function calls must have returned before `evolve()` or `activate_revision()` is called. - Retained revisions are never unmapped, so a violating in-flight call executes stale but still-mapped code rather than UB — the contract remains so a swap cannot publish a torn set of pointers from two different revisions. - This matches the intended usage pattern (run functions, collect results, evolve, repeat) and is enforced with an assertion in debug builds at zero cost in release. - Calls through `RevisionFn` handles are exempt: they pin their revision and never read the swapped pointers, so they may run concurrently with `evolve()` / `activate_revision()` and with each other. - Multi-threading is possible, but requires extra care. - **Same toolchain required**: Rust has no stable ABI. The binary and dylib must be compiled with the same `rustc` version to guarantee matching calling conventions and memory layouts. The harness ensures this by compiling the dylib on the same machine with the same toolchain. - **Shared API crates for custom types**: diff --git a/examples/cuda-softmax/Cargo.toml b/examples/cuda-softmax/Cargo.toml new file mode 100644 index 0000000..956da0e --- /dev/null +++ b/examples/cuda-softmax/Cargo.toml @@ -0,0 +1,31 @@ +[package] +edition = "2024" +license = "MPL-2.0" +name = "cuda-softmax-example" +publish = false +version.workspace = true + +[lints] +workspace = true + +[lib] +crate-type = ["dylib", "rlib"] + +[dependencies] +symbiont = { path = "../../symbiont" } + +# `dynamic-loading` (the default) resolves libcuda/libnvrtc with `dlopen` at +# runtime, so this example builds on machines with no CUDA toolkit and no GPU +# — it only needs a device to *run*. An explicit `cuda-12000` pins the driver +# API subset: the CUDA driver API is backward compatible, so binding to the +# 12.0 symbols works on every 12.x driver, while `cuda-version-from-build-system` +# would shell out to `nvcc` at build time and break GPU-less CI. +cudarc = { version = "0.19", default-features = false, features = [ + "cuda-12000", + "driver", + "dynamic-loading", + "nvrtc", + "std", +] } +tokio.workspace = true +tracing.workspace = true diff --git a/examples/cuda-softmax/README.md b/examples/cuda-softmax/README.md new file mode 100644 index 0000000..6d71667 --- /dev/null +++ b/examples/cuda-softmax/README.md @@ -0,0 +1,174 @@ +# CUDA Softmax — Evolving a GPU Kernel Against a Measured Roofline + +An LLM writes **CUDA kernels**, the host compiles them with NVRTC, checks them +against a CPU oracle, times them on the GPU, and feeds the achieved memory +throughput back into the next round. The best kernel is hot-swapped in; the +rest stay registered as revisions. + +```bash +MODEL=... cargo run -p cuda-softmax-example --release +``` + +The evolvable function is not the kernel itself — it is the thing that *emits* +one: + +```rust +fn plan(rows: usize, cols: usize) -> KernelPlan // { source, grid, block, shared_bytes } +``` + +The agent controls both halves, because both matter: the same source at the +wrong block size is an order of magnitude slower. + +## Why a GPU kernel is a good subject for an evolution loop + +- **The fitness function is not a matter of taste.** A row-wise softmax must + read `rows * cols` floats and write as many, and nothing else. The example + measures the ceiling instead of quoting a spec sheet: a copy kernel over the + *same buffers* on the *same card*, so a candidate's score is a percentage of + something real. (On a card with a large L2 the working set may be cache + resident — which is exactly why the ceiling is measured under the same + conditions as the candidates rather than taken from a datasheet.) +- **The gap is enormous and comes from strategy, not micro-edits.** The naive + kernel the run starts from is one thread per row: the 32 threads of a warp + touch addresses 4 KiB apart, so every load pulls a full cache line to use + four bytes of it, and the row is re-read from memory three times. Closing the + gap means coalescing, block-per-row reductions, warp shuffles, vectorized + `float4` loads, fast-math intrinsics — different *programs*, which is what + `evolve_batch` is for. +- **Correctness is checkable and the trap is real.** The benchmark input has + per-row biases up to ~90, so `expf(x)` overflows to `inf` in f32 unless the + row maximum is subtracted first. Deleting that pass is a tempting way to + drop a third of the memory traffic; the oracle catches it. + +## A real run + +RTX PRO 6000 Blackwell (188 SMs), 4096x1024 f32, one round of two lanes +against a local `Qwen3.6-35B-A3B` at Q4. One lane produced a kernel with an +incomplete reduction and was rejected by the oracle; the other: + +| kernel | time | throughput | % of ceiling | vs baseline | +|---|---|---|---|---| +| naive, one thread per row (the starting point) | 523 us | 64 GB/s | 1.1% | 1x | +| **evolved: one warp per row, pure warp shuffles** | **12.7 us** | **2640 GB/s** | **47%** | **41x** | +| hand-written reference: block per row, `float4` in registers | 7.1 us | 4730 GB/s | 84% | 74x | + +```cuda +// what the agent came up with, verbatim +extern "C" __global__ void softmax(const float* input, float* output, int rows, int cols) { + int row = blockIdx.x; + if (row >= rows) return; + const int tid = threadIdx.x; + const int stride = cols >> 5; + int base = row * cols; + + float local_max = -INFINITY; + for (int i = 0; i < stride; ++i) + local_max = fmaxf(local_max, input[base + tid + (i << 5)]); + for (int m = 16; m > 0; m >>= 1) + local_max = fmaxf(local_max, __shfl_down_sync(0xFFFFFFFF, local_max, m)); + float row_max = __shfl_sync(0xFFFFFFFF, local_max, 0); + + float local_sum = 0.0f; + for (int i = 0; i < stride; ++i) + local_sum += expf(input[base + tid + (i << 5)] - row_max); + for (int m = 16; m > 0; m >>= 1) + local_sum += __shfl_down_sync(0xFFFFFFFF, local_sum, m); + float row_sum = __shfl_sync(0xFFFFFFFF, local_sum, 0); + + float inv_sum = 1.0f / row_sum; + for (int i = 0; i < stride; ++i) + output[base + tid + (i << 5)] = expf(input[base + tid + (i << 5)] - row_max) * inv_sum; +} +``` + +It sidestepped the cross-warp reduction entirely by giving each row a single +warp — 32 threads, `__shfl_down_sync` only, no shared memory, no +`__syncthreads()`, and therefore none of the bugs that killed the other lane. +It is also *not* the best kernel available: at 32 threads per block it leaves +about half the bandwidth on the table, which the last row shows. The harness +reports that honestly rather than declaring victory — which is the point of +scoring against a measured ceiling instead of against the previous attempt. + +Whether a given model gets further is a property of the model. Everything the +loop needs to keep pushing — the incumbent's source, its throughput, and the +exact reason each rejected candidate was rejected — goes into the next round's +prompt. + +## The interesting part: a GPU kernel cannot be `catch_unwind`-ed + +Symbiont contains a misbehaving CPU implementation *in process*. Generated +functions are wrapped in `catch_unwind` inside the dylib, a panic becomes a +default return value, and the message is fed back to the agent. The loop never +stops. + +None of that is available on a GPU. An illegal memory access is a **sticky** +CUDA error: it does not merely fail the launch, it invalidates the context — +and, as this example's own test asserts, afterwards *no* context can be created +in that process at all. `cuDevicePrimaryCtxRetain` hands back the poisoned one +and `cuCtxCreate` fails too, even after every handle has been dropped: + +``` +fault: CUDA_ERROR_ILLEGAL_ADDRESS +primary: Err(CUDA_ERROR_ILLEGAL_ADDRESS) +non-primary: Err(CUDA_ERROR_ILLEGAL_ADDRESS) +``` + +So the process boundary *is* the GPU's `catch_unwind`. The parent never +launches agent-written kernels: it re-executes itself once per candidate with +`CUDA_SOFTMAX_EVAL_PLAN` set, and that child compiles, verifies, times, and +writes a one-line report. A candidate that faults the device, hangs, or +segfaults costs one child process and one table row — the search continues. +`tests/isolation.rs` asserts exactly that, end to end. + +## What the loop feeds back + +Each rejection is classified, because different failures need different +nudges: + +| kind | source | fed back as | +|---|---|---| +| `bad geometry` | host-side check | "block (2048,1,1) is 2048 threads, the device allows at most 1024" | +| `nvrtc error` | NVRTC log | the compiler diagnostics, with line numbers pointing at the agent's own source | +| `missing symbol` | module lookup | "the kernel must be declared `extern \"C\"`" | +| `wrong output` | CPU oracle | "output[0][0] = 6.38e-5, expected 2.04e-3 (tolerance 1e-6)" | +| `launch fault` | dead child | "the evaluation process exited without reporting" | + +Two details worth stealing: + +- NVRTC compiles without the CUDA toolkit headers, so `INFINITY`, `NAN` and + `FLT_MAX` are simply undefined — the single most common way for an otherwise + fine kernel to fail to compile. The host prepends guarded definitions plus a + `#line 1` directive, so NVRTC's diagnostics still point at the agent's line + numbers. +- The output buffer is filled with `NaN` before the correctness run, so a + kernel that leaves elements untouched fails instead of silently passing on + the previous candidate's results. + +## Structure + +- `src/lib.rs` — `KernelPlan`, the naive starting kernel, the deterministic + benchmark input and the `f64` CPU oracle. +- `src/gpu.rs` — every `unsafe` block in the example: NVRTC, launch, + verification, timing, and the measured copy ceiling. +- `src/isolate.rs` — the parent/child protocol. +- `src/main.rs` — the `evolvable!` declaration, the prompts, and the search. + +Agent code stays inside symbiont's policy the whole time: no `unsafe`, no +statics, no FFI. It emits *text* and typed launch parameters, and the host +decides whether that text is even allowed near the device. The generated dylib +is compiled in **debug** on purpose — the Rust it contains just builds a +string, and all the performance lives in the CUDA source. + +## Knobs + +| env | default | meaning | +|---|---|---| +| `MODEL` | required | model slug served at `BASE_URL` | +| `ROUNDS` | 3 | search rounds | +| `LANES` | 4 | candidates per round (8 strategy hints available, then they cycle) | +| `STRICT` | unset | fail instead of skipping when no CUDA device is present | + +Without an NVIDIA GPU the example prints why it is skipping and exits +successfully. It still *builds* anywhere: `cudarc`'s default `dynamic-loading` +resolves `libcuda`/`libnvrtc` with `dlopen` at runtime, so no CUDA toolkit is +needed at build time. diff --git a/examples/cuda-softmax/src/gpu.rs b/examples/cuda-softmax/src/gpu.rs new file mode 100644 index 0000000..b0277bf --- /dev/null +++ b/examples/cuda-softmax/src/gpu.rs @@ -0,0 +1,858 @@ +// SPDX-License-Identifier: MPL-2.0 +//! The host-owned GPU façade: NVRTC compilation, correctness gate, timing. +//! +//! Every `unsafe` block in this example lives in this file. Agent code never +//! sees a device pointer or a launch: it emits a [`KernelPlan`] and this side +//! decides whether that plan is even launchable, whether its output matches +//! the CPU oracle, and how fast it is. + +use std::{ + sync::Arc, + time::{ + Duration, + Instant, + }, +}; + +use cudarc::{ + driver::{ + CudaContext, + CudaSlice, + CudaStream, + DriverError, + LaunchConfig, + PushKernelArg, + sys::CUdevice_attribute, + }, + nvrtc::{ + CompileError, + CompileOptions, + compile_ptx_with_opts, + }, +}; + +use crate::{ + KERNEL_NAME, + KernelPlan, + benchmark_input, + reference_softmax, +}; + +/// How long to keep launching before the timed runs start. +/// +/// Long enough to matter: a fresh context runs at idle clocks, and these +/// kernels are tens of microseconds each, so a handful of warmup launches +/// would time the GPU on its way up to its boost clock. Measured ceilings +/// varied by 3x before this was time-based rather than a fixed count. +const WARMUP: Duration = Duration::from_millis(100); +/// Timed launches per candidate. Averaged, not median: the kernels here run in +/// tens of microseconds, where a single sync-per-run would dominate, so the +/// whole batch is timed between two synchronizations. +const TIMED_RUNS: usize = 50; +/// Largest absolute deviation from the CPU reference a candidate may have. +/// +/// Outputs are probabilities over 1024 columns, so the values themselves are +/// around 1e-3. This tolerance accepts the fast-math `__expf` intrinsic and +/// any sane reduction order, and rejects a missing max-subtraction outright +/// (that produces `inf`/`NaN`, not a small error). +const TOLERANCE: f32 = 1e-6; + +/// Prepended to every candidate before compilation. +/// +/// NVRTC compiles without the CUDA toolkit headers, so `` staples like +/// `INFINITY` and `FLT_MAX` are simply undefined — the most common way for an +/// otherwise fine kernel to fail to compile, and a distraction the search +/// should not spend rounds on. Every definition is guarded, so a candidate +/// that brings its own wins. +/// +/// The trailing `#line 1` resets NVRTC's line counter, so the line numbers in +/// its diagnostics refer to the agent's source rather than to this prelude. +const NVRTC_PRELUDE: &str = "#ifndef INFINITY\n\ + #define INFINITY __int_as_float(0x7f800000)\n\ + #endif\n\ + #ifndef NAN\n\ + #define NAN __int_as_float(0x7fffffff)\n\ + #endif\n\ + #ifndef FLT_MAX\n\ + #define FLT_MAX 3.402823466e+38f\n\ + #endif\n\ + #ifndef FLT_MIN\n\ + #define FLT_MIN 1.175494351e-38f\n\ + #endif\n\ + #line 1\n"; + +/// Static properties of the device the benchmark runs on. +#[derive(Debug, Clone)] +pub struct DeviceInfo { + /// Marketing name, e.g. `NVIDIA GeForce RTX 4090`. + pub name: String, + /// Compute capability as `(major, minor)`. + pub compute_capability: (i32, i32), + /// Number of streaming multiprocessors. + pub multiprocessors: i32, + /// Maximum threads in a single block. + pub max_threads_per_block: i32, + /// Maximum statically declared shared memory per block, in bytes. + pub max_shared_memory_per_block: i32, + /// Threads per warp (32 on every current architecture). + pub warp_size: i32, +} + +/// What a candidate kernel achieved. +#[derive(Debug, Clone, Copy)] +pub struct Measurement { + /// Mean kernel duration in microseconds. + pub micros: f64, + /// Effective memory throughput: one read plus one write of the matrix, + /// divided by the kernel duration. + pub gb_per_s: f64, + /// [`Measurement::gb_per_s`] as a percentage of the device-to-device copy + /// throughput measured on the same buffers — the practical roofline for a + /// memory-bound kernel. + pub pct_of_roofline: f64, + /// Largest absolute deviation from the CPU reference. + pub max_abs_err: f32, +} + +/// Why a candidate kernel did not produce a measurement. +#[derive(Debug, Clone)] +pub enum KernelFailure { + /// The launch geometry was rejected before reaching the driver. + Geometry(String), + /// NVRTC refused the source; carries the compiler log. + Compile(String), + /// The module compiled but did not export [`KERNEL_NAME`]. + MissingSymbol(String), + /// The driver rejected the launch, or the kernel faulted while running. + Launch(String), + /// The kernel ran but its output does not match the CPU reference. + Wrong { + /// Largest absolute deviation found. + max_abs_err: f32, + /// Where the first bad element was and what it looked like. + detail: String, + }, +} + +impl KernelFailure { + /// Whether this failure leaves the CUDA context unusable. + /// + /// An illegal memory access is a *sticky* error: every subsequent call + /// fails too, and not even a brand new context can be created in this + /// process afterwards, so the only recovery is to replace the process + /// ([`crate::Isolated`]). Compile errors and wrong answers, by contrast, + /// leave the device perfectly healthy — which is why they are separate + /// variants rather than one `Error(String)`. + #[must_use] + pub fn poisons_context(&self) -> bool { + matches!(self, Self::Launch(_)) + } + + /// One-line label for report tables. + #[must_use] + pub fn kind(&self) -> &'static str { + match self { + Self::Geometry(_) => "bad geometry", + Self::Compile(_) => "nvrtc error", + Self::MissingSymbol(_) => "missing symbol", + Self::Launch(_) => "launch fault", + Self::Wrong { .. } => "wrong output", + } + } +} + +impl std::fmt::Display for KernelFailure { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Geometry(msg) | Self::MissingSymbol(msg) | Self::Launch(msg) => { + write!(f, "{}: {msg}", self.kind()) + } + Self::Compile(log) => write!(f, "nvrtc error:\n{log}"), + Self::Wrong { + max_abs_err, + detail, + } if max_abs_err.is_finite() => { + write!( + f, + "wrong output (max abs error {max_abs_err:.3e}): {detail}" + ) + } + Self::Wrong { detail, .. } => write!(f, "wrong output: {detail}"), + } + } +} + +impl KernelFailure { + /// The message without the [`std::fmt::Display`] prefix, for a wire format + /// that carries the kind separately — otherwise a failure round-tripped + /// through a child process ends up labelled twice. + #[must_use] + pub fn detail(&self) -> &str { + match self { + Self::Geometry(msg) + | Self::Compile(msg) + | Self::MissingSymbol(msg) + | Self::Launch(msg) + | Self::Wrong { detail: msg, .. } => msg, + } + } +} + +/// Failure to set up the device itself, as opposed to a bad +/// candidate kernel. +#[derive(Debug)] +pub enum GpuError { + /// `libcuda` could not be loaded: no driver installed. + DriverUnavailable, + /// A driver call failed while setting up the benchmark. + Driver(DriverError), + /// The host's own reference kernel misbehaved — a bug in this example, + /// not in anything the agent produced. + Internal(String), +} + +impl From for GpuError { + fn from(err: DriverError) -> Self { + Self::Driver(err) + } +} + +impl std::fmt::Display for GpuError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::DriverUnavailable => write!( + f, + "the CUDA driver library could not be loaded (no NVIDIA driver installed?)" + ), + Self::Driver(err) => write!(f, "CUDA driver error: {err}"), + Self::Internal(msg) => write!(f, "internal error: {msg}"), + } + } +} + +impl std::error::Error for GpuError {} + +/// Owns the CUDA context, the benchmark buffers and the correctness oracle. +/// +/// Lives in the *child* process for the duration of one candidate (see +/// [`crate::Isolated`]), and in the parent only to report the device and the +/// copy ceiling. The evolvable function never touches it. +#[derive(Debug)] +pub struct Gpu { + ctx: Arc, + stream: Arc, + input: CudaSlice, + output: CudaSlice, + rows: usize, + cols: usize, + reference: Vec, + /// Written into the output buffer before the correctness run, so a kernel + /// that leaves elements untouched fails the comparison instead of reading + /// as "the previous candidate's answer" or "zero". + poison: Vec, + info: DeviceInfo, + /// `compute_XX`, leaked once per context because + /// [`CompileOptions::arch`] wants a `&'static str`. + arch: Option<&'static str>, + roofline_gb_per_s: f64, +} + +impl Gpu { + /// Bring up device 0 and stage a `rows` x `cols` benchmark on it. + /// + /// # Errors + /// + /// [`GpuError::DriverUnavailable`] when there is no CUDA driver to load, + /// so a caller can skip the example instead of dying, and + /// [`GpuError::Driver`] for a device that is present but unusable. + pub fn new(rows: usize, cols: usize) -> Result { + if !driver_present() { + return Err(GpuError::DriverUnavailable); + } + let ctx = CudaContext::new(0)?; + let info = device_info(&ctx)?; + let arch = arch_flag(info.compute_capability); + let host_input = benchmark_input(rows, cols); + let reference = reference_softmax(&host_input, rows, cols); + let poison = vec![f32::NAN; rows * cols]; + + let stream = ctx.default_stream(); + let input = stream.clone_htod(&host_input)?; + let output = stream.alloc_zeros::(rows * cols)?; + + let mut gpu = Self { + ctx, + stream, + input, + output, + rows, + cols, + reference, + poison, + info, + arch, + roofline_gb_per_s: 0.0, + }; + gpu.roofline_gb_per_s = match inherited_roofline() { + Some(gb_per_s) => gb_per_s, + None => gpu.measure_copy_roofline()?, + }; + Ok(gpu) + } + + /// Static properties of the device. + #[must_use] + pub fn info(&self) -> &DeviceInfo { + &self.info + } + + /// Device-to-device copy throughput over the benchmark buffers: the + /// practical bandwidth ceiling a memory-bound kernel is measured against. + #[must_use] + pub fn roofline_gb_per_s(&self) -> f64 { + self.roofline_gb_per_s + } + + /// Bytes a correct kernel must move at minimum: read the matrix once, + /// write it once. + #[must_use] + pub fn traffic_bytes(&self) -> u64 { + 2 * (self.rows * self.cols * size_of::()) as u64 + } + + /// Compile, verify and time one candidate. + /// + /// # Errors + /// + /// A [`KernelFailure`] describing exactly which stage rejected the + /// candidate, phrased so it can be fed straight back to the agent. + pub fn evaluate(&mut self, plan: &KernelPlan) -> Result { + self.check_geometry(plan)?; + let ptx = compile_ptx_with_opts( + format!("{NVRTC_PRELUDE}{}", plan.source), + CompileOptions { + arch: self.arch, + ..Default::default() + }, + ) + .map_err(|err| KernelFailure::Compile(compile_log(&err)))?; + + let module = self + .ctx + .load_module(ptx) + .map_err(|err| KernelFailure::Compile(format!("PTX failed to load: {err}")))?; + let func = module.load_function(KERNEL_NAME).map_err(|err| { + KernelFailure::MissingSymbol(format!( + "the module does not export `{KERNEL_NAME}` ({err}); \ + the kernel must be declared `extern \"C\"`" + )) + })?; + + let cfg = LaunchConfig { + grid_dim: plan.grid, + block_dim: plan.block, + shared_mem_bytes: plan.shared_bytes, + }; + let rows = i32::try_from(self.rows).expect("row count fits in i32"); + let cols = i32::try_from(self.cols).expect("column count fits in i32"); + + // Poison the output so untouched elements read as NaN. + self.stream + .memcpy_htod(&self.poison, &mut self.output) + .map_err(|err| KernelFailure::Launch(format!("could not reset the output: {err}")))?; + + // Split the borrow: the builder needs `&input` and `&mut output` at + // the same time, which is only expressible per field. + let Self { + stream, + input, + output, + .. + } = self; + let mut launch = stream.launch_builder(&func); + launch.arg(&*input).arg(&mut *output).arg(&rows).arg(&cols); + + // SAFETY: nothing about a machine-written kernel is safe, which is the + // whole point of running it behind a correctness gate and a context we + // are prepared to throw away. The arguments do match the signature the + // agent is required to implement, and `check_geometry` has already + // ruled out the launch configurations the driver would reject. + unsafe { launch.launch(cfg) } + .map_err(|err| KernelFailure::Launch(format!("launch rejected: {err}")))?; + stream + .synchronize() + .map_err(|err| KernelFailure::Launch(format!("kernel faulted: {err}")))?; + + let warmup_deadline = Instant::now() + WARMUP; + while Instant::now() < warmup_deadline { + // SAFETY: as above; the first launch already completed cleanly. + unsafe { launch.launch(cfg) } + .map_err(|err| KernelFailure::Launch(format!("launch rejected: {err}")))?; + stream + .synchronize() + .map_err(|err| KernelFailure::Launch(format!("kernel faulted: {err}")))?; + } + + let started = Instant::now(); + for _ in 0..TIMED_RUNS { + // SAFETY: as above. + unsafe { launch.launch(cfg) } + .map_err(|err| KernelFailure::Launch(format!("launch rejected: {err}")))?; + } + stream + .synchronize() + .map_err(|err| KernelFailure::Launch(format!("kernel faulted: {err}")))?; + let elapsed = started.elapsed(); + + // Verification reads back the *warm* result, so a kernel cannot pass + // by being correct once and racy afterwards. + let max_abs_err = self.verify()?; + + let seconds = elapsed.as_secs_f64() / TIMED_RUNS as f64; + let gb_per_s = self.traffic_bytes() as f64 / seconds / 1e9; + Ok(Measurement { + micros: seconds * 1e6, + gb_per_s, + pct_of_roofline: 100.0 * gb_per_s / self.roofline_gb_per_s, + max_abs_err, + }) + } + + /// Reject launch geometry the driver would refuse, with a message the + /// agent can act on. Cheaper than a round trip to the device, and it keeps + /// trivially malformed plans out of the fault-recovery path. + fn check_geometry(&self, plan: &KernelPlan) -> Result<(), KernelFailure> { + let threads = plan.threads_per_block(); + let max_threads = u64::try_from(self.info.max_threads_per_block).unwrap_or(1024); + if threads == 0 || plan.blocks() == 0 { + return Err(KernelFailure::Geometry(format!( + "grid {:?} x block {:?} launches no threads", + plan.grid, plan.block + ))); + } + if threads > max_threads { + return Err(KernelFailure::Geometry(format!( + "block {:?} is {threads} threads, the device allows at most {max_threads}", + plan.block + ))); + } + let rows = self.rows as u64; + if plan.blocks() * threads < rows { + return Err(KernelFailure::Geometry(format!( + "grid {:?} x block {:?} is {} threads for {rows} rows: some rows would go \ + unprocessed", + plan.grid, + plan.block, + plan.blocks() * threads + ))); + } + Ok(()) + } + + /// Compare the device output against the CPU oracle. + fn verify(&self) -> Result { + let got = self.stream.clone_dtoh(&self.output).map_err(|err| { + KernelFailure::Launch(format!("could not read the output back: {err}")) + })?; + + let mut max_abs_err = 0.0_f32; + let mut first_bad: Option<(usize, f32, f32)> = None; + for (idx, (&got, &want)) in got.iter().zip(self.reference.iter()).enumerate() { + let err = (got - want).abs(); + if got.is_finite() && err <= max_abs_err { + continue; + } + if got.is_finite() { + max_abs_err = err; + } + if (!got.is_finite() || err > TOLERANCE) && first_bad.is_none() { + first_bad = Some((idx, got, want)); + } + } + + if let Some((idx, got, want)) = first_bad { + let (row, col) = (idx / self.cols, idx % self.cols); + return Err(KernelFailure::Wrong { + max_abs_err, + detail: format!( + "output[{row}][{col}] = {got:e}, expected {want:e} (tolerance {TOLERANCE:e})" + ), + }); + } + Ok(max_abs_err) + } + + /// Time a kernel that only copies the benchmark buffers. + /// + /// A copy moves exactly the traffic a perfect softmax must move — one read + /// and one write — so its throughput is the honest ceiling to report + /// candidates against, and it self-calibrates on whatever card is present + /// instead of trusting a spec sheet. + /// + /// It is deliberately a *kernel* and not `cuMemcpyDtoD`: the driver copy is + /// blocking, so at these sizes its per-call overhead would dominate and + /// produce a ceiling that candidates comfortably exceed. Measuring the + /// ceiling through the same launch path as the candidates keeps the + /// comparison apples to apples. + fn measure_copy_roofline(&mut self) -> Result { + let internal = |what: &str, err: &dyn std::fmt::Display| { + GpuError::Internal(format!("the built-in copy kernel {what}: {err}")) + }; + let ptx = compile_ptx_with_opts( + format!("{NVRTC_PRELUDE}{COPY_KERNEL}"), + CompileOptions { + arch: self.arch, + ..Default::default() + }, + ) + .map_err(|err| internal("did not compile", &compile_log(&err)))?; + let module = self + .ctx + .load_module(ptx) + .map_err(|err| internal("did not load", &err))?; + let func = module + .load_function("copy_kernel") + .map_err(|err| internal("has no entry point", &err))?; + + let elements = i32::try_from(self.rows * self.cols / 4).expect("float4 count fits in i32"); + let cfg = LaunchConfig::for_num_elems(u32::try_from(elements).expect("positive")); + let Self { + stream, + input, + output, + .. + } = self; + let mut launch = stream.launch_builder(&func); + launch.arg(&*input).arg(&mut *output).arg(&elements); + + let warmup_deadline = Instant::now() + WARMUP; + while Instant::now() < warmup_deadline { + // SAFETY: the copy kernel is host-written, its arguments match the + // signature above, and `for_num_elems` covers every element once. + unsafe { launch.launch(cfg) }.map_err(|err| internal("failed to launch", &err))?; + stream + .synchronize() + .map_err(|err| internal("faulted", &err))?; + } + let started = Instant::now(); + for _ in 0..TIMED_RUNS { + // SAFETY: as above. + unsafe { launch.launch(cfg) }.map_err(|err| internal("failed to launch", &err))?; + } + stream + .synchronize() + .map_err(|err| internal("faulted", &err))?; + let seconds = started.elapsed().as_secs_f64() / TIMED_RUNS as f64; + Ok(self.traffic_bytes() as f64 / seconds / 1e9) + } +} + +/// The bandwidth ceiling, expressed as a kernel: read the matrix once as +/// `float4`, write it once, do nothing else. +const COPY_KERNEL: &str = r#" +extern "C" __global__ void copy_kernel(const float4* __restrict__ input, + float4* __restrict__ output, + int n4) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n4) output[i] = input[i]; +} +"#; + +/// The copy ceiling handed down by a parent process, if any. +fn inherited_roofline() -> Option { + std::env::var(crate::ROOFLINE_ENV) + .ok()? + .parse::() + .ok() + .filter(|gb_per_s| *gb_per_s > 0.0) +} + +/// Whether `libcuda` can be loaded at all. +fn driver_present() -> bool { + // SAFETY: this only attempts a `dlopen` of the driver library and reports + // whether it succeeded. It initializes no CUDA state. + unsafe { cudarc::driver::sys::is_culib_present() } +} + +/// NVRTC's `--gpu-architecture` flag for a compute capability. +/// +/// `CompileOptions::arch` is a `&'static str`, so the string is leaked — once +/// per process in practice, since the device does not change under us. Falling +/// back to `None` for an unknown capability is safe: NVRTC then emits PTX for +/// its default architecture and the driver JITs it for the real device. +fn arch_flag(compute_capability: (i32, i32)) -> Option<&'static str> { + let (major, minor) = compute_capability; + if major < 5 { + return None; + } + Some(Box::leak( + format!("compute_{major}{minor}").into_boxed_str(), + )) +} + +/// Query the device properties worth putting in a prompt. +fn device_info(ctx: &Arc) -> Result { + Ok(DeviceInfo { + name: ctx.name()?, + compute_capability: ctx.compute_capability()?, + multiprocessors: ctx + .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)?, + max_threads_per_block: ctx + .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK)?, + max_shared_memory_per_block: ctx + .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK)?, + warp_size: ctx.attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_WARP_SIZE)?, + }) +} + +/// Pull the compiler log out of an NVRTC error. +/// +/// [`CompileError`]'s own `Display` is its `Debug`, which buries the log in +/// escaped `CString` output next to the full option list. The agent needs the +/// diagnostics and nothing else. +fn compile_log(err: &CompileError) -> String { + match err { + CompileError::CompileError { log, .. } => { + let log = log.to_string_lossy().trim().to_string(); + if log.is_empty() { + "nvrtc rejected the source without a diagnostic".to_string() + } else { + log + } + } + other => format!("{other:?}"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + COLS, + NAIVE_KERNEL, + }; + + /// Rows used by the tests. Small enough that the `f64` CPU oracle is cheap + /// in a debug build, large enough to fill a modern GPU. + const TEST_ROWS: usize = 512; + + /// A kernel a competent optimizer would write: one block per row, the row + /// held in registers as one `float4` per thread, and two block-wide + /// reductions over warp shuffles. Test-only on purpose — it is the answer, + /// and the host crate's rustdoc is fed to the agent. + const OPTIMIZED_KERNEL: &str = r#" +extern "C" __global__ void softmax(const float* __restrict__ input, + float* __restrict__ output, + int rows, int cols) { + int row = blockIdx.x; + if (row >= rows) return; + int n4 = cols >> 2; + + const float4* in4 = (const float4*)(input + (size_t)row * cols); + float4* out4 = (float4*)(output + (size_t)row * cols); + + float4 v = make_float4(-INFINITY, -INFINITY, -INFINITY, -INFINITY); + if (threadIdx.x < n4) v = in4[threadIdx.x]; + + float m = fmaxf(fmaxf(v.x, v.y), fmaxf(v.z, v.w)); + for (int off = 16; off > 0; off >>= 1) { + m = fmaxf(m, __shfl_down_sync(0xffffffff, m, off)); + } + __shared__ float warp_max[32]; + __shared__ float warp_sum[32]; + int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + int warps = (blockDim.x + 31) >> 5; + if (lane == 0) warp_max[warp] = m; + __syncthreads(); + if (threadIdx.x == 0) { + float acc = warp_max[0]; + for (int i = 1; i < warps; ++i) acc = fmaxf(acc, warp_max[i]); + warp_max[0] = acc; + } + __syncthreads(); + float row_max = warp_max[0]; + + float ex = __expf(v.x - row_max), ey = __expf(v.y - row_max); + float ez = __expf(v.z - row_max), ew = __expf(v.w - row_max); + float s = ex + ey + ez + ew; + for (int off = 16; off > 0; off >>= 1) { + s += __shfl_down_sync(0xffffffff, s, off); + } + if (lane == 0) warp_sum[warp] = s; + __syncthreads(); + if (threadIdx.x == 0) { + float acc = 0.0f; + for (int i = 0; i < warps; ++i) acc += warp_sum[i]; + warp_sum[0] = acc; + } + __syncthreads(); + float inv = 1.0f / warp_sum[0]; + + if (threadIdx.x < n4) { + out4[threadIdx.x] = make_float4(ex * inv, ey * inv, ez * inv, ew * inv); + } +} +"#; + + fn naive_plan(rows: usize) -> KernelPlan { + KernelPlan { + source: NAIVE_KERNEL.to_string(), + grid: (u32::try_from(rows.div_ceil(256)).expect("fits"), 1, 1), + block: (256, 1, 1), + shared_bytes: 0, + } + } + + fn optimized_plan(rows: usize) -> KernelPlan { + KernelPlan { + source: OPTIMIZED_KERNEL.to_string(), + grid: (u32::try_from(rows).expect("fits"), 1, 1), + block: (256, 1, 1), + shared_bytes: 0, + } + } + + /// Every test needs a device; without one they report and pass, so the + /// suite stays green on machines that cannot run it. + fn gpu() -> Option { + match Gpu::new(TEST_ROWS, COLS) { + Ok(gpu) => Some(gpu), + Err(err) => { + eprintln!("skipping: {err}"); + None + } + } + } + + #[test] + fn the_naive_kernel_is_correct() { + let Some(mut gpu) = gpu() else { return }; + let measurement = gpu + .evaluate(&naive_plan(TEST_ROWS)) + .expect("the naive kernel is correct"); + assert!( + measurement.max_abs_err <= TOLERANCE, + "max abs err {}", + measurement.max_abs_err + ); + assert!(measurement.gb_per_s > 0.0); + eprintln!( + "naive: {:.1} us, {:.0} GB/s ({:.1}% of ceiling)", + measurement.micros, measurement.gb_per_s, measurement.pct_of_roofline + ); + } + + #[test] + fn an_optimized_kernel_measures_faster_than_the_naive_one() { + let Some(mut gpu) = gpu() else { return }; + let naive = gpu + .evaluate(&naive_plan(TEST_ROWS)) + .expect("the naive kernel is correct"); + let optimized = gpu + .evaluate(&optimized_plan(TEST_ROWS)) + .expect("the optimized kernel is correct"); + eprintln!( + "naive {:.1} us -> optimized {:.1} us ({:.0} GB/s, {:.1}% of ceiling)", + naive.micros, optimized.micros, optimized.gb_per_s, optimized.pct_of_roofline + ); + // The point of the metric is that it separates these two by a lot. + assert!( + optimized.micros * 3.0 < naive.micros, + "expected a large speedup, got {:.1} us vs {:.1} us", + optimized.micros, + naive.micros + ); + } + + #[test] + fn a_kernel_that_skips_the_max_subtraction_is_rejected() { + let Some(mut gpu) = gpu() else { return }; + let plan = KernelPlan { + source: r#" +extern "C" __global__ void softmax(const float* input, float* output, int rows, int cols) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= rows) return; + const float* in = input + (size_t)row * cols; + float* out = output + (size_t)row * cols; + float sum = 0.0f; + for (int i = 0; i < cols; ++i) sum += expf(in[i]); + for (int i = 0; i < cols; ++i) out[i] = expf(in[i]) / sum; +} +"# + .to_string(), + ..naive_plan(TEST_ROWS) + }; + let failure = gpu.evaluate(&plan).expect_err("inf/inf is not a softmax"); + assert!( + matches!(failure, KernelFailure::Wrong { .. }), + "{failure:?}" + ); + assert!(!failure.poisons_context()); + } + + #[test] + fn nvrtc_diagnostics_come_back_as_feedback() { + let Some(mut gpu) = gpu() else { return }; + let plan = KernelPlan { + source: "extern \"C\" __global__ void softmax(int a) { this is not c++ }".to_string(), + ..naive_plan(TEST_ROWS) + }; + let failure = gpu.evaluate(&plan).expect_err("that does not compile"); + let KernelFailure::Compile(log) = &failure else { + panic!("{failure:?}") + }; + assert!(log.contains("error"), "log without diagnostics: {log}"); + } + + #[test] + fn unlaunchable_geometry_is_rejected_before_the_driver_sees_it() { + let Some(mut gpu) = gpu() else { return }; + let too_many_threads = KernelPlan { + block: (2048, 1, 1), + ..naive_plan(TEST_ROWS) + }; + assert!(matches!( + gpu.evaluate(&too_many_threads), + Err(KernelFailure::Geometry(_)) + )); + + let too_few_threads = KernelPlan { + grid: (1, 1, 1), + block: (32, 1, 1), + ..naive_plan(TEST_ROWS) + }; + assert!(matches!( + gpu.evaluate(&too_few_threads), + Err(KernelFailure::Geometry(_)) + )); + } + + #[test] + fn an_out_of_bounds_kernel_is_reported_as_poisoning_the_context() { + let Some(mut gpu) = gpu() else { return }; + let plan = KernelPlan { + source: r#" +extern "C" __global__ void softmax(const float* input, float* output, int rows, int cols) { + output[(size_t)(blockIdx.x + 1) * 1000000000ull] = 1.0f; +} +"# + .to_string(), + ..naive_plan(TEST_ROWS) + }; + let failure = gpu.evaluate(&plan).expect_err("that writes out of bounds"); + assert!( + failure.poisons_context(), + "an illegal access is sticky: {failure:?}" + ); + + // The device is unusable from here on. The naive kernel passed a + // moment ago in `the_naive_kernel_is_correct`; now the very same plan + // cannot run, and neither retaining the primary context nor creating + // an independent one recovers it. That is why candidates are evaluated + // in a child process instead of being caught in-place. + assert!( + gpu.evaluate(&naive_plan(TEST_ROWS)).is_err(), + "if this ever starts passing, in-process recovery became possible" + ); + } +} diff --git a/examples/cuda-softmax/src/isolate.rs b/examples/cuda-softmax/src/isolate.rs new file mode 100644 index 0000000..ce172c0 --- /dev/null +++ b/examples/cuda-softmax/src/isolate.rs @@ -0,0 +1,383 @@ +// SPDX-License-Identifier: MPL-2.0 +//! Running candidate kernels in a child process. +//! +//! Symbiont contains a misbehaving *CPU* implementation inside the dylib: the +//! generated code is wrapped in `catch_unwind`, a panic is turned into a +//! default return value, and the message is fed back to the agent. The loop +//! keeps running in the same process. +//! +//! A GPU kernel cannot be contained that way. An illegal memory access is a +//! **sticky** CUDA error: it does not just fail the launch, it invalidates the +//! context — and, as this example's own probe showed, `cuCtxCreate` and +//! `cuDevicePrimaryCtxRetain` then keep returning `CUDA_ERROR_ILLEGAL_ADDRESS` +//! for the rest of the process. Dropping every handle first does not help. +//! There is no in-process recovery; the process is what has to be replaced. +//! +//! So the parent process never launches agent-written kernels. It re-executes +//! itself once per candidate with [`EVAL_ENV`] set, and that child does the +//! compiling, verifying and timing. A candidate that faults the device, or +//! hangs, or segfaults the process outright, costs one child and one table +//! row — the search continues. + +use std::{ + io, + path::{ + Path, + PathBuf, + }, + process::Command, + time::{ + Duration, + Instant, + }, +}; + +use crate::{ + COLS, + Gpu, + KernelFailure, + KernelPlan, + Measurement, + ROWS, +}; + +/// Set on the child process; names the file holding the plan to evaluate. +/// The child writes its verdict to the same path with `.report` appended. +pub const EVAL_ENV: &str = "CUDA_SOFTMAX_EVAL_PLAN"; +/// Set on the child process; the copy ceiling every candidate is scored +/// against, in GB/s. +/// +/// Measured once by the parent and handed down, so that every percentage in a +/// run refers to the same number. Letting each child measure its own ceiling +/// made the reports incoherent — a kernel came back as "2613 GB/s, 46% of a +/// 636 GB/s ceiling" because the two processes had caught the GPU at +/// different clock states. +pub const ROOFLINE_ENV: &str = "CUDA_SOFTMAX_ROOFLINE"; +/// Appended to the plan path to get the report path. +const REPORT_SUFFIX: &str = ".report"; +/// How long a candidate may take before the child is killed. Generous: it +/// covers CUDA init, NVRTC, and the timed runs. A kernel that exceeds it is +/// almost certainly stuck in an unterminated loop. +const EVAL_TIMEOUT: Duration = Duration::from_secs(90); + +/// Spawns child processes to evaluate candidate kernels. +#[derive(Debug)] +pub struct Isolated { + exe: PathBuf, + dir: PathBuf, + next: std::cell::Cell, + roofline_gb_per_s: Option, +} + +impl Isolated { + /// Prepare a scratch directory for plan/report files. + /// + /// # Errors + /// + /// If the running executable cannot be located or the scratch directory + /// cannot be created. + pub fn new() -> io::Result { + Self::with_executable(std::env::current_exe()?) + } + + /// [`Isolated::new`] with an explicit evaluator binary. + /// + /// The binary must call [`evaluator_main`] before doing anything else. + /// Tests use this to drive the real child protocol against the example's + /// binary, which `current_exe` would not point at. + /// + /// # Errors + /// + /// If the scratch directory cannot be created. + pub fn with_executable(exe: impl Into) -> io::Result { + let dir = std::env::temp_dir().join(format!("cuda-softmax-{}", std::process::id())); + std::fs::create_dir_all(&dir)?; + Ok(Self { + exe: exe.into(), + dir, + next: std::cell::Cell::new(0), + roofline_gb_per_s: None, + }) + } + + /// Score every candidate against `gb_per_s` instead of letting each child + /// measure its own ceiling. See [`ROOFLINE_ENV`]. + #[must_use] + pub fn with_roofline(mut self, gb_per_s: f64) -> Self { + self.roofline_gb_per_s = Some(gb_per_s); + self + } + + /// Compile, verify and time `plan` in a child process. + /// + /// # Errors + /// + /// The [`KernelFailure`] the child reported, or a synthesized one when the + /// child died, timed out, or produced no report — all of which mean the + /// candidate took the CUDA context down with it. + pub fn evaluate(&self, plan: &KernelPlan) -> Result { + let id = self.next.get(); + self.next.set(id + 1); + let plan_path = self.dir.join(format!("plan-{id}.cu")); + let report_path = report_path_of(&plan_path); + + write_plan(&plan_path, plan) + .map_err(|err| KernelFailure::Launch(format!("could not stage the plan: {err}")))?; + + let status = self.run_child(&plan_path)?; + let report = std::fs::read_to_string(&report_path).map_err(|_| { + KernelFailure::Launch(format!( + "the evaluation process exited with {status} without reporting — the kernel \ + faulted the CUDA context or crashed the process" + )) + })?; + parse_report(&report) + } + + /// Run one child to completion, killing it if it overruns. + fn run_child(&self, plan_path: &Path) -> Result { + let spawn_failed = + |err: &dyn std::fmt::Display| KernelFailure::Launch(format!("child process: {err}")); + let mut command = Command::new(&self.exe); + command + .env(EVAL_ENV, plan_path) + // The child is a measurement harness, not a participant in the + // evolution loop: keep its logs out of the parent's report. + .env("RUST_LOG", "error"); + if let Some(roofline) = self.roofline_gb_per_s { + command.env(ROOFLINE_ENV, roofline.to_string()); + } + let mut child = command.spawn().map_err(|err| spawn_failed(&err))?; + + let deadline = Instant::now() + EVAL_TIMEOUT; + loop { + match child.try_wait().map_err(|err| spawn_failed(&err))? { + Some(status) => return Ok(status.to_string()), + None if Instant::now() >= deadline => { + let _ = child.kill(); + let _ = child.wait(); + return Err(KernelFailure::Launch(format!( + "the kernel did not finish within {}s and was killed — check for an \ + unterminated loop", + EVAL_TIMEOUT.as_secs() + ))); + } + None => std::thread::sleep(Duration::from_millis(20)), + } + } + } +} + +/// The child half: evaluate the staged plan and write the report. +/// +/// Returns `None` in the parent process (where [`EVAL_ENV`] is unset), so +/// `main` can call it first thing and carry on. Returns `Some(exit_code)` in a +/// child, which should exit with it immediately. +#[must_use] +pub fn evaluator_main() -> Option { + let plan_path = PathBuf::from(std::env::var_os(EVAL_ENV)?); + let report_path = report_path_of(&plan_path); + + let report = match read_plan(&plan_path) { + Ok(plan) => match Gpu::new(ROWS, COLS) { + // Anything past this point may kill the process; that is the + // reason this code is in a child at all. + Ok(mut gpu) => match gpu.evaluate(&plan) { + Ok(measurement) => format_ok(&measurement), + Err(failure) => format_err(&failure), + }, + Err(err) => format!("err\tdevice\t{}", escape(&err.to_string())), + }, + Err(err) => format!("err\tplan\t{}", escape(&err.to_string())), + }; + + // Only reached when the candidate did not take the process down with it. + if let Err(err) = std::fs::write(&report_path, report) { + eprintln!("could not write {}: {err}", report_path.display()); + return Some(2); + } + Some(0) +} + +fn report_path_of(plan_path: &Path) -> PathBuf { + let mut path = plan_path.as_os_str().to_owned(); + path.push(REPORT_SUFFIX); + PathBuf::from(path) +} + +/// Plan file format: three header lines, then the CUDA source verbatim. +fn write_plan(path: &Path, plan: &KernelPlan) -> io::Result<()> { + let (gx, gy, gz) = plan.grid; + let (bx, by, bz) = plan.block; + std::fs::write( + path, + format!( + "grid {gx} {gy} {gz}\nblock {bx} {by} {bz}\nshared {}\nsource\n{}", + plan.shared_bytes, plan.source + ), + ) +} + +fn read_plan(path: &Path) -> io::Result { + let raw = std::fs::read_to_string(path)?; + let (header, source) = raw + .split_once("source\n") + .ok_or_else(|| io::Error::other("malformed plan: no source marker"))?; + let mut plan = KernelPlan { + source: source.to_string(), + ..KernelPlan::default() + }; + for line in header.lines() { + let mut parts = line.split_whitespace(); + let field = parts.next().unwrap_or_default(); + let mut number = || { + parts + .next() + .and_then(|v| v.parse::().ok()) + .unwrap_or(0) + }; + match field { + "grid" => plan.grid = (number(), number(), number()), + "block" => plan.block = (number(), number(), number()), + "shared" => plan.shared_bytes = number(), + _ => return Err(io::Error::other(format!("malformed plan line: {line}"))), + } + } + Ok(plan) +} + +fn format_ok(m: &Measurement) -> String { + format!( + "ok\t{}\t{}\t{}\t{}", + m.micros, m.gb_per_s, m.pct_of_roofline, m.max_abs_err + ) +} + +fn format_err(failure: &KernelFailure) -> String { + let max_abs_err = match failure { + KernelFailure::Wrong { max_abs_err, .. } => max_abs_err.to_string(), + _ => "-".to_string(), + }; + format!( + "err\t{}\t{max_abs_err}\t{}", + failure.kind(), + escape(failure.detail()) + ) +} + +/// Reports are one line, so embedded newlines are escaped. +fn escape(text: &str) -> String { + text.replace('\\', "\\\\").replace('\n', "\\n") +} + +fn unescape(text: &str) -> String { + text.replace("\\n", "\n").replace("\\\\", "\\") +} + +fn parse_report(report: &str) -> Result { + let mut fields = report.trim_end().split('\t'); + match fields.next() { + Some("ok") => { + let mut number = || fields.next().and_then(|v| v.parse::().ok()); + let (Some(micros), Some(gb_per_s), Some(pct_of_roofline), Some(max_abs_err)) = + (number(), number(), number(), number()) + else { + return Err(KernelFailure::Launch(format!( + "unparseable report: {report}" + ))); + }; + Ok(Measurement { + micros, + gb_per_s, + pct_of_roofline, + max_abs_err: max_abs_err as f32, + }) + } + Some("err") => { + let kind = fields.next().unwrap_or("unknown"); + let max_abs_err = fields + .next() + .and_then(|v| v.parse::().ok()) + .unwrap_or(f32::NAN); + let message = unescape(fields.next().unwrap_or_default()); + Err(rehydrate(kind, max_abs_err, message)) + } + _ => Err(KernelFailure::Launch(format!( + "unparseable report: {report}" + ))), + } +} + +/// Turn a report line back into the variant the child produced, so the parent +/// can classify it the same way in its table. +fn rehydrate(kind: &str, max_abs_err: f32, message: String) -> KernelFailure { + match kind { + "bad geometry" => KernelFailure::Geometry(message), + "nvrtc error" => KernelFailure::Compile(message), + "missing symbol" => KernelFailure::MissingSymbol(message), + "wrong output" => KernelFailure::Wrong { + max_abs_err, + detail: message, + }, + _ => KernelFailure::Launch(message), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn plans_round_trip_through_the_staging_file() { + let dir = std::env::temp_dir().join("cuda-softmax-roundtrip"); + std::fs::create_dir_all(&dir).expect("scratch dir"); + let path = dir.join("plan.cu"); + let plan = KernelPlan { + source: "extern \"C\" __global__ void softmax() {\n // source\n}\n".to_string(), + grid: (7, 2, 1), + block: (128, 1, 1), + shared_bytes: 4096, + }; + write_plan(&path, &plan).expect("write"); + assert_eq!(read_plan(&path).expect("read"), plan); + } + + #[test] + fn measurements_round_trip_through_the_report() { + let measurement = Measurement { + micros: 12.5, + gb_per_s: 1234.0, + pct_of_roofline: 88.5, + max_abs_err: 1e-7, + }; + let parsed = parse_report(&format_ok(&measurement)).expect("ok report"); + assert!((parsed.micros - 12.5).abs() < 1e-9); + assert!((parsed.pct_of_roofline - 88.5).abs() < 1e-9); + } + + #[test] + fn failures_keep_their_kind_and_multi_line_message() { + let failure = KernelFailure::Compile("line 1\nline 2".to_string()); + let parsed = parse_report(&format_err(&failure)).expect_err("err report"); + assert_eq!(parsed.kind(), "nvrtc error"); + assert!(parsed.to_string().contains("line 2"), "{parsed}"); + } + + #[test] + fn a_wrong_answer_keeps_its_error_magnitude_and_is_labelled_once() { + let failure = KernelFailure::Wrong { + max_abs_err: 9.25e-3, + detail: "output[0][0] = 1, expected 2".to_string(), + }; + let parsed = parse_report(&format_err(&failure)).expect_err("err report"); + let text = parsed.to_string(); + assert_eq!(text.matches("wrong output").count(), 1, "{text}"); + assert!(text.contains("9.250e-3"), "{text}"); + } + + #[test] + fn a_report_that_never_arrived_is_a_launch_fault() { + let parsed = parse_report("").expect_err("empty report"); + assert!(parsed.poisons_context()); + } +} diff --git a/examples/cuda-softmax/src/lib.rs b/examples/cuda-softmax/src/lib.rs new file mode 100644 index 0000000..3329025 --- /dev/null +++ b/examples/cuda-softmax/src/lib.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: MPL-2.0 +//! Host side of the CUDA softmax example. +//! +//! The evolvable function returns a [`KernelPlan`] — CUDA C source plus the +//! geometry to launch it with — and everything that needs `unsafe`, a GPU, or +//! a driver handle lives here, behind [`Gpu`]. That split is what lets agent +//! code stay inside symbiont's policy (no `unsafe`, no statics, no FFI) while +//! still programming the GPU: the dylib emits *text* and typed launch +//! parameters, the host compiles, verifies, and measures them. + +#![allow( + unused_crate_dependencies, + reason = "symbiont, tokio and tracing are used by this package's binary target." +)] + +mod gpu; +mod isolate; + +pub use gpu::{ + DeviceInfo, + Gpu, + GpuError, + KernelFailure, + Measurement, +}; +pub use isolate::{ + EVAL_ENV, + Isolated, + ROOFLINE_ENV, + evaluator_main, +}; + +/// The kernel entry point the host looks up in every compiled module. +pub const KERNEL_NAME: &str = "softmax"; + +/// The C signature every candidate kernel must have. +pub const KERNEL_SIGNATURE: &str = + "extern \"C\" __global__ void softmax(const float* input, float* output, int rows, int cols)"; + +/// A CUDA kernel and the geometry to launch it with. +/// +/// The host compiles [`KernelPlan::source`] with NVRTC, looks up +/// [`KERNEL_NAME`], and launches it with the given grid/block/shared-memory +/// configuration. Both halves matter: the same source at the wrong block size +/// can be an order of magnitude slower, so the launch geometry is part of what +/// gets evolved rather than something the host picks. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct KernelPlan { + /// CUDA C source defining [`KERNEL_SIGNATURE`]. + /// + /// Compiled by NVRTC at evaluation time, so it may use any device-side + /// CUDA C++ the installed NVRTC understands: shared memory, warp shuffles, + /// vectorized `float4` loads, fast-math intrinsics such as `__expf`. + pub source: String, + + /// Grid dimensions, in blocks. + pub grid: (u32, u32, u32), + + /// Block dimensions, in threads. The product must not exceed the device's + /// maximum threads per block (1024 on every current architecture). + pub block: (u32, u32, u32), + + /// Dynamic shared memory per block, in bytes — the size that + /// `extern __shared__` arrays get. Leave at 0 when the kernel declares its + /// shared memory statically. + pub shared_bytes: u32, +} + +impl KernelPlan { + /// Total threads per block. + #[must_use] + pub fn threads_per_block(&self) -> u64 { + u64::from(self.block.0) * u64::from(self.block.1) * u64::from(self.block.2) + } + + /// Total blocks in the grid. + #[must_use] + pub fn blocks(&self) -> u64 { + u64::from(self.grid.0) * u64::from(self.grid.1) * u64::from(self.grid.2) + } +} + +/// Rows and columns of the benchmark matrix. +/// +/// One softmax per row over `cols` contiguous `f32`s. `cols` is a multiple of +/// 4, so `float4` loads are always legal. +pub const ROWS: usize = 4096; +/// See [`ROWS`]. +pub const COLS: usize = 1024; + +/// The naive kernel the evolvable function starts from: one thread per row. +/// +/// Correct, and about as slow as a softmax can reasonably be. Each thread +/// walks an entire row on its own, so the 32 threads of a warp touch addresses +/// `cols * 4` bytes apart and every load pulls in a fresh cache line to use +/// four bytes of it. The three passes then re-read the row from DRAM twice. +pub const NAIVE_KERNEL: &str = r#" +extern "C" __global__ void softmax(const float* input, float* output, int rows, int cols) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= rows) return; + const float* in = input + (size_t)row * cols; + float* out = output + (size_t)row * cols; + + float maximum = -INFINITY; + for (int i = 0; i < cols; ++i) { + maximum = fmaxf(maximum, in[i]); + } + float sum = 0.0f; + for (int i = 0; i < cols; ++i) { + sum += expf(in[i] - maximum); + } + for (int i = 0; i < cols; ++i) { + out[i] = expf(in[i] - maximum) / sum; + } +} +"#; + +/// Deterministic benchmark input: `rows * cols` values in `[bias, bias + 8]` +/// with a per-row `bias` of up to ~90. +/// +/// The bias is the point. `expf(90.0f)` overflows to `inf` in single +/// precision, so a kernel that skips the max-subtraction — a tempting way to +/// delete a whole pass over the row — produces `inf / inf = NaN` and is caught +/// by the correctness gate instead of winning the benchmark. +#[must_use] +pub fn benchmark_input(rows: usize, cols: usize) -> Vec { + let mut data = Vec::with_capacity(rows * cols); + // SplitMix64, inlined: a fixed stream keeps every run comparable without + // pulling in an RNG dependency. + let mut state = 0x2545_F491_4F6C_DD1D_u64; + let mut next = move || { + state = state.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = state; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + f64::from((z ^ (z >> 31)) as u32) / f64::from(u32::MAX) + }; + for _ in 0..rows { + let bias = next() * 90.0; + for _ in 0..cols { + data.push((bias + next() * 8.0) as f32); + } + } + data +} + +/// Row-wise softmax on the CPU, used as the correctness oracle. +/// +/// Accumulated in `f64` so the tolerance the kernels are held to measures +/// *their* error rather than the reference's. +#[must_use] +pub fn reference_softmax(input: &[f32], rows: usize, cols: usize) -> Vec { + let mut out = vec![0.0_f32; rows * cols]; + for row in 0..rows { + let src = &input[row * cols..(row + 1) * cols]; + let dst = &mut out[row * cols..(row + 1) * cols]; + let maximum = src.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let sum: f64 = src + .iter() + .map(|&v| (f64::from(v) - f64::from(maximum)).exp()) + .sum(); + for (d, &v) in dst.iter_mut().zip(src.iter()) { + *d = ((f64::from(v) - f64::from(maximum)).exp() / sum) as f32; + } + } + out +} + +/// Prelude imported by the generated dylib through [`symbiont::DylibConfig`]. +pub mod prelude { + pub use crate::{ + KERNEL_NAME, + KERNEL_SIGNATURE, + KernelPlan, + NAIVE_KERNEL, + }; +} diff --git a/examples/cuda-softmax/src/main.rs b/examples/cuda-softmax/src/main.rs new file mode 100644 index 0000000..8db2eb6 --- /dev/null +++ b/examples/cuda-softmax/src/main.rs @@ -0,0 +1,529 @@ +// SPDX-License-Identifier: MPL-2.0 +//! Evolving a CUDA kernel: the agent writes GPU code, the host measures it. +//! +//! The evolvable function returns a [`KernelPlan`] — CUDA C source plus the +//! grid/block geometry to launch it with. Each round the host compiles every +//! candidate with NVRTC, checks it against a CPU softmax oracle, times it, and +//! reports the achieved memory throughput as a percentage of the device's +//! measured copy bandwidth. The best kernel is activated; the rest are kept as +//! revisions. +//! +//! ## Why a GPU kernel is a good subject for an evolution loop +//! +//! The fitness function is not a matter of taste. A row-wise softmax has to +//! move `2 * rows * cols * 4` bytes and nothing else, so there is a hard +//! ceiling — measured here as a device-to-device copy of the same buffers — +//! and every candidate can be scored as a fraction of it. The naive kernel the +//! run starts from sits at a few percent of that ceiling, and closing the gap +//! is a matter of *strategy* (coalescing, block-per-row reductions, warp +//! shuffles, vectorized loads, fast-math intrinsics), not of micro-edits. That +//! is exactly the shape of problem [`Runtime::evolve_batch`] exists for: many +//! genuinely different candidates per round, each scored, only the winner +//! activated. +//! +//! ## What runs where +//! +//! Agent code stays inside symbiont's policy: no `unsafe`, no statics, no FFI. +//! It emits *text* and typed launch parameters. Everything that needs a device +//! pointer lives in the host crate behind `Gpu`. +//! +//! Candidate kernels are then launched in a **child process**, one per +//! candidate. Symbiont contains a misbehaving CPU implementation in-process by +//! wrapping the dylib's functions in `catch_unwind`; a GPU kernel cannot be +//! contained that way, because an illegal memory access is a sticky CUDA error +//! that leaves the whole process unable to create *any* context afterwards. +//! The process boundary is the GPU's `catch_unwind`. See +//! `cuda_softmax_example::Isolated`. +//! +//! ## Running +//! +//! ```bash +//! MODEL=... cargo run -p cuda-softmax-example --release +//! ``` +//! +//! `ROUNDS` (default 3) and `LANES` (default 4) size the search. Without a +//! CUDA device the example reports that and exits successfully, unless +//! `STRICT=1` is set. + +#![allow( + unused_crate_dependencies, + reason = "cudarc is used by this package's library target." +)] + +use std::fmt::Write; + +use cuda_softmax_example::{ + COLS, + DeviceInfo, + Gpu, + Isolated, + KERNEL_SIGNATURE, + Measurement, + ROWS, + prelude::*, +}; +use symbiont::{ + DylibConfig, + Revision, + Runtime, +}; +use tracing::{ + info, + warn, +}; + +symbiont::evolvable! { + /// Produce the CUDA kernel that computes a row-wise softmax, and the + /// geometry to launch it with. + /// + /// The host compiles `source` with NVRTC, looks up the `softmax` symbol, + /// and launches it once per measurement with + /// `(input, output, rows, cols)`. `input` and `output` are row-major + /// `rows * cols` `f32` matrices in device memory; `output[r][c]` must be + /// `exp(input[r][c] - max(row r)) / sum(exp(input[r][*] - max(row r)))`. + /// + /// Correctness is checked against a CPU oracle before any timing is + /// reported, and the input contains per-row biases up to ~90, so the + /// max-subtraction is not optional: `expf(90.0f)` is `inf` in single + /// precision. + fn plan(rows: usize, cols: usize) -> KernelPlan { + // One thread per row: correct, and about as slow as a softmax gets. + KernelPlan { + source: NAIVE_KERNEL.to_string(), + grid: ((rows as u32).div_ceil(256), 1, 1), + block: (256, 1, 1), + shared_bytes: 0, + } + } +} + +/// Output cap per lane. Kernels plus a short explanation fit comfortably; the +/// wall clock of a batch is its slowest lane, so an unbounded rambler holds up +/// the whole round. +const MAX_OUTPUT_TOKENS: u64 = 3072; + +/// One strategy hint per lane, appended last so every lane shares the same +/// (cacheable) prompt prefix. Each one points at a different optimization +/// axis, which is what makes the candidates worth comparing. +const STRATEGIES: &[&str] = &[ + "Assign one block per row and reduce the maximum and the sum across the \ + whole block — warp shuffles first, then combine the per-warp partials \ + through shared memory — so consecutive threads read consecutive columns.", + "Assign one warp (32 threads) per row and use only warp shuffles — no \ + shared memory, no __syncthreads.", + "Use vectorized float4 loads and stores; cols is a multiple of 4, so each \ + thread can process four contiguous columns per access.", + "Stage the row in shared memory on the first pass so the exponentials are \ + computed without re-reading global memory.", + "Use the online (single-pass) softmax formulation: keep a running maximum \ + and a running sum, rescaling the sum by exp(old_max - new_max) whenever \ + the maximum grows.", + "Use the fast-math intrinsics __expf and __fdividef, and multiply by a \ + precomputed reciprocal instead of dividing per element.", + "Pick the block size from the row length and maximize occupancy; make \ + sure every global load is coalesced across the warp.", + "Combine float4 vectorized loads with a warp-shuffle reduction and a \ + single fused write pass.", +]; + +/// The best candidate so far, and what it achieved. +struct Champion { + revision: Revision, + plan: KernelPlan, + measurement: Measurement, + label: String, +} + +/// How one lane of a round turned out. +struct LaneOutcome { + lane: usize, + revision: Option, + plan: Option, + measurement: Option, + /// Rejection reason, phrased for the next round's prompt. + failure: Option, + kind: &'static str, +} + +impl LaneOutcome { + fn rejected(lane: usize, revision: Option, kind: &'static str, why: String) -> Self { + Self { + lane, + revision, + plan: None, + measurement: None, + failure: Some(why), + kind, + } + } +} + +/// Call the candidate's `plan`, then compile, verify and time its kernel in a +/// child process. +/// +/// The `plan` call goes through a [`symbiont::RevisionFn`] handle, which pins +/// its own revision: candidates are measured without ever becoming the active +/// implementation, so a round can be evaluated in any order and nothing is +/// hot-swapped between measurements. +fn evaluate_lane(isolated: &Isolated, lane: usize, revision: Revision) -> LaneOutcome { + let Some(handle) = plan_fn(revision) else { + return LaneOutcome::rejected(lane, None, "unregistered", "revision not registered".into()); + }; + + let plan = handle.get()(ROWS, COLS); + if let Some(panic) = handle.take_panic() { + return LaneOutcome::rejected( + lane, + Some(revision), + "panicked", + format!("`plan` panicked: {panic}"), + ); + } + + match isolated.evaluate(&plan) { + Ok(measurement) => LaneOutcome { + lane, + revision: Some(revision), + plan: Some(plan), + measurement: Some(measurement), + failure: None, + kind: "ok", + }, + Err(failure) => LaneOutcome { + lane, + revision: Some(revision), + plan: Some(plan), + measurement: None, + failure: Some(failure.to_string()), + kind: failure.kind(), + }, + } +} + +/// The part of the prompt that never changes: task, contract, device. +/// +/// Kept first and byte-identical across every lane and every round so a server +/// with prefix caching prefills it once for the whole run. +fn task_preamble(signature: &str, info: &DeviceInfo, roofline_gb_per_s: f64) -> String { + format!( + "You are optimizing a CUDA kernel. Implement this Rust function:\n\ + ```rust\n{signature}\n```\n\ + It returns a `KernelPlan {{ source, grid, block, shared_bytes }}`: the CUDA C source \ + of the kernel, and the geometry the host launches it with.\n\n\ + The kernel must have exactly this signature:\n\ + ```cuda\n{KERNEL_SIGNATURE}\n```\n\ + `input` and `output` are row-major {ROWS}x{COLS} f32 matrices in device memory. For \ + every row r and column c:\n\ + `output[r][c] = expf(input[r][c] - rowmax) / sum_c expf(input[r][c] - rowmax)`, where \ + `rowmax` is the maximum of row r.\n\n\ + Rules:\n\ + - The output is compared against a CPU reference before any timing is reported. Rows \ + contain per-row biases up to ~90, so subtracting the row maximum is mandatory: \ + `expf(90.0f)` is `inf` in f32.\n\ + - Every one of the {ROWS} rows must be written. `cols` is {COLS}, a multiple of 4, so \ + float4 loads are legal and aligned.\n\ + - The kernel must be `extern \"C\"` or the host cannot find the symbol.\n\ + - Out-of-bounds accesses are detected and cost the whole candidate; guard your indices.\n\ + - The comparison is strict (1e-6 absolute). The usual cause of a wrong answer is an \ + incomplete reduction: __shfl_down_sync only reduces within one warp, so a block of \ + more than 32 threads must combine the per-warp partials through shared memory before \ + the maximum or the sum is valid for the whole row.\n\ + - Every thread of a block must reach every __syncthreads(). Returning early from \ + out-of-range threads before a barrier deadlocks the block.\n\ + - Return the launch geometry that matches your kernel. A block-per-row kernel wants \ + `grid = (rows, 1, 1)`; a thread-per-row kernel wants `grid = (rows / block, 1, 1)`.\n\ + - Set `shared_bytes` only for `extern __shared__` arrays; statically sized \ + `__shared__` declarations need no dynamic allocation.\n\n\ + Target device: {} (compute capability {}.{}, {} SMs, up to {} threads and {} bytes of \ + shared memory per block, warp size {}).\n\ + This kernel is memory bound: it must read {ROWS}x{COLS} floats and write as many. A \ + device-to-device copy of that same traffic runs at {roofline_gb_per_s:.0} GB/s on this \ + card, which is the practical ceiling.\n", + info.name, + info.compute_capability.0, + info.compute_capability.1, + info.multiprocessors, + info.max_threads_per_block, + info.max_shared_memory_per_block, + info.warp_size, + ) +} + +/// The part of the prompt that changes between rounds: the incumbent and what +/// went wrong last time. +fn round_state(champion: &Champion, feedback: &str) -> String { + let m = &champion.measurement; + format!( + "\nCurrent best kernel ({}): {:.0} GB/s, {:.1}% of the copy ceiling, {:.1} us per \ + launch, launched with grid {:?} block {:?} shared {} bytes:\n```cuda\n{}\n```\n\ + Beat it. A correct kernel that is slower than this one is not an improvement.\n{feedback}", + champion.label, + m.gb_per_s, + m.pct_of_roofline, + m.micros, + champion.plan.grid, + champion.plan.block, + champion.plan.shared_bytes, + champion.plan.source.trim(), + ) +} + +/// Summarize the previous round's rejections so the next one does not repeat +/// them. Truncated per lane: NVRTC logs can run to dozens of lines. +fn feedback_from(outcomes: &[LaneOutcome]) -> String { + let mut text = String::new(); + for outcome in outcomes.iter().filter(|o| o.failure.is_some()) { + let why = outcome.failure.as_deref().unwrap_or_default(); + let why: String = why.lines().take(6).collect::>().join(" | "); + let _ = writeln!( + text, + "A previous attempt was rejected ({}): {}", + outcome.kind, + why.chars().take(400).collect::() + ); + } + if text.is_empty() { + String::new() + } else { + format!("\nRejected attempts from the last round — do not repeat these:\n{text}") + } +} + +/// Print one round's leaderboard. +fn print_report(outcomes: &[LaneOutcome], baseline: &Measurement) { + println!("\n| Lane | Rev | Status | us | GB/s | % ceiling | Speedup"); + println!("|------|-----|----------------|--------|--------|-----------|--------"); + for outcome in outcomes { + let rev = outcome + .revision + .map_or_else(|| "-".to_string(), |r| r.to_string()); + match &outcome.measurement { + Some(m) => println!( + "| {:<4} | {rev:<3} | {:<14} | {:>6.1} | {:>6.0} | {:>8.1}% | {:>6.1}x", + outcome.lane, + "ok", + m.micros, + m.gb_per_s, + m.pct_of_roofline, + baseline.micros / m.micros, + ), + None => println!( + "| {:<4} | {rev:<3} | {:<14} | - | - | - | -", + outcome.lane, outcome.kind, + ), + } + } + for outcome in outcomes { + if let Some(failure) = &outcome.failure { + let first = failure.lines().next().unwrap_or_default(); + println!(" lane {}: {first}", outcome.lane); + } + } +} + +/// Read a positive `usize` from the environment, or fall back to `default`. +fn env_usize(key: &str, default: usize) -> usize { + std::env::var(key) + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(default) +} + +/// Bring up the GPU, or explain why the example is skipping. +fn open_gpu() -> Option { + match Gpu::new(ROWS, COLS) { + Ok(gpu) => Some(gpu), + Err(err) => { + let strict = std::env::var_os("STRICT").is_some(); + assert!( + !strict, + "no usable CUDA device: {err} (STRICT was requested)" + ); + println!( + "Skipping: {err}\nThis example needs an NVIDIA GPU; everything else about it \ + (including `cargo check`) works without one." + ); + None + } + } +} + +fn main() -> symbiont::Result<()> { + // A child process staged by `Isolated` evaluates one kernel and exits; it + // must never fall through into the evolution loop. Checked before anything + // else, including the tokio runtime it would have no use for. + if let Some(code) = cuda_softmax_example::evaluator_main() { + std::process::exit(code); + } + search() +} + +#[expect( + clippy::too_many_lines, + reason = "One linear walkthrough of the search — device, baseline, rounds, winner — which is what the example exists to show" +)] +#[tokio::main] +async fn search() -> symbiont::Result<()> { + symbiont::init_tracing(); + + // Probe the device before spending anything on inference or compilation. + // The parent only ever runs the host's own copy kernel, so its context + // stays healthy for the whole run. + let Some(gpu) = open_gpu() else { + return Ok(()); + }; + let isolated = Isolated::new() + .expect("can stage evaluation child processes") + // One ceiling for the whole run, measured here, so every candidate's + // percentage refers to the same number. + .with_roofline(gpu.roofline_gb_per_s()); + let info = gpu.info().clone(); + println!( + "Device: {} (cc {}.{}, {} SMs). Benchmark: {ROWS}x{COLS} f32 softmax, {:.1} MiB of \ + traffic per launch.\nCopy ceiling: {:.0} GB/s.", + info.name, + info.compute_capability.0, + info.compute_capability.1, + info.multiprocessors, + gpu.traffic_bytes() as f64 / (1024.0 * 1024.0), + gpu.roofline_gb_per_s(), + ); + + // Debug: the evolved Rust is a few lines that build a string, so there is + // nothing for the optimizer to do — all the performance is in the CUDA + // text, which NVRTC compiles at evaluation time regardless. + let host_crate = env!("CARGO_PKG_NAME"); + let runtime = Runtime::new( + SYMBIONT_DECLS, + SYMBIONT_PRELUDE, + DylibConfig::host_package( + symbiont::Profile::Debug, + host_crate, + env!("CARGO_MANIFEST_DIR"), + ), + ) + .await?; + + let model = std::env::var("MODEL").expect("the MODEL env var names the model slug"); + let agent = symbiont::agent_builder_from_env(Some(host_crate), &model) + .await? + .max_tokens(MAX_OUTPUT_TOKENS) + .build(); + + // -- Baseline --------------------------------------------------------- + let baseline_outcome = evaluate_lane(&isolated, 0, Revision::INITIAL); + let (Some(baseline), Some(baseline_plan)) = + (baseline_outcome.measurement, baseline_outcome.plan) + else { + panic!( + "the naive kernel must be correct: {:?}", + baseline_outcome.failure + ); + }; + println!( + "\nBaseline (one thread per row): {:.1} us, {:.0} GB/s, {:.1}% of the copy ceiling.", + baseline.micros, baseline.gb_per_s, baseline.pct_of_roofline + ); + + let mut champion = Champion { + revision: Revision::INITIAL, + plan: baseline_plan, + measurement: baseline, + label: "baseline".to_string(), + }; + + let preamble = task_preamble(&runtime.fn_sigs()[0], &info, gpu.roofline_gb_per_s()); + let rounds = env_usize("ROUNDS", 3); + let lanes = env_usize("LANES", 4); + let mut feedback = String::new(); + + // -- Search ----------------------------------------------------------- + for round in 1..=rounds { + let state = round_state(&champion, &feedback); + let prompts = Vec::from_iter(STRATEGIES.iter().cycle().take(lanes).map(|hint| { + format!("{preamble}{state}\nStrategy for this attempt: {hint}\nCode only.") + })); + + println!("\n=== Round {round}/{rounds}: {lanes} candidate kernels ==="); + let started = std::time::Instant::now(); + let results = runtime.evolve_batch(&agent, &prompts).await; + println!( + "Generated, validated and compiled {} candidates in {:.1}s.", + results.len(), + started.elapsed().as_secs_f64() + ); + + let mut outcomes = Vec::with_capacity(results.len()); + for (lane, result) in results.into_iter().enumerate() { + let outcome = match result { + Ok(revision) => evaluate_lane(&isolated, lane, revision), + Err(err) => LaneOutcome::rejected(lane, None, "no code", err.to_string()), + }; + outcomes.push(outcome); + } + + print_report(&outcomes, &baseline); + feedback = feedback_from(&outcomes); + + // -- Commit to the winner, if it is one --------------------------- + let improvement = outcomes + .iter() + .filter_map(|o| Some((o, o.measurement?))) + .filter(|(_, m)| m.micros < champion.measurement.micros) + .min_by(|a, b| a.1.micros.total_cmp(&b.1.micros)); + if let Some((outcome, measurement)) = improvement { + let revision = outcome.revision.expect("a measured lane has a revision"); + runtime.activate_revision(revision)?; + println!( + "Activated revision {revision} from lane {}: {:.1} us ({:.1}x the baseline, \ + {:.1}% of the ceiling).", + outcome.lane, + measurement.micros, + baseline.micros / measurement.micros, + measurement.pct_of_roofline, + ); + champion = Champion { + revision, + plan: outcome.plan.clone().expect("a measured lane has a plan"), + measurement, + label: format!("round {round}, lane {}", outcome.lane), + }; + } else { + info!( + "round {round} produced no improvement; keeping revision {}", + champion.revision + ); + } + } + + // -- Result ----------------------------------------------------------- + println!( + "\n=== Best kernel: {} ===\n{:.1} us, {:.0} GB/s, {:.1}% of the {:.0} GB/s copy ceiling, \ + {:.1}x faster than the naive baseline.\n```cuda\n{}\n```", + champion.label, + champion.measurement.micros, + champion.measurement.gb_per_s, + champion.measurement.pct_of_roofline, + gpu.roofline_gb_per_s(), + baseline.micros / champion.measurement.micros, + champion.plan.source.trim(), + ); + + // Dispatch now runs the winner: a plain call, not a pinned handle, so this + // really is the hot-swapped implementation. + let active = plan(ROWS, COLS); + assert_eq!( + active, champion.plan, + "the active revision must be the champion" + ); + match isolated.evaluate(&active) { + Ok(m) => println!( + "Re-measured through the active dispatch pointer: {:.1} us.", + m.micros + ), + Err(err) => warn!("the activated kernel failed on re-measurement: {err}"), + } + + Ok(()) +} diff --git a/examples/cuda-softmax/tests/isolation.rs b/examples/cuda-softmax/tests/isolation.rs new file mode 100644 index 0000000..0059c82 --- /dev/null +++ b/examples/cuda-softmax/tests/isolation.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: MPL-2.0 +//! End-to-end test of the process isolation that makes the search survivable. +//! +//! Drives the real child protocol against the example's own binary: a kernel +//! that writes out of bounds takes a child process down with it, and the +//! parent goes on to measure the next candidate as if nothing happened. That +//! is the property the whole `Isolated` design exists for, and it cannot be +//! observed from inside a single process — once a sticky CUDA error lands, +//! everything in that process stays broken. + +#![allow( + unused_crate_dependencies, + reason = "This test only exercises the isolation protocol of the library target." +)] + +use cuda_softmax_example::{ + Isolated, + KernelPlan, + NAIVE_KERNEL, + ROWS, +}; + +/// The example's binary, which dispatches to `evaluator_main` when staged as +/// a child. `current_exe` would point at this test harness instead. +const EVALUATOR: &str = env!("CARGO_BIN_EXE_cuda-softmax-example"); + +fn naive_plan() -> KernelPlan { + KernelPlan { + source: NAIVE_KERNEL.to_string(), + grid: (u32::try_from(ROWS.div_ceil(256)).expect("fits"), 1, 1), + block: (256, 1, 1), + shared_bytes: 0, + } +} + +fn out_of_bounds_plan() -> KernelPlan { + KernelPlan { + source: r#" +extern "C" __global__ void softmax(const float* input, float* output, int rows, int cols) { + output[(size_t)(blockIdx.x + 1) * 1000000000ull] = 1.0f; +} +"# + .to_string(), + ..naive_plan() + } +} + +#[test] +fn a_faulting_kernel_costs_a_child_process_and_nothing_else() { + let isolated = Isolated::with_executable(EVALUATOR).expect("scratch dir"); + + // Establish that this machine can run the benchmark at all; without a GPU + // the child reports a device error and there is nothing to test. + match isolated.evaluate(&naive_plan()) { + Ok(measurement) => assert!(measurement.gb_per_s > 0.0), + Err(failure) => { + eprintln!("skipping: {failure}"); + return; + } + } + + let failure = isolated + .evaluate(&out_of_bounds_plan()) + .expect_err("an out-of-bounds write must not be reported as a measurement"); + assert!( + failure.poisons_context(), + "a dead child is a launch fault: {failure}" + ); + + // The point: the parent is untouched and the search continues. + let after = isolated + .evaluate(&naive_plan()) + .expect("the parent can still evaluate candidates after a fault"); + assert!(after.gb_per_s > 0.0); +}