prefetch-data-load

# Prefetch Data Load Optimization Apply software prefetch (double-buffering) to overlap async data loads with compute in FlyDSL GPU kernel loops. **API note.** The worked examples below are transcribed from the PA decode kernel and are *schematic*: they demonstrate the loop structure — prologue, loop-carried state, epilogue — not a copy-pasteable kernel. The loads are still spelled as raw `buffer_ops.buffer_load`, which now lives in `kernels/common/buffer_ops.py` (moved out of `flydsl.expr` in #880); for new code use `fx.rocdl.make_buffer_tensor` + `fx.copy` instead. The MMA is shown in the current atom form (`mma = fx.make_mma_atom(...)`, then `fx.gemm(mma, d, a, b, c)`); the raw `rocdl.mfma_*` intrinsics take `(result_type, operands)`, not `(a, b, acc)`. See the **kernel-code-cleanup** skill for the full mapping. ## Core Principle GPU global memory loads (`buffer_ops.buffer_load`, `buffer_load_dwordx4`) are **asynchronous** -- the load instruction returns immediately and the hardware fetches data in the background. The data is only needed when a subsequent instruction actually **consumes** it. If we issue the load early enough, the data arrives by the time we need it, effectively hiding the load latency behind compute work. **Without prefetch** (load latency fully exposed): ``` for i in range(N): data = load(ptr + i) # <-- stall: wait for data result = compute(data) # <-- cannot start until load completes ``` Timeline: ``` |--load--|--stall--|--compute--|--load--|--stall--|--compute--| ``` **With prefetch** (load overlapped with compute): ``` # Pre-load first iteration BEFORE the loop next_data = load(ptr + 0) for i in range(N): # Swap: the prefetched data becomes current data = next_data # Issue load for NEXT iteration (async, non-blocking) if i + 1 < N: next_data = load(ptr + i + 1) # Compute using CURRENT data -- overlaps with next load result = compute(data) ``` Timeline: ``` |--load₀--|--compute₀ + load₁--|--compute₁ + load₂--|--compute₂--| ``` The total time drops from `N * (load + compute)` to roughly `load + N * max(load, compute)`. ## FlyDSL Implementation: `range(..., init=...)` with Loop-Carried Prefetch In FlyDSL kernels, Python-level `for _pi in range(N)` gets traced into N flat copies that LLVM re-rolls. This makes the `data = next_data` swap **invisible** to MLIR — both variables alias the same SSA value, so LLVM hoists loads as loop-invariant. **Solution**: Use FlyDSL's runtime `range(..., init=...)` (loop-carried values) to create genuine SSA phi nodes. See the `flydsl-kernel-authoring` skill, section "Runtime Loops with Loop-Carried Values", for the full pattern and three critical pitfalls. ### Transformation Steps Given a loop like: ```python for i in range(START, END): # === LOAD PHASE === offsets = compute_offsets(i) data_A = buffer_ops.buffer_load(rsrc_A, offsets, vec_width=4) data_B = buffer_ops.buffer_load(rsrc_B, offsets, vec_width=4) # === COMPUTE PHASE === fx.gemm(mma, acc, transform(data_A), transform(data_B), acc) # d, a, b, c ``` Apply the following transformation using `range(..., init=...)`: #### Step 1: Prologue — load first iteration before loop ```python offsets_0 = compute_offsets(START) next_data_A = buffer_ops.buffer_load(rsrc_A, offsets_0, vec_width=4) next_data_B = buffer_ops.buffer_load(rsrc_B, offsets_0, vec_width=4) init_state = [_unwrap(v) for v in [next_data_A, next_data_B, acc]] ``` #### Step 2: Runtime loop with loop-carried state ```python _start = fx.Int64(0) _stop = fx.Int64(N - 1) # N-1 iterations; last handled in epilogue _step = fx.Int64(1) for iv, state in range(_start, _stop, _step, init=init_state): # Swap: prefetched -> current data_A = state[0] data_B = state[1] acc = state[2] # Prefetch next iteration (async, non-blocking) offsets_next = compute_offsets(iv + 1) next_data_A = buffer_ops.buffer_load(rsrc_A, offsets_next, vec_width=4) next_data_B = buffer_ops.buffer_load(rsrc_B, offsets_next, vec_width=4) # Compute using current data (overlaps with next load) fx.gemm(mma, acc, transform(data_A), transform(data_B), acc) results = yield [_unwrap(v) for v in [next_data_A, next_data_B, acc]] ``` #### Step 3: Epilogue — process last iteration ```python data_A = results[0] data_B = results[1] acc = results[2] fx.gemm(mma, acc, transform(data_A), transform(data_B), acc) ``` ### Handling auxiliary data (block tables, scales) Any offset calculations, block table lookups, or scale factor loads needed for the *next* iteration's data should also be carried as loop state: ```python init_state = [_unwrap(v) for v in [ next_data_A, next_data_B, next_block_id, next_scale, acc ]] for iv, state in range(_start, _stop, _step, init=init_state): data_A, data_B, block_id, scale, acc = state # Prefetch next iteration next_block_id = load_block_table(iv + 1) offsets_next = compute_offsets(iv + 1, next_block_id) next_data_A = buffer_ops.buffer_load(rsrc_A, offsets_next, vec_width=4) next_data_B = buffer_ops.buffer_load(rsrc_B, offsets_next, vec_width=4) next_scale = buffer_ops.buffer_load(rsrc_scale, next_block_id, vec_width=1) # Compute with current data fx.gemm(mma, acc, transform(data_A) * scale, transform(data_B), acc) results = yield [_unwrap(v) for v in [ next_data_A, next_data_B, next_block_id, next_scale, acc ]] ``` ### PA Decode Kernel Example (verified, 112us, 0.75x vs Gluon) State inventory (15 values carried across iterations): - 8 x `vector<4xi32>` — K data (4 tiles x 2 loads) - 1 x `i32` — partition_start - 2 x `i32` — block table values (phys_block/page_off or phys_0/phys_1) - 2 x `f32` — running_max, running_sum (online softmax) - 2 x `vector<4xf32>` — PV accumulators ```python # Pack/unpack helpers def _pack(kv_flat, part_start, bt_vals, rmax, rsum, acc_pv): raw = kv_flat + [part_start] + bt_vals + [rmax, rsum] + acc_pv return [v.ir_value() if hasattr(v, 'ir_value') else v for v in raw] def _unpack(state): kv_flat = list(state[0:8]) kv = [[kv_flat[t*2], kv_flat[t*2+1]] for t in range(4)] return kv, state[8], list(state[9:11]), state[11], state[12], [state[13], state[14]] # Prologue pf_0 = issue_bt_k_loads(partition_0) init_state = _pack(flatten(pf_0['kv']), pf_0['part_start'], ...) # Runtime loop (bounds MUST be a typed DSL integer like fx.Int64, not Python ints!) for iv, state in range(fx.Int64(0), fx.Int64(N - 1), fx.Int64(1), init=init_state): kv, part_start, bt, rmax, rsum, acc = _unpack(state) rmax, rsum, acc = compute_qk_softmax_pv(kv, part_start, bt, rmax, rsum, acc) pf_next = issue_bt_k_loads(next_partition(iv + 1)) results = yield _pack(flatten(pf_next['kv']), pf_next['part_start'], ...) # Epilogue: compute last partition, write output kv, part_start, bt, rmax, rsum, acc = _unpack(results) compute_qk_softmax_pv(kv, part_start, bt, rmax, rsum, acc) write_output(rmax, rsum, acc) ``` **ISA result**: 8 K-prefetch `buffer_load_dwordx4` appear at the END of the loop body (after PV MFMA), overlapping with the MFMA pipeline drain. The prologue has 8 K loads before the loop. The epilogue has 8 V loads only (no K loads needed). ### Three Critical Pitfalls 1. **Loop bounds must be a typed DSL integer such as `fx.Int64(...)`, NOT a Python int.** A plain int makes the AST rewriter unroll the loop and silently ignore `init=`. If you write `range(0, 15, 1, init=...)`, the AST rewriter treats the constant bounds as a Python `range` and unrolls; only plain Python-int bounds are unrolled, so a typed bound still produces a runtime `scf.for` (the rewriter index-casts non-Python-int bounds into `scf.for`). Use `fx.Int64(0)`, `fx.Int64(15)`, `fx.Int64(1)` instead. 2. **Prefer internal types; unwrap only at hard boundaries.** Most loop-carried values can remain `fx.Int32`, `fx.Float32`, or `fx.Vector`. Prefer these concrete types over wrapping a raw value in `ArithValue` directly -- note `fx.Vector` subclasses `ArithValue`, so this is about which constructor you reach for, not about avoiding the base class. If a low-level helper explicitly expects raw `ir.Value`, unwrap at that boundary. 3. **Build LDS views at the top of the kernel, not inside the runtime loop.** Allocate shared memory with `fx.SharedAllocator().allocate(...).peek()` and build each `.view()` once up front. A view created inside the `scf.for` body is defined in the loop scope; using it in the epilogue (outside the loop) causes an SSA dominance error. Building it once at the top makes it dominate both the loop and the epilogue. ## Applicable Patterns This optimization applies whenever you see this pattern in a kernel: | Signal | Description | |--------|-------------| | `for ... in range(N)` loop with `buffer_load` followed by MFMA | Load-then-compute in a loop body | | Block table lookup inside loop | `buffer_load(block_table_rsrc, idx)` followed by `buffer_load(cache_rsrc, page_id * stride)` | | KV cache iteration | Paged attention, flash attention, any tiled GEMM with paged memory | | Scale factor loads | FP8 per-token quantization scales loaded per KV block | ## Compiler Constraints FlyDSL kernels compile to GCN ISA where `s_waitcnt` insertion is controlled by the **compiler**, not by the programmer. You cannot directly eliminate `s_waitcnt` instructions. Instead, prefetch restructures the code so the compiler places `s_waitcnt` after enough compute work to hide the latency. ### Register Budget **Always check register headroom before adding prefetch buffers:** On CDNA3 (gfx942 MI300X/MI308), VGPRs are tracked as two **physical** files that share **one combined 512-entry occupancy budget** per SIMD: - **arch_vgpr** (up to 256 per SIMD): used by VALU, VMEM loads, LDS ops, and prefetch buffers - **accum_vgpr / AGPR** (up to 256 per SIMD): used by MFMA result writeback Prefetch buffers physically live in **arch_vgpr** and MFMA accumulators in **accum_vgpr**, but occupancy is governed by their **sum** (`arch_vgpr + accum_vgpr`), so growing prefetch buffers *does* compete with MFMA accumulators for the shared 512 budget and can cost occupancy. ```python # Estimate arch_vgpr cost of prefetch buffers: # - Each buffer_load_dwordx4 = 4 arch_vgpr per load # - 8 K-cache loads = 8 x 4 = 32 arch_vgpr for one buffer set # - Double-buffering = 2 x 32 = 64 arch_vgpr (but one set is reused) # - Net additional arch_vgpr ~ 32 (the "next" buffer) # # On MI300X (gfx942): arch_vgpr + accum_vgpr share ONE combined 512 budget/SIMD # Occupancy = 512 / (arch_vgpr_alloc + accum_vgpr_alloc) waves per SIMD # (combined-pool model — NOT 256/max; that was gfx908/CDNA1 only) # # Example: arch=148, accum=148 -> combined 296 -> 512//296 = 1 wave # Adding 32 arch_vgpr -> combined 328 -> still 1 wave (safe) # To reach 2 waves you need combined (arch+accum) <= 256 # arch+accum > 512 -> SPILL (exceeds the combined per-SIMD budget) ``` **Critical thresholds (gfx942, combined arch+accum budget):** | Combined arch_vgpr + accum_vgpr | Max Waves/SIMD | Impact | |--------------|---------------|--------| | <= 128 | 4 | High occupancy | | <= 170 | 3 | Good occupancy | | <= 256 | 2 | Moderate occupancy | | <= 512 | 1 | Minimum occupancy | | > 512 | **SPILL** | Register overflow -> severe perf regression | **How to check current VGPR allocation** (from rocprofv3 database): ```sql SELECT ks.KernelName, ki.arch_vgpr_count, ki.accum_vgpr_count FROM rocpd_kernel_dispatch kd JOIN rocpd_info_kernel_symbol ks ON kd.kernel_symbol_id = ks.id JOIN rocpd_info_kernel ki ON kd.kernel_id = ki.id WHERE ks.KernelName LIKE '%target_kernel%' LIMIT 5; ``` **NOTE**: `maxnreg` has been removed (it never reached LLVM). Forcing `accum_vgpr=0` to free register space for prefetch was never viable anyway: it pushes MFMA results through arch_vgpr via `v_accvgpr_read` spills, measured at a 4.5x GPU kernel regression. Target occupancy with `waves_per_eu`. ### What Prefetch Can and Cannot Do