prefetch-data-load

name	prefetch-data-load
description	Apply prefetch optimization to FlyDSL kernel loops: pre-load the first iteration's data before the loop, issue async loads for the next iteration inside the loop body, and swap buffers at the loop tail via runtime loop-carried values. This overlaps data load latency with compute instructions. Use when a kernel has a loop where buffer_load feeds into MFMA/compute and load latency is exposed. Usage: /prefetch-data-load
allowed-tools	Read Edit Bash Grep Glob Agent

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:

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

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

_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

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:

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

# 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: clear SmemPtr caches, compute last partition, write output
smem_ptr._view_cache = None
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. Clear SmemPtr._view_cache before epilogue. SmemPtr.get() caches the view it creates. If called inside the runtime loop body, the cached view is defined in the loop scope. Using it in the epilogue (outside the loop) causes an SSA dominance error. Fix:

   my_smem_ptr._view_cache = None
   

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.

# 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):

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;

WARNING: Do NOT use maxnreg to force accum_vgpr=0 in hopes of freeing register space for prefetch. This forces MFMA results through arch_vgpr via v_accvgpr_read spills, causing massive slowdown (measured 4.5x GPU kernel regression).

What Prefetch Can and Cannot Do