cuda-oxide: eliding bounds checks with proof carrying views

.md

Bounds-check elision,

gemm vs gemm_views

cargo oxide run gemm            # the baseline + its GFLOPS line
cargo oxide run gemm_views      # correctness, contract rejection, bench table
cargo oxide run gemm_views -- --verify-ptx

Numbers to remember: 2,942 -> 7,159 GFLOPS (2.43x), safety cost ~0.1%, results bit-identical.

---

Intro

• lets run run cargo oxide run gemm. Look at Performance: 0.730 ms, 2942 GFLOPS.
• the inner loop in gemm/src/main.rs: a[row*k+i] * b[i*n+col]. Looks normal. It's safe Rust on a GPU, and most of that runtime is the safety tax.

Hidden safety tax

• Look at gemm.ptx, you’ll see the inner loop:
  • two setp + guarded bra to trap;
  • blocks per fma.
  • Every a[i] is "is i inside the slice? if not, stop the kernel", per access, per thread.
• The checks also split the loop into multiple basic blocks, which kills unrolling. Roughly 13 instructions per fma; ~5 are check overhead.
• Why LLVM can't fix it: a.len() is an opaque parameter. Nothing in the IR connects m * k to it. You know the buffers are big enough; that fact lives in the host program. Everything that follows is about handing that fact to the compiler.

Part 1: prove it instead of checking it

Lets see sgemm_naive_views next to sgemm_naive_raw in gemm_views.

• One slide of mental model: a matrix is one flat array, stride = row width, element (row, col) lives at row * stride + col.
• a_mat.row(row, k): one check covers A's whole row. b_mat.col(col, k): one check covers B's whole column (elements a full row width apart).
• a_row.zip_exact(b_col): one length check; after that the dot loop's only compare is its own exit test. Same six instructions as the raw twin: load, load, fma, advance, compare, branch.
• The C write: tile_2d32_rt(coord, n) with a 1x1 tile. "Tile" just means "the rectangle this thread owns"; naive GEMM's rectangle is one cell. One documented unsafe obligation: every thread passes the same row width. n is a kernel argument, so that holds by construction.
• The contract line above the kernel: requires = (k >= 1, a.len() >= m * k, b.len() >= k * n, c.len() >= m * n). Sizes are now part of the kernel's interface, checked once per launch on the CPU. Point at the run output: the deliberately oversized-K launch is rejected with the relation text, before the GPU does anything.

Proof,

• The gemm_views run: matches CPU reference, safe vs raw bit-identical, then the bench table: 7,159 vs 7,161 GFLOPS. Safety costs ~0.1%. Against gemm's 2,942: 2.43x.
• Run --verify-ptx: the harness pins same branch counts, same load/store instruction multisets, zero traps, no leaked markers. The safe kernel doesn't just perform like the unsafe one; it compiles to the same machine code.

Part 2: the escape hatch (+ a compiler nugget)

• #[kernel(unchecked_indexing)]: for shapes the views can't express. Deletes the indexing checks outright; the contract is get_unchecked's (out of bounds = UB). Only indexing; overflow and div-by-zero still trap. And it never leaks: a kernel calling an opted-in kernel's helper keeps its own checks (regression-tested at the PTX level).
• Nugget if time allows: the safe loop initially refused to unroll. The iterator's Option<(f32, f32)> merged as an aggregate PHI, which LLVM cannot split (SROA only handles memory). New mir-lower pass scalarize_block_args splits those into scalar PHIs; that's the moment safe PTX became identical to raw.

Interesting design notes

• requires is deliberately tiny:
  • one comparison per relation, .len(), unsigned scalars, + - *, literals.
  • being tiny makes it trustworthy: every name is validated against the kernel signature (typos = compile error) and all arithmetic is checked u64 (no wraparound false-passes).
  • Natural extensions: divisibility (n % 16 == 0 would delete grid-edge guards), alignment requirements, and ultimately feeding the proof to the device so even the per-thread view Option checks disappear.
• Same trajectory for tile_2d32_rt's one unsafe: kernel scalars are uniform across threads by construction, so a macro-minted Uniform<u32> witness type would make it fully safe.
• A limitation: partial edge tiles. Tile types have a fixed shape; if the rectangle doesn't fully fit (a 100x100 matrix with 16x16 tiles leaves 4-wide strips at the edges), the constructor returns None rather than a clipped tile. Today: pick sizes that divide, use 1x1 tiles (naive GEMM sidesteps this entirely), or guard edges per element. A clipped-tile type is future work.
• Close on the ladder: check every access -> check once per thread -> check once per launch. Each rung is the same check, done less often, and the type system is what lets you climb down without losing safety.