Chunked-dot Reduce lowering: shorten the reduction dependency chain (measured up to 1.09x, low-occupancy only)

[!WARNING]
Superseded in part by measurement — see this comment. Run on an A100, the register argument below is wrong: ptxas allocates 34 registers for the wide form and 44 for the chunked one, with zero spill in both. The "257 registers" figure is PTX virtual registers (SSA numbering), not pressure. The surviving argument is the reduction's dependency-chain depth, which pays only at low occupancy — measured up to 1.09×, not the ~1.2× estimated here, and 0.98–1.00× once bandwidth-saturated. Chunking also raises register use, so it risks regressing a kernel already at 255-with-spill. Body kept unedited below for the record.

---

Splitting this out of hiraditya/Vx.1#378 so it does not get buried under the attention-routing work. It is the last unclaimed FP32 lever measured in that campaign, it is a self-contained codegen change, and it is worth roughly 1.2× on every slice reduction — not just attention.

The claim

dot(a, b) on the flat path materializes both operand rows as whole vectors before reducing. At HD=64 f32 that is ~192 registers of live state for one scalar result, which is what pins the split-K prefill kernel to the 255-register architectural ceiling with spill. Accumulating the dot in chunks instead should hold ~24 registers, clear the spill, and plausibly free enough budget for a third softmax chain.

What the lowering does today

src/codegen/flat.rs:2235, the Opcode::Reduce arm. For dot(q[i], k[j]) over a 64-wide f32 row it emits:

%vl = vector.load %q[%c0] : memref<64xf32>, vector<64xf32>   // 64 registers
%vr = vector.load %k[%c0] : memref<64xf32>, vector<64xf32>   // 64 registers
%vp = arith.mulf %vl, %vr : vector<64xf32>                   // 64 registers
%v  = vector.reduction <add>, %vp : vector<64xf32> into f32

Three 64-lane f32 vectors live at once. NVPTX renders each vector<64xf32> load as 16 × ld.global.v4.f32, so the dot alone carries ~32 v4 temporaries. The record kernel runs two of these chains, each with its own [1, HD] accumulator (64 floats), so the working set is comfortably past what an SM can give one thread.

The evidence that this is the binding constraint

From the hiraditya/Vx.1#378 campaign, A100, SQ=8192 SK=2048 HD=64, all FP32, device-event timed.

The record kernel (split-K C=8 + two softmax chains, 1.125 ms total = 3.82 TF/s = 19.6% of FP32 FMA peak) compiles to 255 registers/thread — the architectural maximum — with 128 B of spill traffic (ptxas -v). Register cap sweep via VX_MAXRREG (the JIT knob added in d8945492):

reg cap	blocks/SM	region1
none (255)	2	1.083 ms
200	2	2.283
168	3	2.554
128	4	6.006
96	5	8.822

Every cap loses badly. The working set is real — this is not an occupancy tuning problem, it is a "the values genuinely do not fit" problem, and spills swamp any occupancy gained. Corroborating: four softmax chains ran ~3× slower than two, which is the same wall from the other side (4 × 64 accumulator floats spill outright).

So the way forward is not to cap registers or reduce ILP. It is to stop needing the registers.

Proposed change

Lower a slice reduction as a chunked accumulation loop rather than a whole-row materialization:

%zero = arith.constant dense<0.0> : vector<8xf32>
%acc = scf.for %i = %c0 to %c64 step %c8 iter_args(%a = %zero) -> (vector<8xf32>) {
  %x = vector.load %q[%i] : memref<64xf32>, vector<8xf32>
  %y = vector.load %k[%i] : memref<64xf32>, vector<8xf32>
  %n = vector.fma %x, %y, %a : vector<8xf32>
  scf.yield %n : vector<8xf32>
}
%v = vector.reduction <add>, %acc : vector<8xf32> into f32

Live state drops from ~192 registers to 3 × 8 = 24. vector<8xf32> is still two 128-bit loads, so the v4 vectorization that Vx#378 R1 added the vector dialect for is preserved — this trades breadth for depth, it does not go scalar.

Notes on scope:

• Chunk width should be a constant that is easy to retune (8 is the obvious first try; 4 and 16 are worth sweeping).
• Guard it: rows at or below the chunk width keep the current single-shot form, which is already optimal for them.
• Applies to all four slice reductions (dot/sum/max/min), not just dot — the same arm handles them via ins.imm.
• Interacts cleanly with the f16 widening contracts (Vx#320): the arith.extf moves inside the loop and widens a chunk at a time instead of a whole row, which is also strictly less register pressure.
• This changes floating-point association order for the reduction. vector.reduction <add> over 64 lanes already has unspecified order, so chunked accumulation is a different-but-equally-valid order rather than a regression — but it should be stated in the commit, and the differential suite is the check.

Expected payoff, and how to falsify it

Estimate: ~1.2×, landing the fused FP32 record around 4.5–5 TF/s (from 3.82). Two independent sources of gain, and the estimate is wrong if neither materializes:

1. Spill elimination alone recovers the 128 B of spill traffic in the inner loop.
2. If the freed budget fits a third softmax chain, ILP3 becomes testable for the first time — it was previously dead on arrival purely for lack of registers.

Falsifiable: if ptxas -v still reports 255 registers after the change, the diagnosis is wrong and this should be closed rather than tuned.

Acceptance criteria

• ptxas -v on the split-K prefill kernel reports < 255 registers and 0 bytes spilled
• region1 improves on 1.083 ms (target ≤ 0.9 ms); total beats the 1.125 ms record
• ILP3 re-measured now that it is register-affordable (may or may not win — record either way)
• flat-vs-AST differential suite green
• closed forms still exact: prefill 10.235 (f32), smoke 0.075, decode 81.915
• chunk-width sweep {4, 8, 16} recorded, so the constant is chosen and not guessed

Pointers

• Lowering to change: src/codegen/flat.rs:2235 (Opcode::Reduce)
• Benchmark that exercises it: scripts/templates/flash_splitk_bench.vx (record kernel, f0282e57)
• Register knob for the sweep: VX_MAXRREG (d8945492)
• Parent campaign: hiraditya/Vx.1#378. Cooperative-kernel context: hiraditya/Vx.1#379. Related but distinct: hiraditya/Vx.1#337 (contracting dimensions), hiraditya/Vx.1#327 (f16 slice ops, since resolved by the Vx#320 widening contracts).

Why this is worth doing even though attention now routes to a vendor kernel

flash_attention_into (#378) hands fused f16 attention to FlashAttention-2 / cuDNN at 0.066 ms, so attention no longer depends on this. But the Reduce lowering is used by every slice reduction the compiler emits, in every program that is not a routed vendor shape — and "the compiler's own kernels do not fit in registers" is a general codegen weakness that will keep surfacing. This is the cheapest known fix for it, and it is the difference between the compiler's own FP32 output sitting at 19.6% of peak versus ~25%.

Start in src/codegen/flat.rs:2235 at the Opcode::Reduce arm, then inspect scripts/templates/flash_splitk_bench.vx and the flat-vs-AST differential suite. Test chunk widths 4, 8, and 16 with VX_MAXRREG, and use ptxas -v plus the listed benchmarks to compare registers, spills, timing, ILP3, and exact closed forms; done requires the stated acceptance criteria.