From c3983a4697721afa25afbc112d86b6dc8793e101 Mon Sep 17 00:00:00 2001 From: Muhamed Husic Date: Tue, 29 Sep 2026 09:17:32 +0000 Subject: [PATCH 1/2] [EXTERNAL] Do not store undefined lanes into a shared spill slot --- .../llvm/lib/CodeGen/InlineSpiller.cpp | 53 +++++++++-- llvm-patches/llvm-patch-content.txt | 75 ++++++++++++++++ ...inline-spiller-partial-sibling-spill.patch | 87 +++++++++++++++++++ mlir/test/e2e/CMakeLists.txt | 1 + mlir/test/e2e/PrAttentionBF16Gfx950.cfg | 4 + mlir/test/e2e/PrAttentionBF16Gfx950.toml | 25 ++++++ 6 files changed, 240 insertions(+), 5 deletions(-) create mode 100644 llvm-patches/patch-inline-spiller-partial-sibling-spill.patch create mode 100644 mlir/test/e2e/PrAttentionBF16Gfx950.cfg create mode 100644 mlir/test/e2e/PrAttentionBF16Gfx950.toml diff --git a/external/llvm-project/llvm/lib/CodeGen/InlineSpiller.cpp b/external/llvm-project/llvm/lib/CodeGen/InlineSpiller.cpp index 5ed3b36289..b5a2b6fbda 100644 --- a/external/llvm-project/llvm/lib/CodeGen/InlineSpiller.cpp +++ b/external/llvm-project/llvm/lib/CodeGen/InlineSpiller.cpp @@ -320,6 +320,21 @@ static void getVDefInterval(const MachineInstr &MI, LiveIntervals &LIS) { LIS.getInterval(MO.getReg()); } +/// The lanes of LI's register that hold a value at Idx. Live range splitting +/// copies only the lanes that are live, so a sibling may have lanes with no +/// subrange at all. +static LaneBitmask liveLanesAt(const LiveInterval &LI, SlotIndex Idx, + const MachineRegisterInfo &MRI) { + if (!LI.hasSubRanges()) + return LI.liveAt(Idx) ? MRI.getMaxLaneMaskForVReg(LI.reg()) + : LaneBitmask::getNone(); + LaneBitmask LiveLanes = LaneBitmask::getNone(); + for (const LiveInterval::SubRange &SR : LI.subranges()) + if (SR.liveAt(Idx)) + LiveLanes |= SR.LaneMask; + return LiveLanes; +} + /// isSnippet - Identify if a live interval is a snippet that should be spilled. /// It is assumed that SnipLI is a virtual register with the same original as /// Edit->getReg(). @@ -454,6 +469,14 @@ bool InlineSpiller::hoistSpillInsideBB(LiveInterval &SpillLI, if (DefMBB != CopyMI.getParent() || !SrcQ.isKill()) return false; + // The hoisted store writes all of SrcReg into the slot every sibling shares. + // A lane SrcReg leaves undefined while the original value is live in it + // would overwrite that value, which another sibling may still reload. + LaneBitmask OrigLanes = + liveLanesAt(LIS.getInterval(Original), SrcVNI->def, MRI); + if ((OrigLanes & ~liveLanesAt(SrcLI, SrcVNI->def, MRI)).any()) + return false; + MachineBasicBlock *MBB = DefMBB; MachineBasicBlock::iterator MII; if (SrcVNI->isPHIDef()) @@ -1368,20 +1391,40 @@ void InlineSpiller::spillAroundUses(Register Reg) { // FIXME: Infer regclass from instruction alone. Register NewVReg = Edit->createFrom(Reg); - if (RI.Reads) + // The spill below stores every lane of NewVReg. When MI defines only some + // of them without reading the rest, lanes the original value holds across + // MI would be overwritten in the slot, where a sibling may still reload + // them. Fill them from the slot first so the store writes that value back. + LaneBitmask DefLanes = LaneBitmask::getNone(); + bool hasLiveDef = false; + for (const auto &[OpMI, OpIdx] : Ops) { + const MachineOperand &MO = OpMI->getOperand(OpIdx); + if (!MO.isDef()) + continue; + DefLanes |= MO.getSubReg() ? TRI.getSubRegIndexLaneMask(MO.getSubReg()) + : MRI.getMaxLaneMaskForVReg(Reg); + hasLiveDef |= !MO.isDead(); + } + bool PreserveSlotLanes = false; + if (RI.Writes && !RI.Reads && hasLiveDef && + StackInt->liveAt(Idx.getBaseIndex())) { + LaneBitmask OrigLanes = + liveLanesAt(LIS.getInterval(Original), Idx.getBaseIndex(), MRI); + PreserveSlotLanes = (OrigLanes & ~DefLanes).any(); + } + + if (RI.Reads || PreserveSlotLanes) insertReload(NewVReg, Idx, &MI); // Rewrite instruction operands. - bool hasLiveDef = false; for (const auto &OpPair : Ops) { MachineOperand &MO = OpPair.first->getOperand(OpPair.second); MO.setReg(NewVReg); if (MO.isUse()) { if (!OpPair.first->isRegTiedToDefOperand(OpPair.second)) MO.setIsKill(); - } else { - if (!MO.isDead()) - hasLiveDef = true; + } else if (PreserveSlotLanes) { + MO.setIsUndef(false); } } LLVM_DEBUG(dbgs() << "\trewrite: " << Idx << '\t' << MI << '\n'); diff --git a/llvm-patches/llvm-patch-content.txt b/llvm-patches/llvm-patch-content.txt index c64e4ff54e..d0dd7db6da 100644 --- a/llvm-patches/llvm-patch-content.txt +++ b/llvm-patches/llvm-patch-content.txt @@ -858,3 +858,78 @@ Applies on top of patch201186.patch and does not apply to a tree without it. Drop this patch on the next LLVM/Triton bump that includes upstream commit 077ec06128e27e3dce2b97563dab703f89024ddf. + +---------------------------------------------------------------------- +patch-inline-spiller-partial-sibling-spill.patch +---------------------------------------------------------------------- +- Upstream status : Downstream fix; submit to LLVM before the next bump +- Upstream title : [InlineSpiller] Do not store undefined lanes into a + shared spill slot +- File Fixed : llvm/lib/CodeGen/InlineSpiller.cpp +- Contents : Adds liveLanesAt(), which returns the lanes of a register + that hold a value at a slot index, counting a lane with no + subrange as dead. hoistSpillInsideBB now declines to hoist + a spill onto a sibling that leaves undefined a lane the + original value is live in at the sibling's def. When + spillAroundUses spills a def that writes only some lanes + without reading the rest, while the stack slot is live and + the original value is live in other lanes across the def, + it reloads the slot first and drops the undef flag, so the + full-width store writes those lanes back unchanged. Lanes + the original value never holds are free to carry garbage, + so upstream's hoisting of partially defined values + (splitkit.mir, splitkit-copy-bundle.mir) is unchanged. + +Symptom without the patch +------------------------- +A gfx950 attention kernel from the weekly parameterSweeps run (f16, g=5, +seq_len_q=1, seq_len_k=361, 4 Q heads over 2 KV heads, attention bias, +sliding_window_look_back=262, last_valid_kv_index=141,134,279,29,84; perf +config mPerBlockG0=64, nPerBlockG0=128, kPerBlock=64, numWaves=1, +numStages=3, wavesPerEU=2) dies with a GPU memory access fault on the page +16 KiB past the bias buffer. The kernel needs all 512 VGPRs, so SGPRs are +spilled into VGPR lanes. The K buffer descriptor is spilled before the KV +loop and reloaded at the top of every iteration, but inside the loop the +same four lanes are overwritten with the bias buffer's base address. From +the second iteration on, the K loads address the bias buffer with K's +offsets. + +Later rocMLIR changes altered that kernel's IR and it no longer reaches the +hazard. A bf16 kernel from the same sweep (g=6, seq_len_q=1, seq_len_k=151, +16 Q heads over 4 KV heads, attention scale and bias, transV, +sliding_window_look_back=150, last_valid_kv_index=120,143,78,69,87,150; perf +config mPerBlockG0=16, nPerBlockG0=128, nPerBlockG1=256, kPerBlock=512, +numWaves=16, matrixInstrNonkdim=16, numStages=1) faults the same way: a +descriptor rebuilt inside the KV loop from another's sub2_sub3 is stored +full-width over the spill slot of a descriptor reloaded at the top of every +iteration. + +Root cause +---------- +The bias descriptor borrows num_records and flags (sub2_sub3) from the K +descriptor. Splitting K's live range inside the loop therefore produced a +sibling defined only by `undef %a.sub2_sub3 = COPY %b.sub2_sub3`, since the +base (sub0_sub1) is no longer needed in a register there; its value lives +on in the stack slot that all siblings share, for the next iteration's +reload. When the following piece was spilled, spillAroundUses treated its +defining subregister copy as a sibling copy and hoistSpillInsideBB stored +the whole partial sibling into that slot, writing whatever the physical +registers held in sub0_sub1 over K's base. Without the hoist, insertSpill +would store the undefined lanes the same way. Upstream #177703 fixed the +same hazard for the cross-block hoist in isSpillCandBB; upstream main still +has the in-block path unchanged. The #177703 check only inspects subranges +that exist, so it would accept a partial copy like the one above, which has +no subrange for the lanes it does not define. It did not fire here and is +left as upstream wrote it, since it is also what hoists the partially +defined value in splitkit.mir out of its loop. + +CI coverage +----------- +mlir/test/e2e/PrAttentionBF16Gfx950.toml pins the bf16 shape, perf config +and chip counts above and runs them on gfx950 with -pv --pv-f64 on the fixed +input pattern, for both transO values. The kernel LLVM IR it hands the backend +is identical to the faulting sweep's, so without this patch the run dies +with the memory access fault and FileCheck fails on the missing verification +line. + +Drop or reconcile this patch once an equivalent fix lands upstream. diff --git a/llvm-patches/patch-inline-spiller-partial-sibling-spill.patch b/llvm-patches/patch-inline-spiller-partial-sibling-spill.patch new file mode 100644 index 0000000000..b19be78a22 --- /dev/null +++ b/llvm-patches/patch-inline-spiller-partial-sibling-spill.patch @@ -0,0 +1,87 @@ +diff --git a/llvm/lib/CodeGen/InlineSpiller.cpp b/llvm/lib/CodeGen/InlineSpiller.cpp +index 5ed3b3628..b5a2b6fbd 100644 +--- a/llvm/lib/CodeGen/InlineSpiller.cpp ++++ b/llvm/lib/CodeGen/InlineSpiller.cpp +@@ -320,6 +320,21 @@ static void getVDefInterval(const MachineInstr &MI, LiveIntervals &LIS) { + LIS.getInterval(MO.getReg()); + } + ++/// The lanes of LI's register that hold a value at Idx. Live range splitting ++/// copies only the lanes that are live, so a sibling may have lanes with no ++/// subrange at all. ++static LaneBitmask liveLanesAt(const LiveInterval &LI, SlotIndex Idx, ++ const MachineRegisterInfo &MRI) { ++ if (!LI.hasSubRanges()) ++ return LI.liveAt(Idx) ? MRI.getMaxLaneMaskForVReg(LI.reg()) ++ : LaneBitmask::getNone(); ++ LaneBitmask LiveLanes = LaneBitmask::getNone(); ++ for (const LiveInterval::SubRange &SR : LI.subranges()) ++ if (SR.liveAt(Idx)) ++ LiveLanes |= SR.LaneMask; ++ return LiveLanes; ++} ++ + /// isSnippet - Identify if a live interval is a snippet that should be spilled. + /// It is assumed that SnipLI is a virtual register with the same original as + /// Edit->getReg(). +@@ -454,6 +469,14 @@ bool InlineSpiller::hoistSpillInsideBB(LiveInterval &SpillLI, + if (DefMBB != CopyMI.getParent() || !SrcQ.isKill()) + return false; + ++ // The hoisted store writes all of SrcReg into the slot every sibling shares. ++ // A lane SrcReg leaves undefined while the original value is live in it ++ // would overwrite that value, which another sibling may still reload. ++ LaneBitmask OrigLanes = ++ liveLanesAt(LIS.getInterval(Original), SrcVNI->def, MRI); ++ if ((OrigLanes & ~liveLanesAt(SrcLI, SrcVNI->def, MRI)).any()) ++ return false; ++ + MachineBasicBlock *MBB = DefMBB; + MachineBasicBlock::iterator MII; + if (SrcVNI->isPHIDef()) +@@ -1368,20 +1391,40 @@ void InlineSpiller::spillAroundUses(Register Reg) { + // FIXME: Infer regclass from instruction alone. + Register NewVReg = Edit->createFrom(Reg); + +- if (RI.Reads) ++ // The spill below stores every lane of NewVReg. When MI defines only some ++ // of them without reading the rest, lanes the original value holds across ++ // MI would be overwritten in the slot, where a sibling may still reload ++ // them. Fill them from the slot first so the store writes that value back. ++ LaneBitmask DefLanes = LaneBitmask::getNone(); ++ bool hasLiveDef = false; ++ for (const auto &[OpMI, OpIdx] : Ops) { ++ const MachineOperand &MO = OpMI->getOperand(OpIdx); ++ if (!MO.isDef()) ++ continue; ++ DefLanes |= MO.getSubReg() ? TRI.getSubRegIndexLaneMask(MO.getSubReg()) ++ : MRI.getMaxLaneMaskForVReg(Reg); ++ hasLiveDef |= !MO.isDead(); ++ } ++ bool PreserveSlotLanes = false; ++ if (RI.Writes && !RI.Reads && hasLiveDef && ++ StackInt->liveAt(Idx.getBaseIndex())) { ++ LaneBitmask OrigLanes = ++ liveLanesAt(LIS.getInterval(Original), Idx.getBaseIndex(), MRI); ++ PreserveSlotLanes = (OrigLanes & ~DefLanes).any(); ++ } ++ ++ if (RI.Reads || PreserveSlotLanes) + insertReload(NewVReg, Idx, &MI); + + // Rewrite instruction operands. +- bool hasLiveDef = false; + for (const auto &OpPair : Ops) { + MachineOperand &MO = OpPair.first->getOperand(OpPair.second); + MO.setReg(NewVReg); + if (MO.isUse()) { + if (!OpPair.first->isRegTiedToDefOperand(OpPair.second)) + MO.setIsKill(); +- } else { +- if (!MO.isDead()) +- hasLiveDef = true; ++ } else if (PreserveSlotLanes) { ++ MO.setIsUndef(false); + } + } + LLVM_DEBUG(dbgs() << "\trewrite: " << Idx << '\t' << MI << '\n'); diff --git a/mlir/test/e2e/CMakeLists.txt b/mlir/test/e2e/CMakeLists.txt index b8ca74bbed..1d7be2c32f 100644 --- a/mlir/test/e2e/CMakeLists.txt +++ b/mlir/test/e2e/CMakeLists.txt @@ -39,6 +39,7 @@ if (ROCMLIR_DRIVER_PR_E2E_TEST_ENABLED) PrAttentionF32 PrAttentionF16 PrAttentionBF16 + PrAttentionBF16Gfx950 PrAttentionI8 PrGemmSplitK PrGemmElementwiseGemmF32 diff --git a/mlir/test/e2e/PrAttentionBF16Gfx950.cfg b/mlir/test/e2e/PrAttentionBF16Gfx950.cfg new file mode 100644 index 0000000000..5777ffdbcf --- /dev/null +++ b/mlir/test/e2e/PrAttentionBF16Gfx950.cfg @@ -0,0 +1,4 @@ +# The perf config requires 128 KiB of LDS, and the register allocation pinned +# down by these tests is the one produced for gfx950 +if "gfx950" not in config.arch: + config.unsupported = True diff --git a/mlir/test/e2e/PrAttentionBF16Gfx950.toml b/mlir/test/e2e/PrAttentionBF16Gfx950.toml new file mode 100644 index 0000000000..61a25a4a42 --- /dev/null +++ b/mlir/test/e2e/PrAttentionBF16Gfx950.toml @@ -0,0 +1,25 @@ +directory = "PrAttentionBF16Gfx950" +prefix = "rocmlir-gen" +# The fault does not depend on the input values. With random data this shape +# misses the bf16 tolerance against the CPU reference for any perf config, so +# the inputs are pinned to the fixed pattern instead of %random_data. +suffix = "--operation attention -t bf16 --arch %arch -pv --pv-f64 -rand fixed %rocmlir_gen_flags | rocmlir-driver --host-pipeline=highlevel | rocmlir-driver -c | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void | FileCheck %s --check-prefix=" + +[[axis]] +name = "transO" +values = ["true", "false"] +prefix = "--transO=" + +[[suite]] +name = "pr_attention_bf16_gfx950" + +# SGPRs spill into VGPR lanes, and inside the KV loop a buffer descriptor is +# rebuilt from num_records/flags (sub2_sub3) copied out of another one. That +# made InlineSpiller store the rebuilt descriptor full-width into the spill +# slot of a descriptor reloaded at the top of every iteration, replacing its +# base address with garbage, so the second iteration's loads fault +# (llvm-patches/patch-inline-spiller-partial-sibling-spill.patch). Changing +# the shape, perf config or chip counts changes register allocation and can +# hide the bug. +[[suite.test]] +config = "--num_cu 256 --num_chiplets 8 -g 6 -seq_len_q 1 -seq_len_k 151 -num_heads_q 16 -num_heads_kv 4 -head_dim_qk 196 -head_dim_v 183 --with-attn-scale --with-attn-bias --transV=true -sliding_window_look_back=150 -last_valid_kv_index=120,143,78,69,87,150 -perf_config attn:mPerBlockG0=16,nPerBlockG0=128,nPerBlockG1=256,kPerBlock=512,kpack=1,numCTAs=1,numWaves=16,matrixInstrNonkdim=16,splitKFactor=1,numStages=1,wavesPerEU=0,gridGroupSize=0" From ad5fa0a5076f967c646e776625a3166c8b460276 Mon Sep 17 00:00:00 2001 From: Muhamed Husic Date: Wed, 30 Sep 2026 14:04:58 -0500 Subject: [PATCH 2/2] Link the upstream issue in the InlineSpiller patch record --- llvm-patches/llvm-patch-content.txt | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/llvm-patches/llvm-patch-content.txt b/llvm-patches/llvm-patch-content.txt index d0dd7db6da..568f6d7d57 100644 --- a/llvm-patches/llvm-patch-content.txt +++ b/llvm-patches/llvm-patch-content.txt @@ -862,7 +862,8 @@ Drop this patch on the next LLVM/Triton bump that includes upstream commit ---------------------------------------------------------------------- patch-inline-spiller-partial-sibling-spill.patch ---------------------------------------------------------------------- -- Upstream status : Downstream fix; submit to LLVM before the next bump +- Upstream status : Downstream fix. Reported upstream as + https://github.com/llvm/llvm-project/issues/225054 - Upstream title : [InlineSpiller] Do not store undefined lanes into a shared spill slot - File Fixed : llvm/lib/CodeGen/InlineSpiller.cpp @@ -932,4 +933,5 @@ is identical to the faulting sweep's, so without this patch the run dies with the memory access fault and FileCheck fails on the missing verification line. -Drop or reconcile this patch once an equivalent fix lands upstream. +#225054 is still open upstream. Drop or reconcile this patch once a fix for +it lands in the pinned LLVM.