From 94875aea7eec805d54055f55fad99b57e0ad2317 Mon Sep 17 00:00:00 2001 From: Lucas Ramirez <11032120+lucas-rami@users.noreply.github.com> Date: Tue, 7 Apr 2026 01:26:16 +0200 Subject: [PATCH] [CodeGen] Fix multiple connected component issue in rematerializer (#186674) This fixes a rematerializer issue wherein re-creating the interval of a non-rematerializable super-register defined over multiple MIs, some of which defining entirely dead sub-registers, could cause a crash when changing the order of sub-definitions (for example during scheduling) because the re-created interval could end up with multiple connected components, which is illegal. The solution is to split separate components of the interval in such cases. The added unit test crashes without that added behavior. --- llvm/include/llvm/CodeGen/LiveIntervals.h | 6 ++ llvm/lib/CodeGen/Rematerializer.cpp | 17 ++++- llvm/unittests/CodeGen/RematerializerTest.cpp | 71 +++++++++++++++++++ 3 files changed, 93 insertions(+), 1 deletion(-) diff --git a/llvm/include/llvm/CodeGen/LiveIntervals.h b/llvm/include/llvm/CodeGen/LiveIntervals.h index b618e0b778ae..e79d5b32eae4 100644 --- a/llvm/include/llvm/CodeGen/LiveIntervals.h +++ b/llvm/include/llvm/CodeGen/LiveIntervals.h @@ -160,6 +160,12 @@ public: return LI; } + LiveInterval &createAndComputeVirtRegInterval(Register Reg, bool &NeedSplit) { + LiveInterval &LI = createEmptyInterval(Reg); + NeedSplit = computeVirtRegInterval(LI); + return LI; + } + /// Return an existing interval for \p Reg. /// If \p Reg has no interval then this creates a new empty one instead. /// Note: does not trigger interval computation. diff --git a/llvm/lib/CodeGen/Rematerializer.cpp b/llvm/lib/CodeGen/Rematerializer.cpp index 4125c36b4720..c6d90479300f 100644 --- a/llvm/lib/CodeGen/Rematerializer.cpp +++ b/llvm/lib/CodeGen/Rematerializer.cpp @@ -204,6 +204,9 @@ void Rematerializer::updateLiveIntervals() { Register DefReg = UpdateReg.getDefReg(); if (LIS.hasInterval(DefReg)) LIS.removeInterval(DefReg); + // Rematerializable registers have a single definition by construction so + // re-creating their interval cannot yield a live interval with multiple + // connected components. LIS.createAndComputeVirtRegInterval(DefReg); LLVM_DEBUG({ @@ -218,7 +221,19 @@ void Rematerializer::updateLiveIntervals() { if (!SeenUnrematRegs.insert(UnrematReg).second) continue; LIS.removeInterval(UnrematReg); - LIS.createAndComputeVirtRegInterval(UnrematReg); + bool NeedSplit = false; + + // Unrematerializable registers may end up with multiple connected + // components in their live interval after it is re-created. It needs to + // be split in such cases. We don't track unrematerializable registers by + // their actual register index (just by operand index) so we do not need + // to update any state in the rematerializer. + LiveInterval &LI = + LIS.createAndComputeVirtRegInterval(UnrematReg, NeedSplit); + if (NeedSplit) { + SmallVector SplitLIs; + LIS.splitSeparateComponents(LI, SplitLIs); + } LLVM_DEBUG( dbgs() << " Re-computed interval for register " << printReg(UnrematReg, &TRI, diff --git a/llvm/unittests/CodeGen/RematerializerTest.cpp b/llvm/unittests/CodeGen/RematerializerTest.cpp index 6b1373377816..3ca6b02a7538 100644 --- a/llvm/unittests/CodeGen/RematerializerTest.cpp +++ b/llvm/unittests/CodeGen/RematerializerTest.cpp @@ -590,6 +590,77 @@ body: | EXPECT_TRUE(getMF().verify()); } +/// The rematerializer had a bug where re-creating the interval of a +/// non-rematerializable super-register defined over multiple MIs, some of which +/// defining entirely dead subregisters, could cause a crash when changing the +/// order of sub-definitions (for example during scheduling) because the +/// re-created interval could end up with multiple connected components, which +/// is illegal. The solution is to split separate components of the interval in +/// such cases. +TEST_F(RematerializerTest, SplitSubRegDeadDef) { + StringRef MIR = R"( +name: SplitSubRegDeadDef +tracksRegLiveness: true +machineFunctionInfo: + isEntryFunction: true +body: | + bb.0: + undef %0.sub0:vreg_64 = IMPLICIT_DEF + %0.sub1:vreg_64 = IMPLICIT_DEF + %1:vgpr_32 = V_ADD_U32_e32 %0.sub0, %0.sub0, implicit $exec + + bb.1: + S_NOP 0, implicit %1 + S_ENDPGM 0 +... +)"; + ASSERT_TRUE(parseMIRAndInit(MIR, "SplitSubRegDeadDef")); + LiveIntervals &LIS = MFAM.getResult(*MF); + + // Replicates the scheduler's effect on LIS on an intra-block move of MI. + auto MoveMIAndAdjustLiveness = [&](MachineInstr &MI) { + LIS.handleMove(MI); + const MachineRegisterInfo &MRI = MF->getRegInfo(); + const TargetRegisterInfo &TRI = *MF->getSubtarget().getRegisterInfo(); + RegisterOperands RegOpers; + RegOpers.collect(MI, TRI, MRI, true, /*IgnoreDead=*/false); + SlotIndex Sub1Slot = LIS.getInstructionIndex(MI).getRegSlot(); + RegOpers.adjustLaneLiveness(LIS, MRI, Sub1Slot, &MI); + }; + + MachineBasicBlock &MBB0 = *MF->getBlockNumbered(0); + MachineInstr &Sub0Def = *MBB0.begin(); + MachineInstr &Sub1Def = *MBB0.begin()->getNextNode(); + + // Flip %0's subdefinition order. After the move, the definitions look like: + // undef %0.sub1:vreg_64 = IMPLICIT_DEF + // undef %0.sub0:vreg_64 = IMPLICIT_DEF + MBB0.splice(Sub0Def.getIterator(), &MBB0, Sub1Def.getIterator()); + MoveMIAndAdjustLiveness(Sub1Def); + + // Rematerialize %1 to bb.1. This triggers a live-interval update of %0 when + // calling Remater.updateLiveIntervals(), during which its interval is split. + Rematerializer &Remater = getRematerializer(); + Rematerializer::DependencyReuseInfo DRI; + const unsigned MBB1 = 1; + const RegisterIdx Add = 0; + Remater.rematerializeToRegion(Add, MBB1, DRI); + Remater.updateLiveIntervals(); + + // If we didn't split %0 before, its definitions would now look like: + // dead undef %0.sub1:vreg_64 = IMPLICIT_DEF + // undef %0.sub0:vreg_64 = IMPLICIT_DEF + // + // Trying to flip back %0's definition order then triggers an + // error in LIS.handleMove because its live interval is made up of multiple + // connected components. + ASSERT_NE(Sub0Def.getOperand(0).getReg(), Sub1Def.getOperand(0).getReg()); + MBB0.splice(MBB0.end(), &MBB0, Sub1Def.getIterator()); + MoveMIAndAdjustLiveness(Sub1Def); + + EXPECT_TRUE(getMF().verify()); +} + /// Checks that rollback works as expected when the rollback listener is added /// mid-rematerializations. TEST_F(RematerializerTest, Rollback) {