[Perf] Keep runtime base and static tail separate in add_offset [wip]

[Phil-amd]

Summary

The nested-add_offset rewrite in MemrefLowering.td fused unconditionally, collapsing add_offset(add_offset(smem, sw), const) into add_offset(smem, sw + const). With smem a relocatable LDS symbol, every read then computed its own address off the runtime base and needed its own base VGPR. This replaces the rule with a canonicalization that keeps the runtime base and the static tail separate.

Motivation

Reported in #898. On gfx950 the SWA attention kernel showed 38 distinct ds_read_b64_tr_b16 base registers, 44 VGPR spills and 44/27 scratch traffic. The recast_iter form — which the old pattern did not match, so it acted as an accidental fusion barrier — showed 3 bases and no spills, and ran ~26% faster. Kernels have been working around this by inserting recast_iter; this removes the need.

One correction to the issue's diagnosis, from the measured ISA: constants do fold into the ds_read offset: immediate in both forms (252/256 either way). The defect is base-register scatter, not failed immediate folding. The relevant metric is the distinct base-register count.

Changes

• Remove the unconditional fusion rule from MemrefLowering.td.
• Add AddOffsetCanonicalization in LayoutLowering.cpp, converging an add_offset chain to add_offset(add_offset(ptr, dyn), static) — at most one dynamic layer plus one static tail.
  • static+static and dynamic+dynamic still fuse; neither merges a runtime value with a constant.
  • Only the swap is gated on a single-use inner op, because it rebuilds that op. Gating fusion the same way was tried and regressed the buffer path badly: chains were left unfused, costing +18 VGPR on fused_add_rmsnorm_kernel_0 (14 add_offset ops became 115, 6 LLVM adds became 102).
• Termination: each rewrite strictly reduces the static layer's depth or the chain length, and the resulting dyn -> static form is rejected, so the pattern cannot ping-pong. Three-level chains are covered by tests.

Performance

gfx950, SWA attention repro from the issue, use_recast=False (the plain form):

N	before	after	speedup
4096	256 VGPR / 44 spill / 38 bases, 0.0906 ms	242 / 0 / 3, 0.0719 ms	+26.0%
8192	256 / 44 / 38, 0.1748 ms	242 / 0 / 3, 0.1383 ms	+26.4%
16384	256 / 44 / 38, 0.3443 ms	242 / 0 / 3, 0.2740 ms	+25.7%

The plain form now matches the recast_iter form exactly, with the two within measurement noise of each other.

The benefit is bimodal rather than proportional: it comes from crossing the spill threshold, not from the handful of saved VALU adds. Kernels that were not close to spilling see no change — a per-kernel resource comparison against this base shows 32 kernels all unchanged on gfx942, no regressions and no improvements.

Applicability boundary: a non-trivial swizzle is applied to the final address at the load, so the static tail is consumed by the XOR and cannot become an immediate. The canonicalization is neutral there, not harmful. This is asserted by a test so the boundary stays visible.

Testing

• Unit tests added/updated — 9 transform cases in layout_lowering.mlir (four static/dynamic combinations, three-level chains, multi-use inner, shared runtime base, zero/negative constants) and 3 conversion cases in nested_add_offset.mlir (LDS, buffer, swizzle). Both files were verified to fail with the source change reverted.
• Performance benchmarks run — table above; per-kernel VGPR/spill/scratch delta against this base, 32 kernels, 0 worsened.
• Tested on MI300X — full MLIR FileCheck suite 37/37 on gfx942; clang-format clean.

Breaking Changes

None. Kernels using recast_iter as a workaround keep working unchanged — recast_iter lowers to a no-op — they simply no longer need it for this.