The nested-add_offset rewrite in MemrefLowering.td:15-16 collapses
add_offset(add_offset(smem, sw), const) into add_offset(smem, sw + const).
smem is a relocatable LDS symbol, so fusing the runtime sw with the
compile-time const before the GEP produces a distinct dynamic index per
read:
%idx = add i32 %sw, 1024
%p = getelementptr bfloat, ptr addrspace(3) @__shared_alloc_2, i32 %idx
The backend can no longer pull const back into the ds_read offset: immediate,
so it materializes one base VGPR per read → register scatter → spilling.
recast_iter is not matched by the nested-add_offset pattern, so it acts as a
barrier: gep(gep(@sym, sw), const) survives and each constant folds into the
immediate off a single base VGPR.
Files
kernel_swa_repro.py — the full GQA SWA kernel. Identical to
../kernel_swa_v3.py except build_gqa_attn takes a use_recast flag that
selects the load_v_regs LDS-read base form:
base_ptr = fx.add_offset(smem_ptr, fx.make_int_tuple(sw_elems)) # runtime base
if use_recast:
base_ptr = fx.recast_iter(bf16, base_ptr) # fusion barrier
...
view = fx.make_view(fx.add_offset(base_ptr, off // 2), _v_tr_layout) # const/readcompare_full_kernel.py — builds the kernel twice (use_recast=False and
True), dumps the final gfx950 ISA for each, and reports VGPRs / spills /
scratch traffic / distinct ds_read base registers / runtime / TFLOPS.
The two compiled kernels differ only by that one recast_iter, so every
metric delta is attributable to the add_offset fusion.
compare_full_kernel.py
kernel_swa_repro.py
The nested-
add_offsetrewrite inMemrefLowering.td:15-16collapsesadd_offset(add_offset(smem, sw), const)intoadd_offset(smem, sw + const).smemis a relocatable LDS symbol, so fusing the runtimeswwith thecompile-time
constbefore the GEP produces a distinct dynamic index perread:
The backend can no longer pull
constback into theds_read offset:immediate,so it materializes one base VGPR per read → register scatter → spilling.
recast_iteris not matched by the nested-add_offsetpattern, so it acts as abarrier:
gep(gep(@sym, sw), const)survives and each constant folds into theimmediate off a single base VGPR.
Files
kernel_swa_repro.py— the full GQA SWA kernel. Identical to../kernel_swa_v3.pyexceptbuild_gqa_attntakes ause_recastflag thatselects the
load_v_regsLDS-read base form:compare_full_kernel.py— builds the kernel twice (use_recast=FalseandTrue), dumps the final gfx950 ISA for each, and reports VGPRs / spills /scratch traffic / distinct
ds_readbase registers / runtime / TFLOPS.The two compiled kernels differ only by that one
recast_iter, so everymetric delta is attributable to the
add_offsetfusion.compare_full_kernel.py
kernel_swa_repro.py