[EXTERNAL] Do not store undefined lanes into a shared spill slot

[Muhamed-Husic]

Upstream issue: llvm/llvm-project#225054

Motivation

Fixes the GPU memory access faults (exit -13) that the weekly parameterSweeps on #361 hit on gfx950 attention kernels. Targets #361, since that branch enables the weekly CI stages that found this fault.

Root cause

A bug in LLVM's InlineSpiller (register allocator). Under heavy register pressure, SGPRs are spilled into VGPR lanes. Inside the KV loop, a buffer descriptor is rebuilt from the upper half (sub2_sub3) of another descriptor, so only that half is written. The spiller then stored all four registers of this half-written descriptor, including the unwritten lower half, into a spill slot that is reloaded at the top of every iteration. The unwritten half overwrote the reloaded descriptor's base address with garbage, so loads from the second iteration on hit an unmapped page.

Fix

Two changes in external/llvm-project/llvm/lib/CodeGen/InlineSpiller.cpp. Both stop the spiller from writing garbage lanes into a spill slot that still holds a live value.

• hoistSpillInsideBB: this function moves a spill earlier, to where the register being stored is defined. It now refuses when that register leaves some lanes undefined that the original value still needs. Before, the full-width store wrote whatever those lanes happened to hold into the shared slot and overwrote the value another piece of the register would reload later.
• spillAroundUses: when an instruction writes only some lanes of a spilled register (for example, only sub2_sub3) and the slot still holds the value, the spiller now reloads the slot before that instruction. The full-width store after it then writes the untouched lanes back unchanged, instead of garbage.

Lanes the original value never uses are still allowed to hold garbage, so LLVM's existing hoisting of partially defined values (splitkit.mir, splitkit-copy-bundle.mir) is unaffected.

Test

The new mlir/test/e2e/PrAttentionBF16Gfx950.toml pins the bf16 attention config that faulted in the weekly replay, for both transO values. It is gfx950-only because the perf config needs 128 KiB of LDS. It hands the backend the same kernel LLVM IR as the faulting sweep command. It uses the fixed input pattern instead of %random_data (with random data this shape misses the bf16 tolerance for any perf config, which is unrelated to this bug).

Verified locally on gfx950 with two drivers that differ only in InlineSpiller.cpp:

• Without the fix: memory access fault, 3 of 3 runs of the sweep command, and both E2E cases.
• With the fix: [1 1 1], 3 of 3, and both E2E cases pass through lit.
• The generated code differs only at the spill site.

Validation

• Weekly parameterSweeps pass on this branch:
  • Before the fix: https://ml-ci-internal.amd.com/blue/organizations/jenkins/MLIR%2FrocmlirTriton/detail/PR-361/60/pipeline/442/. This is [AIROCMLIR-899] Enable weekly CI stages (parameter sweeps + tuning) #361 at 42afd25 (the base of this PR).
  • After the fix: https://ml-ci-internal.amd.com/blue/organizations/jenkins/MLIR%2FrocmlirTriton/detail/PR-552/2/pipeline/.
• PR CI passes, including the new PrAttentionBF16Gfx950 E2E test on gfx950.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Additional details and impacted files

@@                                Coverage Diff                                @@
##           users/bpetkovi/AIROCMLIR-899-enable-weekly-ci     #552      +/-   ##
=================================================================================
- Coverage                                          87.39%   84.65%   -2.74%     
=================================================================================
  Files                                                208      208              
  Lines                                              40860    36263    -4597     
  Branches                                            9121     8233     -888     
=================================================================================
- Hits                                               35709    30697    -5012     
- Misses                                              3052     3347     +295     
- Partials                                            2099     2219     +120     

Flag	Coverage Δ
gfx950	84.43% <ø> (+0.01%)	⬆️
mfma	84.46% <ø> (+1.04%)	⬆️
navi4x	84.54% <ø> (+0.07%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 42 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[umangyadav]

Looks like upstream has a PR to fix
llvm/llvm-project#226689

Does that work to solve the weekly issue ?

[Muhamed-Husic]

Looks like upstream has a PR to fix llvm/llvm-project#226689

Does that work to solve the weekly issue ?

I tested it and llvm/llvm-project#226689 doesn't fix our fault. I applied its current commit (llvm/llvm-project@c3e394a8) to our LLVM in place of my fix, built rocmlir-driver with it, and ran the bf16 config that faults in the weekly. It hit the same memory access fault in 3 of 3 runs. My fix passed 3 of 3 on the same GPU right after.

The reason is that our spill goes through the in-block hoist, and neither of the PR's two changes catches it there. The PR's guard on that hoist only checks the parts of a register that LLVM tracks separately. LLVM doesn't track the unwritten half of our half-written descriptor separately, so the guard has nothing to reject and lets the hoist through. The PR's lane mask, which makes a spill store only the registers that were written, is only set on the normal spill path, not on the hoist. The PR did narrow other spills in the same kernel, so it was active, just not on the one that faults. So we still need our downstream patch.

[Muhamed-Husic]

Closing in favour of #559