# Kernel Architect Notes — Turn 5

## Current state

Fastest solution (2.061x, ~17.9 ms) is:
1. `x.contiguous(channels_last)` (~0.8 ms, 604 MB read+write of input)
2. cuDNN NHWC TF32 conv (~9.3 ms, must write 4.77 GB conv output)
3. fused pool+sigmoid+sum kernel reading the NHWC output (~3.5 ms, 4.77 GB read)

Turn-4 profile shows the breakdown clearly: conv 9.3 ms, elementwise (layout transforms) 7.1 ms, pool 4.5 ms, sigmoid 0.4, reduce 0.2.

## Where the remaining time goes

- The conv output (128×64×382×382 fp32 = **4.77 GB**) is written by cuDNN and then read back by the epilogue — ~6.5–7 ms of pure DRAM traffic.
- The conv itself (344 GFLOP) runs at only ~37 TFLOPS (TF32 peak ~156) because C_in=8 gives a skinny GEMM (K=72).
- The input NCHW→NHWC transform costs another ~0.8 ms.

## The plan (mega-kernel)

Fuse **everything** — conv + avgpool + sigmoid + per-batch sum — into a single kernel so the 4.77 GB conv output never touches DRAM. Traffic drops to: read input once (~0.77 GB with halo overlap) + write 128 floats.

Design (specialized for C_in=8, 3×3 s1 p0, pool=4, C_out multiple of 8 ≤ 64):

- Block = 256 threads (8 warps), tile = 4×4 pooling windows = 16×16 conv outputs. smem: input tile 18×18×8 as CHW (`xs[ic*324 + y*18 + x]`), weights `[9][8][66]` padded, bias.
- Pixel ordering m = dy*4+dx (a 4×4 pool window per m16 tile), k ordered as tap*8+ic so one `mma.m16n8k8.tf32` k-step covers all 8 channels of one tap.
- Warp handles 2 horizontally adjacent windows × all 64 channels: 2×8×4 = 64 accumulator regs, 144 mma over 9 k-steps.
- Epilogue: pool is a shuffle reduction over the 8 lanes sharing `c4` (masks 4,8,16) of d0+d2 / d1+d3 — the 16 m-rows *are* the 16 window pixels. Bias added after pooling (linear), sigmoid, reduce lanes 0–3, block-reduce 16 window sums in smem, one atomicAdd per block.
- Edge tiles: zero-fill OOB input loads (only affects invalid windows, which are discarded).

Theoretical floor: ~85 M mma ≈ 1.1 ms of tensor-core time + ~0.5 ms input reads ⇒ even at 30–40% efficiency this should land ~3–5 ms total vs ~13.5 ms now → ~6–9x speedup potential.

## Risks & mitigations

- **Fragment index bugs** (turn 4's attempt had addressing bugs and was never wired up). Mitigation: keep the proven 2.06x path as fallback and add a **one-time runtime self-check**: on the first forward for a given shape, compute both mega and fallback outputs, compare with `torch.allclose` (both TF32, diffs ~1e-3 abs on ~2.9e5 sums), and permanently disable the mega path if it mismatches. One extra fallback forward, amortized over timed runs.
- Guard the mma/cvt code with `__CUDA_ARCH__ >= 800`, check device capability in Python; non-applicable shapes (C_in≠8, k≠3, stride≠1, pad≠0, pool≠4, C_out%8≠0 or >64) go to the fallback.
- Template the kernel on `NT = C_out/8` (1–8) so other channel counts still hit the fast path.
- TF32 rounding via `cvt.rna.tf32.f32` on smem loads to match cuDNN numerics.