[llvm-branch-commits] [llvm] [CodeGen] Allow rematerializer to rematerialize at the end of a block (PR #184339)
Lucas Ramirez via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Mar 3 05:43:43 PST 2026
https://github.com/lucas-rami created https://github.com/llvm/llvm-project/pull/184339
This makes the rematerializer able to rematerialize MIs at the end of a basic block. We achieve this by tracking the parent basic block of every region inside the rematerializer and adding an explicit target region to some of the class's methods. The latter removes the requirement that we track the MI of every region (`Rematerializer::MIRegion`) after the analysis phase; the class member is therefore deleted.
This new ability will be used shortly to improve the design of the rollback mechanism.
>From 8b1ec1c40b63ea351e7323fe287e8ecf7c987a58 Mon Sep 17 00:00:00 2001
From: Lucas Ramirez <lucas.rami at proton.me>
Date: Thu, 26 Feb 2026 16:41:11 +0000
Subject: [PATCH] [CodeGen] Allow rematerializer to rematerialize at the end of
a block
This makes the rematerializer able to rematerialize MIs at the end of a
basic block. We achive this by tracking the parent basic block of every
region inside the rematerializer and adding an explicit target region to
some of the class's methods. The latter removes the requirement that we
track the MI of every region (`Rematerializer::MIRegion`) after the
analysis phase; the class member is therefore deleted.
This new ability will be used shortly to improve the design of the
rollback mechanism.
---
llvm/include/llvm/CodeGen/Rematerializer.h | 59 +++++++++---------
llvm/lib/CodeGen/Rematerializer.cpp | 60 +++++++++++--------
llvm/unittests/CodeGen/RematerializerTest.cpp | 30 +++++-----
3 files changed, 83 insertions(+), 66 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/Rematerializer.h b/llvm/include/llvm/CodeGen/Rematerializer.h
index a04dedb36a6b7..11ac95fc9368d 100644
--- a/llvm/include/llvm/CodeGen/Rematerializer.h
+++ b/llvm/include/llvm/CodeGen/Rematerializer.h
@@ -332,9 +332,10 @@ class Rematerializer {
};
/// Rematerializes register \p RootIdx just before its first user inside
- /// region \p UseRegion, transfers all its users in the region to the new
- /// register, and returns the latter's index. The root's dependency DAG is
- /// rematerialized or re-used according to \p DRI.
+ /// region \p UseRegion (or at the end of the region if it has no user),
+ /// transfers all its users in the region to the new register, and returns the
+ /// latter's index. The root's dependency DAG is rematerialized or re-used
+ /// according to \p DRI.
///
/// When the method returns, \p DRI contains additional entries for non-root
/// registers of the root's dependency DAG that needed to be rematerialized
@@ -343,15 +344,15 @@ class Rematerializer {
RegisterIdx rematerializeToRegion(RegisterIdx RootIdx, unsigned UseRegion,
DependencyReuseInfo &DRI);
- /// Rematerializes register \p RootIdx before position \p InsertPos and
- /// returns the new register's index. The root's dependency DAG is
- /// rematerialized or re-used according to \p DRI.
+ /// Rematerializes register \p RootIdx before position \p InsertPos in \p
+ /// UseRegion and returns the new register's index. The root's dependency DAG
+ /// is rematerialized or re-used according to \p DRI.
///
/// When the method returns, \p DRI contains additional entries for non-root
/// registers of the root's dependency DAG that needed to be rematerialized
/// along the root. References to \ref Rematerializer::Reg should be
/// considered invalidated by calls to this method.
- RegisterIdx rematerializeToPos(RegisterIdx RootIdx,
+ RegisterIdx rematerializeToPos(RegisterIdx RootIdx, unsigned UseRegion,
MachineBasicBlock::iterator InsertPos,
DependencyReuseInfo &DRI);
@@ -384,12 +385,12 @@ class Rematerializer {
void transferRegionUsers(RegisterIdx FromRegIdx, RegisterIdx ToRegIdx,
unsigned UseRegion);
- /// Transfers user \p UserMI from register \p FromRegIdx to \p ToRegIdx,
- /// the latter of which must be a rematerialization of the former or have the
- /// same origin register. \p UserMI must be a direct user of \p FromRegIdx. \p
- /// UserMI must be reachable from \p ToRegIdx.
+ /// Transfers user \p UserMI in region \p UserRegion from register \p
+ /// FromRegIdx to \p ToRegIdx, the latter of which must be a rematerialization
+ /// of the former or have the same origin register. \p UserMI must be a direct
+ /// user of \p FromRegIdx. \p UserMI must be reachable from \p ToRegIdx.
void transferUser(RegisterIdx FromRegIdx, RegisterIdx ToRegIdx,
- MachineInstr &UserMI);
+ unsigned UserRegion, MachineInstr &UserMI);
/// Recomputes all live intervals that have changed as a result of previous
/// rematerializations/rollbacks.
@@ -455,9 +456,8 @@ class Rematerializer {
/// data in the \ref Regs vector. This includes registers that no longer exist
/// in the MIR.
DenseMap<Register, RegisterIdx> RegToIdx;
- /// Maps all MIs to their parent region. Region terminators are considered
- /// part of the region they terminate.
- DenseMap<MachineInstr *, unsigned> MIRegion;
+ /// Parent block of each region, in order.
+ SmallVector<MachineBasicBlock *> RegionMBB;
/// Set of registers whose live-range may have changed during past
/// rematerializations/rollbacks.
DenseSet<RegisterIdx> LISUpdates;
@@ -470,8 +470,13 @@ class Rematerializer {
bool SupportRollback = false;
/// During the analysis phase, creates a \ref Rematerializer::Reg object for
- /// virtual register \p VirtRegIdx if it
- void addRegIfRematerializable(unsigned VirtRegIdx, BitVector &SeenRegs);
+ /// virtual register \p VirtRegIdx if it is rematerializable. \p MIRegion maps
+ /// all MIs to their parent region. Set bits in \p SeenRegs indicate virtual
+ /// register indices that have already been visited.
+ void
+ addRegIfRematerializable(unsigned VirtRegIdx,
+ const DenseMap<MachineInstr *, unsigned> &MIRegion,
+ BitVector &SeenRegs);
/// Determines whether \p MI is considered rematerializable. This further
/// restricts constraints imposed by the TII on rematerializable instructions,
@@ -479,16 +484,16 @@ class Rematerializer {
/// defined once.
bool isMIRematerializable(const MachineInstr &MI) const;
- /// Rematerializes register \p RegIdx at \p InsertPos, adding the new
- /// rematerializable register to the backing vector \ref Regs and returning
- /// its index inside the vector. Sets the new registers' rematerializable
- /// dependencies to \p Dependencies (these are assumed to already exist in the
- /// MIR) and its unrematerializable dependencies to the same as \p RegIdx. The
- /// new register initially has no user. Since the method appends to \ref Regs,
- /// references to elements within it should be considered invalidated across
- /// calls to this method unless the vector can be guaranteed to have enough
- /// space for an extra element.
- RegisterIdx rematerializeReg(RegisterIdx RegIdx,
+ /// Rematerializes register \p RegIdx at \p InsertPos in \p UseRegion, adding
+ /// the new rematerializable register to the backing vector \ref Regs and
+ /// returning its index inside the vector. Sets the new registers'
+ /// rematerializable dependencies to \p Dependencies (these are assumed to
+ /// already exist in the MIR) and its unrematerializable dependencies to the
+ /// same as \p RegIdx. The new register initially has no user. Since the
+ /// method appends to \ref Regs, references to elements within it should be
+ /// considered invalidated across calls to this method unless the vector can
+ /// be guaranteed to have enough space for an extra element.
+ RegisterIdx rematerializeReg(RegisterIdx RegIdx, unsigned UseRegion,
MachineBasicBlock::iterator InsertPos,
SmallVectorImpl<Reg::Dependency> &&Dependencies);
diff --git a/llvm/lib/CodeGen/Rematerializer.cpp b/llvm/lib/CodeGen/Rematerializer.cpp
index 739c59d53c189..e07e7d0745d91 100644
--- a/llvm/lib/CodeGen/Rematerializer.cpp
+++ b/llvm/lib/CodeGen/Rematerializer.cpp
@@ -74,18 +74,22 @@ RegisterIdx Rematerializer::rematerializeToRegion(RegisterIdx RootIdx,
DependencyReuseInfo &DRI) {
MachineInstr *FirstMI =
getReg(RootIdx).getRegionUseBounds(UseRegion, LIS).first;
- RegisterIdx NewRegIdx = rematerializeToPos(RootIdx, FirstMI, DRI);
+ // If there are no users in the region, rematerialize the register at the very
+ // end of the region.
+ MachineBasicBlock::iterator InsertPos =
+ FirstMI ? FirstMI : Regions[UseRegion].second;
+ RegisterIdx NewRegIdx =
+ rematerializeToPos(RootIdx, UseRegion, InsertPos, DRI);
transferRegionUsers(RootIdx, NewRegIdx, UseRegion);
return NewRegIdx;
}
RegisterIdx
-Rematerializer::rematerializeToPos(RegisterIdx RootIdx,
+Rematerializer::rematerializeToPos(RegisterIdx RootIdx, unsigned UseRegion,
MachineBasicBlock::iterator InsertPos,
DependencyReuseInfo &DRI) {
assert(!DRI.DependencyMap.contains(RootIdx));
- LLVM_DEBUG(dbgs() << "Rematerializing " << printID(RootIdx) << " to "
- << printUser(&*InsertPos) << '\n');
+ LLVM_DEBUG(dbgs() << "Rematerializing " << printID(RootIdx) << '\n');
// Traverse the root's dependency DAG depth-first to find the set of
// registers we must rematerialize along with it and a legal order to
@@ -115,7 +119,8 @@ Rematerializer::rematerializeToPos(RegisterIdx RootIdx,
SmallVector<Reg::Dependency, 2> Dependencies;
for (const Reg::Dependency &Dep : getReg(RegIdx).Dependencies)
Dependencies.emplace_back(Dep.MOIdx, DRI.DependencyMap.at(Dep.RegIdx));
- LastNewIdx = rematerializeReg(RegIdx, InsertPos, std::move(Dependencies));
+ LastNewIdx =
+ rematerializeReg(RegIdx, UseRegion, InsertPos, std::move(Dependencies));
DRI.DependencyMap.insert({RegIdx, LastNewIdx});
}
@@ -203,9 +208,12 @@ void Rematerializer::reviveRegIfDead(RegisterIdx RootIdx) {
}
void Rematerializer::transferUser(RegisterIdx FromRegIdx, RegisterIdx ToRegIdx,
- MachineInstr &UserMI) {
+ unsigned UserRegion, MachineInstr &UserMI) {
+ assert(getReg(FromRegIdx).Uses.contains(UserRegion) && "no user in region");
+ assert(getReg(FromRegIdx).Uses.at(UserRegion).contains(&UserMI) &&
+ "not a region user");
+
transferUserImpl(FromRegIdx, ToRegIdx, UserMI);
- unsigned UserRegion = MIRegion.at(&UserMI);
Regs[FromRegIdx].eraseUser(&UserMI, UserRegion);
Regs[ToRegIdx].addUser(&UserMI, UserRegion);
deleteRegIfUnused(FromRegIdx);
@@ -230,9 +238,6 @@ void Rematerializer::transferRegionUsers(RegisterIdx FromRegIdx,
void Rematerializer::transferUserImpl(RegisterIdx FromRegIdx,
RegisterIdx ToRegIdx,
MachineInstr &UserMI) {
- assert(MIRegion.contains(&UserMI) && "unknown user");
- assert(getReg(FromRegIdx).Uses.at(MIRegion.at(&UserMI)).contains(&UserMI) &&
- "not a user");
assert(FromRegIdx != ToRegIdx && "identical registers");
assert(getOriginOrSelf(FromRegIdx) == getOriginOrSelf(ToRegIdx) &&
"unrelated registers");
@@ -413,7 +418,6 @@ void Rematerializer::deleteReg(RegisterIdx RegIdx) {
RegionBegin = std::next(MachineBasicBlock::iterator(DeleteReg.DefMI));
LIS.RemoveMachineInstrFromMaps(*DeleteReg.DefMI);
DeleteReg.DefMI->eraseFromParent();
- MIRegion.erase(DeleteReg.DefMI);
DeleteReg.DefMI = nullptr;
}
@@ -445,7 +449,7 @@ bool Rematerializer::analyze(bool SupportRollback) {
UnrematableOprds.clear();
Origins.clear();
Rematerializations.clear();
- MIRegion.clear();
+ RegionMBB.clear();
RegToIdx.clear();
LISUpdates.clear();
Revivable.clear();
@@ -453,7 +457,12 @@ bool Rematerializer::analyze(bool SupportRollback) {
if (Regions.empty())
return false;
+ /// Maps all MIs to their parent region. Region terminators are considered
+ /// part of the region they terminate.
+ DenseMap<MachineInstr *, unsigned> MIRegion;
+
// Initialize MI to containing region mapping.
+ RegionMBB.reserve(Regions.size());
for (unsigned I = 0, E = Regions.size(); I < E; ++I) {
RegionBoundaries Region = Regions[I];
assert(Region.first != Region.second && "empty cannot be region");
@@ -461,9 +470,11 @@ bool Rematerializer::analyze(bool SupportRollback) {
assert(!MIRegion.contains(&*MI) && "regions should not intersect");
MIRegion.insert({&*MI, I});
}
+ MachineBasicBlock &MBB = *Region.first->getParent();
+ RegionMBB.push_back(&MBB);
// A terminator instruction is considered part of the region it terminates.
- if (Region.second != Region.first->getParent()->end()) {
+ if (Region.second != MBB.end()) {
MachineInstr *RegionTerm = &*Region.second;
assert(!MIRegion.contains(RegionTerm) && "regions should not intersect");
MIRegion.insert({RegionTerm, I});
@@ -474,7 +485,7 @@ bool Rematerializer::analyze(bool SupportRollback) {
BitVector SeenRegs(NumVirtRegs);
for (unsigned I = 0, E = NumVirtRegs; I != E; ++I) {
if (!SeenRegs[I])
- addRegIfRematerializable(I, SeenRegs);
+ addRegIfRematerializable(I, MIRegion, SeenRegs);
}
assert(Regs.size() == UnrematableOprds.size());
@@ -485,8 +496,9 @@ bool Rematerializer::analyze(bool SupportRollback) {
return !Regs.empty();
}
-void Rematerializer::addRegIfRematerializable(unsigned VirtRegIdx,
- BitVector &SeenRegs) {
+void Rematerializer::addRegIfRematerializable(
+ unsigned VirtRegIdx, const DenseMap<MachineInstr *, unsigned> &MIRegion,
+ BitVector &SeenRegs) {
assert(!SeenRegs[VirtRegIdx] && "register already seen");
Register DefReg = Register::index2VirtReg(VirtRegIdx);
SeenRegs.set(VirtRegIdx);
@@ -530,7 +542,7 @@ void Rematerializer::addRegIfRematerializable(unsigned VirtRegIdx,
continue;
unsigned DepRegIdx = DepReg.virtRegIndex();
if (!SeenRegs[DepRegIdx])
- addRegIfRematerializable(DepRegIdx, SeenRegs);
+ addRegIfRematerializable(DepRegIdx, MIRegion, SeenRegs);
if (auto DepIt = RegToIdx.find(DepReg); DepIt != RegToIdx.end())
RematReg.Dependencies.push_back(Reg::Dependency(MOIdx, DepIt->second));
else
@@ -575,12 +587,12 @@ RegisterIdx Rematerializer::getDefRegIdx(const MachineInstr &MI) const {
}
RegisterIdx Rematerializer::rematerializeReg(
- RegisterIdx RegIdx, MachineBasicBlock::iterator InsertPos,
+ RegisterIdx RegIdx, unsigned UseRegion,
+ MachineBasicBlock::iterator InsertPos,
SmallVectorImpl<Reg::Dependency> &&Dependencies) {
- RegisterIdx NewRegIdx = Regs.size();
+ const RegisterIdx NewRegIdx = Regs.size();
notifyListeners(&Listener::beforeRegRematerialized, RegIdx, NewRegIdx);
- unsigned UseRegion = MIRegion.at(&*InsertPos);
Reg &NewReg = Regs.emplace_back();
Reg &FromReg = Regs[RegIdx];
NewReg.Mask = FromReg.Mask;
@@ -590,15 +602,14 @@ RegisterIdx Rematerializer::rematerializeReg(
// Track rematerialization link between registers. Origins are always
// registers that existed originally, and rematerializations are always
// attached to them.
- RegisterIdx OriginIdx =
- isRematerializedRegister(RegIdx) ? getOriginOf(RegIdx) : RegIdx;
+ const RegisterIdx OriginIdx = getOriginOrSelf(RegIdx);
Origins.push_back(OriginIdx);
Rematerializations[OriginIdx].insert(NewRegIdx);
// Use the TII to rematerialize the defining instruction with a new defined
// register.
Register NewDefReg = MRI.cloneVirtualRegister(FromReg.getDefReg());
- TII.reMaterialize(*InsertPos->getParent(), InsertPos, NewDefReg, 0,
+ TII.reMaterialize(*RegionMBB[UseRegion], InsertPos, NewDefReg, 0,
*FromReg.DefMI);
NewReg.DefMI = &*std::prev(InsertPos);
RegToIdx.insert({NewDefReg, NewRegIdx});
@@ -608,7 +619,6 @@ RegisterIdx Rematerializer::rematerializeReg(
if (Bounds.first == std::next(MachineBasicBlock::iterator(NewReg.DefMI)))
Bounds.first = NewReg.DefMI;
LIS.InsertMachineInstrInMaps(*NewReg.DefMI);
- MIRegion.emplace_or_assign(NewReg.DefMI, UseRegion);
LISUpdates.insert(NewRegIdx);
// Replace dependencies as needed in the rematerialized MI. All dependencies
@@ -785,7 +795,7 @@ Printable Rematerializer::printUser(const MachineInstr *MI) const {
if (RegIdx != NoReg)
OS << printID(RegIdx);
else
- OS << "(-/-)[" << MIRegion.at(MI) << ']';
+ OS << "(-/-)[?]";
OS << ' ';
MI->print(OS, /*IsStandalone=*/true, /*SkipOpers=*/false,
/*SkipDebugLoc=*/false, /*AddNewLine=*/false);
diff --git a/llvm/unittests/CodeGen/RematerializerTest.cpp b/llvm/unittests/CodeGen/RematerializerTest.cpp
index a697a4a7f3312..ca2bc3b86d47c 100644
--- a/llvm/unittests/CodeGen/RematerializerTest.cpp
+++ b/llvm/unittests/CodeGen/RematerializerTest.cpp
@@ -271,7 +271,7 @@ body: |
DRI.clear().reuse(Cst0).reuse(Cst1);
const RegisterIdx RematAdd01 =
- Remater.rematerializeToPos(/*RootIdx=*/Add01, NopMI, DRI);
+ Remater.rematerializeToPos(Add01, MBB1, NopMI, DRI);
// This adds an additional user to the used constants, and does not change
// existing users for the original register.
EXPECT_NO_USERS(RematAdd01);
@@ -281,14 +281,14 @@ body: |
DRI.clear();
const RegisterIdx RematCst3 =
- Remater.rematerializeToPos(/*RootIdx=*/Cst3, NopMI, DRI);
+ Remater.rematerializeToPos(Cst3, MBB1, NopMI, DRI);
// This does not change existing users for the original register.
EXPECT_NO_USERS(RematCst3);
EXPECT_NUM_USERS(Cst3, 1);
DRI.clear().useRemat(Add01, RematAdd01).useRemat(Cst3, RematCst3);
const RegisterIdx RematAdd23 =
- Remater.rematerializeToPos(/*RootIdx=*/Add23, NopMI, DRI);
+ Remater.rematerializeToPos(Add23, MBB1, NopMI, DRI);
// This adds a user to used rematerializations, and does not change existing
// users for the original register.
EXPECT_NO_USERS(RematAdd23);
@@ -298,7 +298,7 @@ body: |
// Finally transfer the NOP user from the original to the rematerialized
// register.
- Remater.transferUser(Add23, RematAdd23, *NopMI);
+ Remater.transferUser(Add23, RematAdd23, MBB1, *NopMI);
EXPECT_NO_USERS(Add23);
EXPECT_NUM_USERS(RematAdd23, 1);
@@ -516,24 +516,26 @@ body: |
// regions are empty. %2's region ends at the end of its parent block, whereas
// %3's region ends at a terminator MI (S_BRANCH).
Remater.rematerializeToRegion(/*RootIdx=*/Cst2, /*UseRegion=*/MBB3, DRI);
- Remater.rematerializeToRegion(/*RootIdx=*/Cst3, /*UseRegion=*/MBB3, DRI);
+ Remater.rematerializeToRegion(/*RootIdx=*/Cst3, /*UseRegion=*/MBB3,
+ DRI.clear());
RegionSizes[MBB1] -= 1;
RegionSizes[MBB2] -= 1;
RegionSizes[MBB3] += 2;
ASSERT_REGION_SIZES(RegionSizes);
- // We can move %0 and %1 to bb.2 because the boundary of bb.2's region points
- // to the region terminator (S_BRANCH), which is a valid position to insert
- // before. We couldn't move them to bb.1 however, since after %2 is
- // rematerialized there is no MI left to reference inside the region.
- RegisterIdx RematCst0 =
- Remater.rematerializeToPos(/*RootIdx=*/Cst0, (*Regions)[MBB2].first, DRI);
- RegisterIdx RematCst1 =
- Remater.rematerializeToPos(/*RootIdx=*/Cst1, (*Regions)[MBB2].first, DRI);
+ // Move %0 to the empty MBB1 block/region.
+ const RegisterIdx RematCst0 =
+ Remater.rematerializeToRegion(Cst0, MBB1, DRI.clear());
Remater.transferRegionUsers(Cst0, RematCst0, MBB3);
+
+ // Move %1 to the empty MBB2 region, right before the S_BRANCH terminator.
+ const RegisterIdx RematCst1 = Remater.rematerializeToPos(
+ Cst1, MBB2, (*Regions)[MBB2].first, DRI.clear());
Remater.transferRegionUsers(Cst1, RematCst1, MBB3);
+
RegionSizes[MBB0] -= 2;
- RegionSizes[MBB2] += 2;
+ RegionSizes[MBB1] += 1;
+ RegionSizes[MBB2] += 1;
ASSERT_REGION_SIZES(RegionSizes);
Remater.updateLiveIntervals();
More information about the llvm-branch-commits
mailing list