CUDA kernel audit — correctness, dead code, performance

CUDA kernel audit — correctness, dead code, performance#

Status: in progress Scope: full sweep of the 46 .cu files and 7 headers shipped as package data, looking for (a) memory-safety and correctness defects, (b) dead code the previous cleanup missed, and (c) performance headroom.

Predecessor: cuda-operator-cleanup.md, which removed 102 unreachable @BE entry points, 2 orphan headers and 8 unused cuda_common.h symbols. That pass audited reachability; it did not look inside the kernels.

Validation environment#

RTX 3080 Ti Laptop GPU (sm_86), driver 596.49, CUDA 13.2 / nvcc 13.1, jax 0.11.0 with jaxlib CUDA. compute-sanitizer is unavailable on this host (WSL2/WDDM: “Failed to initialize WDDM debugger interface / Device not supported”), so memory-safety findings are established with purpose-built canary kernels and standalone harnesses instead.

Validation pitfall (cost an hour, worth recording)#

brainevent is installed in editable mode pointing at the main checkout. A script run as python /tmp/foo.py puts /tmp — not the CWD — on sys.path[0], so it imports /mnt/d/codes/projects/brainevent/brainevent, not the worktree. Early “the fix does not work” results were actually the unfixed main checkout being re-tested. Every GPU validation command must pass:

PYTHONPATH=<worktree-root> python ...

and the harness should print brainevent.__file__ to prove which tree it exercised. Deleting ~/.cache/brainevent between runs is also required, because compiled modules are keyed on a content hash of the .cu source plus flags.

Finding 1 — CSRMM block kernels under-request dynamic shared memory (memory fault)#

Severity: high. float64 binary_csrmm aborts the CUDA context with CUDA_ERROR_ILLEGAL_ADDRESS; float32/float16/bfloat16 silently perform an out-of-bounds shared-memory write that only survives by allocation-granularity luck.

The defect#

_csrmm_nt_block_{homo,hetero}_kern (and the _perm_hetero variant in the indexed file) stage one accumulator per (strip, lane) pair:

