[Dialect][Perf] Don't merge mixed static/runtime offsets on LDS pointers

[Arist12]

Summary

The nested-add_offset merge in MemrefLowering.td folds a compile-time offset into a runtime one before the getelementptr is formed. On a shared (LDS) pointer that hides the constant inside a dynamic index, so the AMDGPU backend can no longer sink it into the ds_read ... offset: immediate and every access has to materialize its own base address. In a kernel with many live LDS reads that register scatter becomes spilling.

This adds a single guard to that pattern: on a shared pointer, merge only when both offsets are of the same kind — constant with constant (pure constant folding) or runtime with runtime (nothing constant is lost). A mixed chain is left nested, in either operand order. Pointers in other address spaces are unaffected.

Fixes #898.

Motivation

Reported in #898 with a GQA sliding-window attention kernel. Written the natural way it spills 44 VGPRs and issues 71 scratch operations in its inner loop; inserting an fx.recast_iter between the two add_offset calls — a no-op for the address arithmetic that happens to block the merge — removes all of it and runs ~26% faster. Kernel authors should not need to know about that workaround.

Changes

• include/flydsl/Dialect/Fly/IR/FlyOps.td: new Fly_MergeableOffsetChain constraint — the pointer is not shared/LDS, or both offsets have the same staticness.
• include/flydsl/Dialect/Fly/Transforms/MemrefLowering.td: the existing nested-offset pattern gains that constraint. The rewrite itself is unchanged and no new pattern is added, so the diff is one guard plus a comment.
• tests/mlir/Transforms/layout_lowering.mlir: merge-vs-leave-nested at the fly level, for both operand orders and for global pointers.
• tests/mlir/Conversion/add_offset_reassociation.mlir (new): the resulting address formation after --convert-fly-to-rocdl, i.e. the shape the backend actually sees.
• tests/kernels/test_lds_offset_chain.py (new): device test running both mixed orderings plus the already-merged control, checking the values read against a torch reference.

Performance (MI355X, gfx950, ROCm 7.2)

Reproducer from the issue, B=1, window (512,0), D=128, H=32, H_KV=16, bf16. "plain" is the natural nested add_offset form; "recast" is the recast_iter workaround.

N	variant	VGPR	spill	scratch	ds_read bases	ms	TFLOPS
8192	plain, before	256	44	71	38	0.1752	381
8192	plain, after	242	0	0	3	0.1376	485
8192	recast, before/after	242	0	0	3	0.1385/0.1374	482/485

Across shapes, plain form before → after: 4096 0.0916 → 0.0741 ms (352 → 436 TFLOPS), 8192 0.1752 → 0.1376 ms (381 → 485 TFLOPS), 16384 0.3500 → 0.2780 ms (387 → 488 TFLOPS). The plain form now matches the workaround on every counter.

Testing

• Unit tests added/updated — 4 cases in layout_lowering.mlir, 5 in the new conversion test, 3 device cases in test_lds_offset_chain.py. MLIR FileCheck suite 36 → 37 files, all passing. The new conversion test was confirmed to fail against an unpatched fly-opt, and the device test to fail if an address component is dropped, so neither passes vacuously.
• Performance benchmarks run — table above.
• Tested on MI355X (gfx950). Cross-compiled the new device test for gfx942, gfx1201 and gfx1100 as well; all lower and compile.
• RUN_TESTS_FULL=1 scripts/run_tests.sh on an unpatched tree and on this branch, same machine: 2248 passed / 0 failed and 2251 passed / 0 failed respectively, the difference being exactly the three new device cases. (Multi-GPU tests excluded; they need more than one visible device.)
• Codegen impact on existing kernels: compiled 22 in-tree kernels across preshuffle GEMM fp8, flash attention forward, blockscale preshuffle GEMM and MoE GEMM in both trees and diffed the emitted ISA — byte-identical in every case.

Dependencies

• No new third-party dependencies added.

Breaking Changes

None. The guard only declines a merge; it introduces no new rewrite, and non-shared address spaces are untouched.

Made with Cursor

[Arist12]

Additional evidence, in case it is useful for review: since the headline numbers in the description are all register/spill counters, here is the rest of the emitted ISA for the same kernel (N=8192, plain add_offset form), to show the win is not just displacing cost somewhere I wasn't counting.

	before	after
total instructions	4087	3837	−250
VALU	3085	2971	−114
SALU	491	426	−65
s_waitcnt	49	24	−25
v_mfma	256	256	unchanged
ds_read	384	384	unchanged
LDS bytes (group_segment_fixed_size)	65536	65536	unchanged

The MFMA and ds_read counts being identical is the relevant part: the arithmetic and the LDS traffic are the same kernel, so nothing was skipped or hoisted out, and LDS usage is flat, so the spills were eliminated rather than moved into shared memory. The instruction and wait-count reductions are the address-materialisation code that is no longer needed once the reads share a base register.