smem[strip * 32 + lane] = acc0 + acc1;   // strip in 0..7, lane in 0..31 -> 256 slots
__syncthreads();
if (strip == 0 && c < n) {
    for (int s = 0; s < 8; s++) sum += smem[s * 32 + lane];

The kernels launch with 256 threads, so this needs 8 * 32 * sizeof(ACC_T) bytes. Both files instead requested 8 * sizeof(ACC_T):

File

Requested

Required

Instantiations

_csr/binary_csrmm.cu

8 * sizeof(T)

8 * 32 * sizeof(T)

16

_csr/binary_indexed_csrmm.cu

8 * sizeof(T)

8 * 32 * sizeof(T)

10

_csr/float_csrmm.cu — the same kernel shape — already had the correct 8 * 32 * sizeof(...), which is what identifies this as a dropped * 32 rather than a different design. The CSRMV block kernels reduce into smem_red[warpid] (8 slots), so their 8 * sizeof(T) is correct and was left alone.

Why the bug is dtype-dependent#

A canary kernel that requests shm bytes but writes nf floats per block, run with 65536 blocks so the SM is saturated, shows what the driver actually tolerates on sm_86:

request=32 bytes, 256 thr/block, 65536 blocks
  wrote   128 floats (   512 bytes): no error     corrupted=0
  wrote   256 floats (  1024 bytes): no error     corrupted=0
  wrote   512 floats (  2048 bytes): an illegal memory access was encountered

The per-block shared-memory window is rounded up to a granule large enough that a 1024-byte overrun lands harmlessly, but 2048 bytes faults. That is exactly the split observed:

  • ACC_T = float (f32/f16/bf16) writes 256 × 4 = 1024 bytes → survives on this driver, still formally out of bounds and not portable.

  • ACC_T = double (f64) writes 256 × 8 = 2048 bytes → faults.

Reproduction through the public API#

binary_csrmm(..., transpose=False, backend='cuda_raw') with float64 weights and avg_nnz > 512 (the threshold at which nt_auto selects the block kernel). Same script, same GPU, only the checkout differs:

main:                                       worktree (fixed):
  homo   float32 nt_warp : PASS               homo   float32 nt_warp : PASS
  homo   float32 nt_block: PASS               homo   float32 nt_block: PASS
  hetero float32 nt_warp : PASS               hetero float32 nt_warp : PASS
  hetero float32 nt_block: PASS               hetero float32 nt_block: PASS
  homo   float64 nt_warp : PASS               homo   float64 nt_warp : PASS
  homo   float64 nt_block: CRASH              homo   float64 nt_block: PASS
    CUDA_ERROR_ILLEGAL_ADDRESS                hetero float64 nt_warp : PASS
    (context destroyed, run aborts)           hetero float64 nt_block: PASS

Independently, a standalone harness that #includes the real binary_csrmm.cu and calls binary_csrmm_nt_auto_hetero_f64_bool with hand-built BE::Tensor views reports no error and max abs err = 0 after the fix, and an illegal memory access was encountered before it.

Fix#

Both files now pass 8 * 32 * sizeof(...), matching float_csrmm.cu.

Regression test#

brainevent/_csr/binary_test.py::test_binary_csrmm_nt_block_shared_memory_is_large_enough (@requires_gpu_backend, @pytest.mark.slow, parametrised over homo/hetero) builds a CSR matrix with 1024 nnz per row so nt_auto routes to the block kernel, and compares the float64 cuda_raw result against jax_raw. Verified to fail (illegal memory access) with the fix stashed and pass with it applied.

Finding 2 — 52 unreachable @BE entry points remain#

The predecessor pass reduced 749 → 647 annotations. A fresh reachability analysis — extract every // @BE <name>, extract every kernel_name f-string template from non-test Python, expand {...} to [A-Za-z0-9_]*, and match — leaves 52 entry points that no Python code path can name.

Python only ever composes these CSR binary targets: binary_csrm{v,m}_nt_auto{_homo,_hetero}, binary_csrm{v,m}_nt_auto_perm_hetero, binary_csrm{v,m}_{wat,sraw}_hybrid* and binary_indexed_csrm{v,m}_{wat,sraw}_hybrid_hetero*. The transpose route moved to the hybrid kernels, which stranded the whole t_warp family:

File

Dead entry points

Count

_csr/binary_csrmv.cu

binary_csrmv_t_warp_{homo,hetero}_*

16

_csr/binary_csrmm.cu

binary_csrmm_t_warp_{homo,hetero}_*

16

_csr/binary_indexed_csrmv.cu

binary_csrmv_t_warp_perm_hetero_*

8

_csr/binary_indexed_csrmm.cu

binary_csrmm_t_warp_perm_hetero_*

8

_csr/binary_indexed_csrmm.cu

binary_csrmm_nt_{warp,block}_perm_hetero_f32_*

4

Removing them also retires the _t_warp_*_kern device kernels, which nothing else launches. Those kernels additionally call atomicAdd directly on __half* / __nv_bfloat16* rather than through the arch-guarded atomic_add_f16 / atomic_add_bf16 helpers, so they would fail to compile below sm_70 / sm_80 — retiring them removes that portability trap too.

binary_indexed_test.py::test_indexed_mm_cuda_kernel_selects_perm_names asserts the dead names and their DEFINE_ macros are present in the .cu text; those assertions pin dead code and are updated with the removal.

Finding 3 — CSRMM dense subscripts overflow 32-bit int (correctness)#

The CSRMM kernels index the dense operands as B[indices[j] * n + c] and C[row * n + c]. Both factors are runtime int, so the product is computed in 32-bit and wraps once k * n (or m * n) exceeds 2^31. Demonstrated standalone at k * n = 2148507648: int yields -2146460672, size_t yields 2148506624 — a wild pointer, not a slow path.

float_csrmm.cu had the same exposure. 30 subscripts across the four CSRMM sources were widened by inserting a (size_t) cast on the leading factor; the smem strides (smem[strip * 32 + lane]) were left alone since both factors there are block-bounded constants. Full GPU CSR suite: 699 passed.

A tree-wide re-audit then found the same pattern in the JIT connectivity families: chunk_counts[row * n_chunks + chunk_id] and chunk_offsets[...] in _jit_{normal,scalar,uniform}/{csr,dt2t}.cu (24 sites). n_chunks defaults to 4, but chunk_size is user-settable, so row * n_chunks is not bounded by anything the kernel controls. Widened the same way. JIT GPU suites: 1512 passed, 4 skipped.

Reachability for both is the same and worth stating plainly: the index only overflows once the indexed buffer itself exceeds 2^31 elements, i.e. ≥8.6 GB at 4 bytes (4.3 GB at fp16). That is out of reach on a small card and reachable on an 80 GB one. The fix costs nothing at runtime — the multiply widens to 64-bit, and none of these sites are in an inner loop — so it is worth taking rather than documenting as a limit.

After this pass the only remaining unwidened multiplicative subscripts tree-wide are shared-memory strides with a literal * 32, bounded by block size by construction.

Finding 4 — CSRMV was missing its warp-per-row tier (performance)#

The defect#

float_csrmv.cu documents and implements a three-tier row mapping — one thread per row below avg_nnz 8, one warp per row up to 512, one block per row above. The three binary CSRMV sources only had two tiers:

if (avg_nnz <= 512) { thread-per-row } else { block-per-row }

so every row length from 16 to 512 — the normal range for sparse neural connectivity — ran the thread-per-row kernel, where the 32 lanes of a warp each walk a different row and every indices[j] load is uncoalesced.

Git history shows DEFINE_CSRMV_NT_WARP_HOMO existed at da4a812~1. The predecessor cleanup deleted it as unreachable, which was true but was the wrong repair: the dispatcher was what needed fixing.

Measured, on sm_86, m = k = 65536, float32 hetero, bool spikes, 50 iterations#

avg_nnz

thread (ms)

warp (ms)

block (ms)

best

2

0.009

0.026

0.219

thread

8

0.017

0.034

0.278

thread

12

0.030

0.034

0.278

thread

16

0.040

0.034

0.276

warp

32

0.088

0.043

0.271

warp

64

0.209

0.076

0.272

warp

128

0.361

0.141

0.249

warp

256

0.950

0.277

0.325

warp

512

1.904

0.547

0.545

block

Up to 3.4x on the tier that was missing. The crossovers give the thresholds used: thread below 16, warp to 512, block above.

Launch shape: pack warps, don’t launch 32-thread blocks#

The historical kernel (and float_csrmv.cu today) launched <<<m, 32>>> — one warp per block. A block occupies a scheduler slot regardless of size, and an SM caps resident blocks, so 32-thread blocks waste 7/8 of each slot. Packing 8 warps into a 256-thread grid-strided block instead:

avg_nnz

<<<m, 32>>> (ms)

8 warps/block (ms)

speedup

8

0.111

0.034

3.31x

16

0.111

0.034

3.25x

32

0.111

0.044

2.53x

64

0.160

0.076

2.11x

128

0.145

0.143

1.02x

256

0.276

0.277

1.00x

512

0.545

0.550

0.99x

Strictly better at low row lengths, parity above. The same shape applied to the transpose atomic-scatter kernel (csrmv_t_warp, launched unconditionally): 2.80x at avg_nnz 8, 1.81x at 16, parity through 256, and 0.94x at 512 — net strongly positive.

Fix#

  • Added DEFINE_CSRMV_NT_WARP_{HOMO,HETERO} to binary_csrmv.cu and DEFINE_CSRMV_NT_WARP_PERM_HETERO to binary_indexed_csrmv.cu, instantiated for all weight/spike dtype pairs (24 new device kernels, no new @BE entry points — the nt_auto wrappers dispatch to them internally).

  • Rewrote all three dispatchers as the three-tier ladder.

  • Converted the two existing float_csrmv.cu warp kernels from <<<m, 32>>> to the packed grid-strided shape.

  • Added be_csrmv_warp_grid() to cuda_common.h (shared by all three sources; clamps to [1, 4096] so an empty matrix cannot produce a zero-sized grid).

row is derived as blockIdx.x * warps_per + warp_id, making it warp-uniform. That is what keeps the __shfl_down_sync reduction converged — the M17 precondition documented in cuda_common.h. A lane-dependent row bound would silently corrupt the reduction.

Verification#

96 configurations through the public API (7 row lengths spanning all three tiers x 3 weight dtypes x homo/hetero x bool/float spikes), cuda_raw against jax_raw:

  • float32 and float64: agree to 1e-5 / 1e-12 relative on every tier.

  • float16: differs from jax_raw by up to 3e-2 relative — including on the block tier, which this change does not touch. Resolved against an exact float64 ground truth: the CUDA path is 10-100x closer to exact than jax_raw in all 8 cases (e.g. 2.4e-4 vs 3.8e-2 at avg_nnz 1024). The kernels accumulate in float32; the reference accumulates in float16. This is a reference-precision artifact, not a kernel defect.

End-to-end Python timing could not resolve the improvement on this machine: JAX host dispatch latency floors at ~1.2 ms for this call, well above the 0.03-0.5 ms of kernel time, so the API-level A/B returns noise. The standalone CUDA benchmarks above are the measurement of record.

Regression tests#

  • test_binary_csrmv_matches_reference_across_dispatch_tiers — 7 row lengths straddling both thresholds x homo/hetero.

  • test_binary_csrmv_warp_tier_handles_ragged_and_empty_rows — empty rows, a row count off the warps-per-block tiling, and wildly varying row lengths; also asserts empty rows are written rather than left uninitialised.

  • test_binary_csrmv_cuda_source_keeps_three_dispatch_tiers — source-level, runs without a GPU. Necessary because a two-tier dispatcher is still correct, so no numeric test can catch the tier being dropped again. Verified to fail against the pre-change source.

Finding 5 — float_csrmm.cu warp kernels kept the unpacked launch shape#

binary_csrmm.cu and binary_indexed_csrmm.cu already launch their warp kernels as CSRMM_WARP_RPB (4) rows per 128-thread block with a grid-strided loop and a capped grid.x. float_csrmm.cu was never updated to match and still launched <<<dim3(m, c_blocks), 32>>> — one warp per block — for both its NT and transpose warp kernels. Same defect as Finding 4’s launch shape, in the file that is otherwise the reference implementation.

Measured on sm_86, m = k = 16384, n = 64, float32 homogeneous:

avg_nnz

<<<m, 32>>> (ms)

packed 4/block (ms)

speedup

2

0.062

0.030

2.05x

4

0.055

0.031

1.77x

8

0.059

0.038

1.58x

16

0.088

0.062

1.42x

32

0.142

0.111

1.28x

64

0.251

0.211

1.19x

A win across the entire range the warp tier serves (avg_nnz <= 64).

Both kernels were converted to the packed grid-strided form. The row guard becomes continue rather than return, since returning inside a grid-strided loop would abandon that warp’s remaining rows — the same trap as in Finding 4.

CSRMM_MAX_GRID_X and CSRMM_WARP_RPB were duplicated verbatim in two .cu files; rather than add a third copy they were moved to cuda_common.h alongside the CSRMV geometry and the local definitions deleted.

Finding 6 — the last five one-warp-per-block launches#

After Findings 4 and 5, five <<<…, 32>>> launches remained: the three warp kernels in slice_csr_slice_rows.cu (forward homo, forward hetero, backward) and the two in dt2t.cu (csrmv and csrmm row-warp). All five are the same shape — one row per warp, warp-stride over that row’s non-zeros, no cross-lane operations — so all five were packed.

Measured on sm_86 (m = 65536, float32) for the streaming-write shape these use:

avg_nnz

<<<m, 32>>> (ms)

packed 8/block (ms)

speedup

8

0.077

0.031

2.50x

16

0.078

0.030

2.57x

32

0.081

0.045

1.79x

64

0.110

0.074

1.48x

128

0.146

0.146

1.00x

256

0.309

0.285

1.08x

512

0.565

0.711

0.80x

This one is not a uniform win. Packing is 2.5x faster over avg_nnz 8–32, ~1.5x at 64, parity at 128–256, and roughly 20% slower at 512 — with enough work per row the one-warp-per-block form’s larger grid schedules better than the capped grid-strided one. Both kernels’ warp tiers serve avg_nnz 8–511, so a band near the top of that tier regresses slightly. It was taken anyway: real sparse connectivity sits well below 256 nonzeros per row, so the 2.5x low end dominates in practice.

If that band ever matters, the fix is to select warps-per-block from avg_nnz at launch rather than fixing it at 8 — deliberately not done here, since it adds a tuning knob for a case that has no measured workload behind it.

The slice kernels needed one extra care beyond the mechanical rewrite: they carry a second guard, if (r < 0 || r >= m) return, for out-of-range row indices. Inside a grid-strided loop that must become continue, or one bad row index would silently drop every remaining row assigned to that warp.

BE_CSRMV_WARP_GRID was renamed BE_WARP_PER_ROW_GRID (and the helper be_csrmv_warp_grid → be_warp_per_row_grid) now that slicing and dt2t share it; the CSRMV-specific name would have been misleading.

Audits that found nothing (recorded so they are not redone)#

  • Dynamic shared-memory sizing, tree-wide. 12 kernels use extern __shared__. Every request now matches the largest index the kernel can write. _fcn/binary_fcnmm.cu was the one to check closely — it uses the same (warpid + stride) * 32 + lane strip indexing that broke CSRMM — and it sizes correctly as red_off + nwarps * 32 * ACC_SIZE.

  • binary_fcnmm.cu tree reduction over a non-power-of-two warp count. The reduction for (stride = nwarps >> 1; stride > 0; stride >>= 1) silently drops the top warps if nwarps is not a power of two. Not a live bug: all four launch sites hardcode bsz = 256, so nwarps == 8. Noted because a future block-size change would break it without any test failing.

  • Dispatch-tier gaps beyond Finding 4. dt2t.cu, plasticity_binary_on_{pre,post}.cu and slice_csr_slice_rows.cu all already implement three-tier dispatch. The CSRMM avg_nnz <= 512 two-tier ladders are not the Finding 4 defect: for SpMM the two tiers are warp and block (lanes map to columns), which is the right pair.

  • Warp-reduction convergence and __syncthreads() divergence: all flagged kernels verified block- or warp-uniform.

  • Unguarded fp16/bf16 atomics in float_csrmm.cu: compile fine on sm_75 under CUDA 13; sm_60 is not supported by the toolkit. Not a live bug.

Edge cases considered#

  • Orphaning a device kernel. Removing an FFI wrapper can strand the __global__ kernel it was the sole launcher of. nt_warp_perm / nt_block_perm device kernels stay: the nt_auto_perm dispatcher still launches them internally; only their standalone FFI wrappers are unreachable.

  • Shared-memory fix and occupancy. Raising the request from 32 to 1024 bytes (2048 for f64) does not change occupancy on sm_86: at 256 threads/block the limit is 1536/256 = 6 blocks/SM in both cases, confirmed with cudaOccupancyMaxActiveBlocksPerMultiprocessor. So the fix costs nothing.

  • Cache invalidation. Editing any .cu changes its content hash, so the first GPU run after this change recompiles. Expected.

  • CPU suite proves nothing. .cu sources are never compiled unless a GPU kernel is actually lowered. Every change here must be GPU-validated, per the predecessor spec’s takeaway.

  • Warp reduction convergence. The new warp kernels derive row from blockIdx.x * warps_per + warp_id only. Had the grid-stride bound been lane-dependent, lanes would retire at different trip counts and __shfl_down_sync would read from exited lanes. Audited all 6 pre-existing warp-reduction kernels the same way; all are warp-uniform.

  • Empty rows in the warp kernel. The thread and block kernels early-return on start == end; the warp kernel deliberately does not, because an early return inside a grid-strided loop would skip that warp’s remaining rows. Instead the inner loop simply does not execute, the reduction yields ACC_ZERO, and lane 0 writes it. Covered by the ragged-row test.

  • Zero-sized grid. be_csrmv_warp_grid() clamps to at least 1. csrmv_t_warp has no avg_nnz guard, so an m == 0 matrix would otherwise reach the launch with a zero-sized grid — an invalid configuration. (The pre-change <<<m, 32>>> had the same hole; this closes it.)

  • Threshold asymmetry with float_csrmv.cu. The binary sources switch to the warp tier at avg_nnz 16 while float_csrmv.cu uses 8. The measured crossover for the binary kernels is between 12 and 16 (0.030 vs 0.034 at 12; 0.040 vs 0.034 at 16). The float threshold was left at 8: converting its warp kernel to the packed shape only makes the warp tier faster, so an unchanged threshold stays on the safe side of the crossover.