Doomhowl-Interactive/llvm-project
0
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity
Files
main
llvm-project/llvm/lib/CodeGen/Rematerializer.cpp
Lucas Ramirez 77a360663d [CodeGen] Fix incorrect index in rematerialization tracking (#194387) 
When deleting the last rematerialization of a register, we should delete
the rematerializer's remat tracking map's entry that corresponds to the
index of the *original* register, not the rematerialized register.

The existing typo has no impact on correctness at the moment because
entries with rematerialized register indices are never created (so there
is nothing to erase), and having an empty set in a value does not break
any code invariant; it just wastes memory.

Assisted-by: Claude Code
2026-04-27 17:12:53 +02:00

834 lines
31 KiB
C++
Raw Permalink Blame History

//=====-- Rematerializer.cpp - MIR rematerialization support ----*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//==-----------------------------------------------------------------------===//
//
/// \file
/// Implements helpers for target-independent rematerialization at the MIR
/// level.
//
//===----------------------------------------------------------------------===//
#include "llvm/CodeGen/Rematerializer.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/CodeGen/LiveIntervals.h"
#include "llvm/CodeGen/MachineBasicBlock.h"
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/Register.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/Support/Debug.h"
#include <optional>
#define DEBUG_TYPE "rematerializer"
using namespace llvm;
using RegisterIdx = Rematerializer::RegisterIdx;
// Pin the vtable to this file.
void Rematerializer::Listener::anchor() {}
/// Checks whether the value in \p LI at \p UseIdx is identical to \p OVNI (this
/// implies it is also live there). When \p LI has sub-ranges, checks that
/// all sub-ranges intersecting with \p Mask are also live at \p UseIdx.
static bool isIdenticalAtUse(const VNInfo &OVNI, LaneBitmask Mask,
SlotIndex UseIdx, const LiveInterval &LI) {
if (&OVNI != LI.getVNInfoAt(UseIdx))
return false;
if (LI.hasSubRanges()) {
// Check that intersecting subranges are live at user.
for (const LiveInterval::SubRange &SR : LI.subranges()) {
if ((SR.LaneMask & Mask).none())
continue;
if (!SR.liveAt(UseIdx))
return false;
// Early exit if all used lanes are checked. No need to continue.
Mask &= ~SR.LaneMask;
if (Mask.none())
break;
}
}
return true;
}
/// If \p MO is a virtual read register, returns it. Otherwise returns the
/// sentinel register.
static Register getRegDependency(const MachineOperand &MO) {
if (!MO.isReg() || !MO.readsReg())
return Register();
Register Reg = MO.getReg();
if (Reg.isPhysical()) {
// By the requirements on trivially rematerializable instructions, a
// physical register use is either constant or ignorable.
return Register();
}
return Reg;
}
RegisterIdx Rematerializer::rematerializeToRegion(RegisterIdx RootIdx,
unsigned UseRegion,
DependencyReuseInfo &DRI) {
MachineInstr *FirstMI =
getReg(RootIdx).getRegionUseBounds(UseRegion, LIS).first;
// 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, unsigned UseRegion,
MachineBasicBlock::iterator InsertPos,
DependencyReuseInfo &DRI) {
assert(!DRI.DependencyMap.contains(RootIdx));
LLVM_DEBUG(dbgs() << "Rematerializing " << printID(RootIdx) << '\n');
SmallVector<Reg::Dependency, 2> NewDeps;
// Copy all dependencies because recursive rematerialization of dependencies
// may invalidate references to the backing vector of registers.
SmallVector<Reg::Dependency, 2> OldDeps(getReg(RootIdx).Dependencies);
for (const Reg::Dependency &Dep : OldDeps) {
// Recursively rematerialize required dependencies at the same position as
// the root. Registers form a DAG so the recursion is guaranteed to
// terminate.
auto RematIdx = DRI.DependencyMap.find(Dep.RegIdx);
RegisterIdx NewDepRegIdx;
if (RematIdx == DRI.DependencyMap.end())
NewDepRegIdx = rematerializeToPos(Dep.RegIdx, UseRegion, InsertPos, DRI);
else
NewDepRegIdx = RematIdx->second;
NewDeps.emplace_back(Dep.MOIdx, NewDepRegIdx);
}
RegisterIdx NewIdx =
rematerializeReg(RootIdx, UseRegion, InsertPos, std::move(NewDeps));
DRI.DependencyMap.insert({RootIdx, NewIdx});
return NewIdx;
}
void Rematerializer::transferUser(RegisterIdx FromRegIdx, RegisterIdx ToRegIdx,
unsigned UserRegion, MachineInstr &UserMI) {
transferUserImpl(FromRegIdx, ToRegIdx, UserMI);
Regs[FromRegIdx].eraseUser(&UserMI, UserRegion);
Regs[ToRegIdx].addUser(&UserMI, UserRegion);
deleteRegIfUnused(FromRegIdx);
}
void Rematerializer::transferRegionUsers(RegisterIdx FromRegIdx,
RegisterIdx ToRegIdx,
unsigned UseRegion) {
auto &FromRegUsers = Regs[FromRegIdx].Uses;
auto UsesIt = FromRegUsers.find(UseRegion);
if (UsesIt == FromRegUsers.end())
return;
const SmallDenseSet<MachineInstr *, 4> &RegionUsers = UsesIt->getSecond();
for (MachineInstr *UserMI : RegionUsers)
transferUserImpl(FromRegIdx, ToRegIdx, *UserMI);
Regs[ToRegIdx].addUsers(RegionUsers, UseRegion);
FromRegUsers.erase(UseRegion);
deleteRegIfUnused(FromRegIdx);
}
void Rematerializer::transferAllUsers(RegisterIdx FromRegIdx,
RegisterIdx ToRegIdx) {
Reg &FromReg = Regs[FromRegIdx], &ToReg = Regs[ToRegIdx];
for (const auto &[UseRegion, RegionUsers] : FromReg.Uses) {
for (MachineInstr *UserMI : RegionUsers)
transferUserImpl(FromRegIdx, ToRegIdx, *UserMI);
ToReg.addUsers(RegionUsers, UseRegion);
}
FromReg.Uses.clear();
deleteRegIfUnused(FromRegIdx);
}
void Rematerializer::transferUserImpl(RegisterIdx FromRegIdx,
RegisterIdx ToRegIdx,
MachineInstr &UserMI) {
assert(FromRegIdx != ToRegIdx && "identical registers");
assert(getOriginOrSelf(FromRegIdx) == getOriginOrSelf(ToRegIdx) &&
"unrelated registers");
LLVM_DEBUG(dbgs() << "User transfer from " << printID(FromRegIdx) << " to "
<< printID(ToRegIdx) << ": " << printUser(&UserMI) << '\n');
UserMI.substituteRegister(getReg(FromRegIdx).getDefReg(),
getReg(ToRegIdx).getDefReg(), 0, TRI);
LISUpdates.insert(FromRegIdx);
LISUpdates.insert(ToRegIdx);
// If the user is rematerializable, we must change its dependency to the
// new register.
if (RegisterIdx UserRegIdx = getDefRegIdx(UserMI); UserRegIdx != NoReg) {
// Look for the user's dependency that matches the register.
for (Reg::Dependency &Dep : Regs[UserRegIdx].Dependencies) {
if (Dep.RegIdx == FromRegIdx) {
Dep.RegIdx = ToRegIdx;
return;
}
}
llvm_unreachable("broken dependency");
}
}
void Rematerializer::updateLiveIntervals() {
DenseSet<Register> SeenUnrematRegs;
for (RegisterIdx RegIdx : LISUpdates) {
const Reg &UpdateReg = getReg(RegIdx);
assert(UpdateReg.isAlive() && "dead register");
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({
dbgs() << "Re-computed interval for " << printID(RegIdx) << ": ";
LIS.getInterval(DefReg).print(dbgs());
dbgs() << '\n' << printRegUsers(RegIdx);
});
// Update intervals for unrematerializable operands.
for (unsigned MOIdx : getUnrematableOprds(RegIdx)) {
Register UnrematReg = UpdateReg.DefMI->getOperand(MOIdx).getReg();
if (!SeenUnrematRegs.insert(UnrematReg).second)
continue;
LIS.removeInterval(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<LiveInterval *> SplitLIs;
LIS.splitSeparateComponents(LI, SplitLIs);
}
LLVM_DEBUG(
dbgs() << " Re-computed interval for register "
<< printReg(UnrematReg, &TRI,
UpdateReg.DefMI->getOperand(MOIdx).getSubReg(),
&MRI)
<< '\n');
}
}
LISUpdates.clear();
}
bool Rematerializer::isMOIdenticalAtUses(MachineOperand &MO,
ArrayRef<SlotIndex> Uses) const {
if (Uses.empty())
return true;
Register Reg = MO.getReg();
unsigned SubIdx = MO.getSubReg();
LaneBitmask Mask = SubIdx ? TRI.getSubRegIndexLaneMask(SubIdx)
: MRI.getMaxLaneMaskForVReg(Reg);
const LiveInterval &LI = LIS.getInterval(Reg);
const VNInfo *DefVN =
LI.getVNInfoAt(LIS.getInstructionIndex(*MO.getParent()).getRegSlot(true));
for (SlotIndex Use : Uses) {
if (!isIdenticalAtUse(*DefVN, Mask, Use, LI))
return false;
}
return true;
}
RegisterIdx Rematerializer::findRematInRegion(RegisterIdx RegIdx,
unsigned Region,
SlotIndex Before) const {
auto It = Rematerializations.find(getOriginOrSelf(RegIdx));
if (It == Rematerializations.end())
return NoReg;
const RematsOf &Remats = It->getSecond();
SlotIndex BestSlot;
RegisterIdx BestRegIdx = NoReg;
for (RegisterIdx RematRegIdx : Remats) {
const Reg &RematReg = getReg(RematRegIdx);
if (RematReg.DefRegion != Region || RematReg.Uses.empty())
continue;
SlotIndex RematRegSlot =
LIS.getInstructionIndex(*RematReg.DefMI).getRegSlot();
if (RematRegSlot < Before &&
(BestRegIdx == NoReg || RematRegSlot > BestSlot)) {
BestSlot = RematRegSlot;
BestRegIdx = RematRegIdx;
}
}
return BestRegIdx;
}
void Rematerializer::deleteRegIfUnused(RegisterIdx RootIdx) {
if (!getReg(RootIdx).Uses.empty())
return;
// Traverse the root's dependency DAG depth-first to find the set of registers
// we can delete and a legal order to delete them in.
SmallVector<RegisterIdx, 4> DepDAG{RootIdx};
SmallSetVector<RegisterIdx, 8> DeleteOrder;
DeleteOrder.insert(RootIdx);
do {
// A deleted register's dependencies may be deletable too.
const Reg &DeleteReg = getReg(DepDAG.pop_back_val());
for (const Reg::Dependency &Dep : DeleteReg.Dependencies) {
// All dependencies loose a user (the deleted register).
Reg &DepReg = Regs[Dep.RegIdx];
DepReg.eraseUser(DeleteReg.DefMI, DeleteReg.DefRegion);
if (DepReg.Uses.empty()) {
DeleteOrder.insert(Dep.RegIdx);
DepDAG.push_back(Dep.RegIdx);
}
}
} while (!DepDAG.empty());
for (RegisterIdx RegIdx : reverse(DeleteOrder)) {
Reg &DeleteReg = Regs[RegIdx];
// It is possible that the defined register we are deleting doesn't have an
// interval yet if the LIS hasn't been updated since it was created.
Register DefReg = DeleteReg.getDefReg();
if (LIS.hasInterval(DefReg))
LIS.removeInterval(DefReg);
LISUpdates.erase(RegIdx);
deleteReg(RegIdx);
if (isRematerializedRegister(RegIdx)) {
// Delete rematerialized register from its origin's rematerializations.
const RegisterIdx OriginIdx = getOriginOf(RegIdx);
RematsOf &OriginRemats = Rematerializations.at(OriginIdx);
assert(OriginRemats.contains(RegIdx) && "broken remat<->origin link");
OriginRemats.erase(RegIdx);
if (OriginRemats.empty())
Rematerializations.erase(OriginIdx);
}
LLVM_DEBUG(dbgs() << "** Deleted " << printID(RegIdx) << "\n");
}
}
void Rematerializer::deleteReg(RegisterIdx RegIdx) {
noteRegDeleted(RegIdx);
Reg &DeleteReg = Regs[RegIdx];
assert(DeleteReg.DefMI && "register was already deleted");
// It is not possible for the deleted instruction to be the upper region
// boundary since we don't ever consider them rematerializable.
MachineBasicBlock::iterator &RegionBegin = Regions[DeleteReg.DefRegion].first;
if (RegionBegin == DeleteReg.DefMI)
RegionBegin = std::next(MachineBasicBlock::iterator(DeleteReg.DefMI));
LIS.RemoveMachineInstrFromMaps(*DeleteReg.DefMI);
DeleteReg.DefMI->eraseFromParent();
DeleteReg.DefMI = nullptr;
}
Rematerializer::Rematerializer(MachineFunction &MF,
SmallVectorImpl<RegionBoundaries> &Regions,
LiveIntervals &LIS)
: Regions(Regions), MRI(MF.getRegInfo()), LIS(LIS),
TII(*MF.getSubtarget().getInstrInfo()), TRI(TII.getRegisterInfo()) {
#ifdef EXPENSIVE_CHECKS
// Check that regions are valid.
DenseSet<MachineInstr *> SeenMIs;
for (const auto &[RegionBegin, RegionEnd] : Regions) {
assert(RegionBegin != RegionEnd && "empty region");
for (auto MI = RegionBegin; MI != RegionEnd; ++MI) {
bool IsNewMI = SeenMIs.insert(&*MI).second;
assert(IsNewMI && "overlapping regions");
assert(!MI->isTerminator() && "terminator in region");
}
if (RegionEnd != RegionBegin->getParent()->end()) {
bool IsNewMI = SeenMIs.insert(&*RegionEnd).second;
assert(IsNewMI && "overlapping regions (upper bound)");
}
}
#endif
}
bool Rematerializer::analyze() {
Regs.clear();
UnrematableOprds.clear();
Origins.clear();
Rematerializations.clear();
RegionMBB.clear();
RegToIdx.clear();
LISUpdates.clear();
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");
for (auto MI = Region.first; MI != Region.second; ++MI) {
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 != MBB.end()) {
MachineInstr *RegionTerm = &*Region.second;
assert(!MIRegion.contains(RegionTerm) && "regions should not intersect");
MIRegion.insert({RegionTerm, I});
}
}
const unsigned NumVirtRegs = MRI.getNumVirtRegs();
BitVector SeenRegs(NumVirtRegs);
for (unsigned I = 0, E = NumVirtRegs; I != E; ++I) {
if (!SeenRegs[I])
addRegIfRematerializable(I, MIRegion, SeenRegs);
}
assert(Regs.size() == UnrematableOprds.size());
LLVM_DEBUG({
for (RegisterIdx I = 0, E = getNumRegs(); I < E; ++I)
dbgs() << printDependencyDAG(I) << '\n';
});
return !Regs.empty();
}
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);
MachineOperand *MO = MRI.getOneDef(DefReg);
if (!MO)
return;
MachineInstr &DefMI = *MO->getParent();
if (!isMIRematerializable(DefMI))
return;
auto DefRegion = MIRegion.find(&DefMI);
if (DefRegion == MIRegion.end())
return;
Reg RematReg;
RematReg.DefMI = &DefMI;
RematReg.DefRegion = DefRegion->second;
unsigned SubIdx = DefMI.getOperand(0).getSubReg();
RematReg.Mask = SubIdx ? TRI.getSubRegIndexLaneMask(SubIdx)
: MRI.getMaxLaneMaskForVReg(DefReg);
// Collect the candidate's direct users, both rematerializable and
// unrematerializable. MIs outside provided regions cannot be tracked so the
// registers they use are not safely rematerializable.
for (MachineInstr &UseMI : MRI.use_nodbg_instructions(DefReg)) {
if (auto UseRegion = MIRegion.find(&UseMI); UseRegion != MIRegion.end())
RematReg.addUser(&UseMI, UseRegion->second);
else
return;
}
if (RematReg.Uses.empty())
return;
// Collect the candidate's dependencies. If the same register is used
// multiple times we just need to consider it once.
SmallDenseSet<Register, 4> AllDepRegs;
SmallVector<unsigned, 2> UnrematDeps;
for (const auto &[MOIdx, MO] : enumerate(RematReg.DefMI->operands())) {
Register DepReg = getRegDependency(MO);
if (!DepReg || !AllDepRegs.insert(DepReg).second)
continue;
unsigned DepRegIdx = DepReg.virtRegIndex();
if (!SeenRegs[DepRegIdx])
addRegIfRematerializable(DepRegIdx, MIRegion, SeenRegs);
if (auto DepIt = RegToIdx.find(DepReg); DepIt != RegToIdx.end())
RematReg.Dependencies.push_back(Reg::Dependency(MOIdx, DepIt->second));
else
UnrematDeps.push_back(MOIdx);
}
// The register is rematerializable.
RegToIdx.insert({DefReg, Regs.size()});
Regs.push_back(RematReg);
UnrematableOprds.push_back(UnrematDeps);
}
bool Rematerializer::isMIRematerializable(const MachineInstr &MI) const {
if (!TII.isReMaterializable(MI))
return false;
assert(MI.getOperand(0).getReg().isVirtual() && "should be virtual");
assert(MRI.hasOneDef(MI.getOperand(0).getReg()) && "should have single def");
for (const MachineOperand &MO : MI.all_uses()) {
// We can't remat physreg uses, unless it is a constant or an ignorable
// use (e.g. implicit exec use on VALU instructions)
if (MO.getReg().isPhysical()) {
if (MRI.isConstantPhysReg(MO.getReg()) || TII.isIgnorableUse(MO))
continue;
return false;
}
}
return true;
}
RegisterIdx Rematerializer::getDefRegIdx(const MachineInstr &MI) const {
if (!MI.getNumOperands() || !MI.getOperand(0).isReg() ||
MI.getOperand(0).readsReg())
return NoReg;
Register Reg = MI.getOperand(0).getReg();
auto UserRegIt = RegToIdx.find(Reg);
if (UserRegIt == RegToIdx.end())
return NoReg;
return UserRegIt->second;
}
RegisterIdx Rematerializer::rematerializeReg(
RegisterIdx RegIdx, unsigned UseRegion,
MachineBasicBlock::iterator InsertPos,
SmallVectorImpl<Reg::Dependency> &&Dependencies) {
RegisterIdx NewRegIdx = Regs.size();
Reg &NewReg = Regs.emplace_back();
Reg &FromReg = Regs[RegIdx];
NewReg.Mask = FromReg.Mask;
NewReg.DefRegion = UseRegion;
NewReg.Dependencies = std::move(Dependencies);
// Track rematerialization link between registers. Origins are always
// registers that existed originally, and rematerializations are always
// attached to them.
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(*RegionMBB[UseRegion], InsertPos, NewDefReg, 0,
*FromReg.DefMI);
NewReg.DefMI = &*std::prev(InsertPos);
RegToIdx.insert({NewDefReg, NewRegIdx});
postRematerialization(RegIdx, NewRegIdx, InsertPos);
noteRegCreated(NewRegIdx);
LLVM_DEBUG(dbgs() << "** Rematerialized " << printID(RegIdx) << " as "
<< printRematReg(NewRegIdx) << '\n');
return NewRegIdx;
}
void Rematerializer::recreateReg(
RegisterIdx RegIdx, unsigned DefRegion,
MachineBasicBlock::iterator InsertPos, Register DefReg,
SmallVectorImpl<Reg::Dependency> &&Dependencies) {
assert(RegToIdx.contains(DefReg) && "unknown defined register");
assert(RegToIdx.at(DefReg) == RegIdx && "incorrect defined register");
assert(!getReg(RegIdx).DefMI && "register is still alive");
Reg &OriginReg = Regs[RegIdx];
OriginReg.DefRegion = DefRegion;
OriginReg.Dependencies = std::move(Dependencies);
// Re-establish the link between origin and rematerialization if necessary.
const bool RecreateOriginalReg = isOriginalRegister(RegIdx);
if (!RecreateOriginalReg)
Rematerializations[getOriginOf(RegIdx)].insert(RegIdx);
// Rematerialize from one of the existing rematerializations or from the
// origin. We expect at least one to exist, otherwise it would mean the value
// held by the original register is no longer available anywhere in the MF.
RegisterIdx ModelRegIdx;
if (RecreateOriginalReg) {
assert(Rematerializations.contains(RegIdx) && "expected remats");
ModelRegIdx = *Rematerializations.at(RegIdx).begin();
} else {
assert(getReg(getOriginOf(RegIdx)).DefMI && "expected alive origin");
ModelRegIdx = getOriginOf(RegIdx);
}
const MachineInstr &ModelDefMI = *getReg(ModelRegIdx).DefMI;
TII.reMaterialize(*RegionMBB[DefRegion], InsertPos, DefReg, 0, ModelDefMI);
OriginReg.DefMI = &*std::prev(InsertPos);
postRematerialization(ModelRegIdx, RegIdx, InsertPos);
LLVM_DEBUG(dbgs() << "** Recreated " << printID(RegIdx) << " as "
<< printRematReg(RegIdx) << '\n');
}
void Rematerializer::postRematerialization(
RegisterIdx ModelRegIdx, RegisterIdx RematRegIdx,
MachineBasicBlock::iterator InsertPos) {
// The start of the new register's region may have changed.
Reg &ModelReg = Regs[ModelRegIdx], &RematReg = Regs[RematRegIdx];
LIS.InsertMachineInstrInMaps(*RematReg.DefMI);
MachineBasicBlock::iterator &RegionBegin = Regions[RematReg.DefRegion].first;
if (RegionBegin == std::next(MachineBasicBlock::iterator(RematReg.DefMI)))
RegionBegin = RematReg.DefMI;
// Replace dependencies as needed in the rematerialized MI. All dependencies
// of the latter gain a new user.
auto ZipedDeps = zip_equal(ModelReg.Dependencies, RematReg.Dependencies);
for (const auto &[OldDep, NewDep] : ZipedDeps) {
assert(OldDep.MOIdx == NewDep.MOIdx && "operand mismatch");
LLVM_DEBUG(dbgs() << " Operand #" << OldDep.MOIdx << ": "
<< printID(OldDep.RegIdx) << " -> "
<< printID(NewDep.RegIdx) << '\n');
Reg &NewDepReg = Regs[NewDep.RegIdx];
if (OldDep.RegIdx != NewDep.RegIdx) {
Register OldDefReg = ModelReg.DefMI->getOperand(OldDep.MOIdx).getReg();
RematReg.DefMI->substituteRegister(OldDefReg, NewDepReg.getDefReg(), 0,
TRI);
LISUpdates.insert(OldDep.RegIdx);
}
NewDepReg.addUser(RematReg.DefMI, RematReg.DefRegion);
LISUpdates.insert(NewDep.RegIdx);
}
}
std::pair<MachineInstr *, MachineInstr *>
Rematerializer::Reg::getRegionUseBounds(unsigned UseRegion,
const LiveIntervals &LIS) const {
auto It = Uses.find(UseRegion);
if (It == Uses.end())
return {nullptr, nullptr};
const RegionUsers &RegionUsers = It->getSecond();
assert(!RegionUsers.empty() && "empty userset in region");
auto User = RegionUsers.begin(), UserEnd = RegionUsers.end();
MachineInstr *FirstMI = *User, *LastMI = FirstMI;
SlotIndex FirstIndex = LIS.getInstructionIndex(*FirstMI),
LastIndex = FirstIndex;
while (++User != UserEnd) {
SlotIndex UserIndex = LIS.getInstructionIndex(**User);
if (UserIndex < FirstIndex) {
FirstIndex = UserIndex;
FirstMI = *User;
} else if (UserIndex > LastIndex) {
LastIndex = UserIndex;
LastMI = *User;
}
}
return {FirstMI, LastMI};
}
void Rematerializer::Reg::addUser(MachineInstr *MI, unsigned Region) {
Uses[Region].insert(MI);
}
void Rematerializer::Reg::addUsers(const RegionUsers &NewUsers,
unsigned Region) {
Uses[Region].insert_range(NewUsers);
}
void Rematerializer::Reg::eraseUser(MachineInstr *MI, unsigned Region) {
RegionUsers &RUsers = Uses.at(Region);
assert(RUsers.contains(MI) && "user not in region");
if (RUsers.size() == 1)
Uses.erase(Region);
else
RUsers.erase(MI);
}
Printable Rematerializer::printDependencyDAG(RegisterIdx RootIdx) const {
return Printable([&, RootIdx](raw_ostream &OS) {
DenseMap<RegisterIdx, unsigned> RegDepths;
std::function<void(RegisterIdx, unsigned)> WalkTree =
[&](RegisterIdx RegIdx, unsigned Depth) -> void {
unsigned MaxDepth = std::max(RegDepths.lookup_or(RegIdx, Depth), Depth);
RegDepths.emplace_or_assign(RegIdx, MaxDepth);
for (const Reg::Dependency &Dep : getReg(RegIdx).Dependencies)
WalkTree(Dep.RegIdx, Depth + 1);
};
WalkTree(RootIdx, 0);
// Sort in decreasing depth order to print root at the bottom.
SmallVector<std::pair<RegisterIdx, unsigned>> Regs(RegDepths.begin(),
RegDepths.end());
sort(Regs, [](const auto &LHS, const auto &RHS) {
return LHS.second > RHS.second;
});
OS << printID(RootIdx) << " has " << Regs.size() - 1 << " dependencies\n";
for (const auto &[RegIdx, Depth] : Regs) {
OS << indent(Depth, 2) << (Depth ? '|' : '*') << ' '
<< printRematReg(RegIdx, /*SkipRegions=*/Depth) << '\n';
}
OS << printRegUsers(RootIdx);
});
}
Printable Rematerializer::printID(RegisterIdx RegIdx) const {
return Printable([&, RegIdx](raw_ostream &OS) {
const Reg &PrintReg = getReg(RegIdx);
OS << '(' << RegIdx << '/';
if (!PrintReg.DefMI) {
OS << "<dead>";
} else {
OS << printReg(PrintReg.getDefReg(), &TRI,
PrintReg.DefMI->getOperand(0).getSubReg(), &MRI);
}
OS << ")[" << PrintReg.DefRegion << "]";
});
}
Printable Rematerializer::printRematReg(RegisterIdx RegIdx,
bool SkipRegions) const {
return Printable([&, RegIdx, SkipRegions](raw_ostream &OS) {
const Reg &PrintReg = getReg(RegIdx);
if (!SkipRegions) {
OS << printID(RegIdx) << " [" << PrintReg.DefRegion;
if (!PrintReg.Uses.empty()) {
assert(PrintReg.DefMI && "dead register cannot have uses");
const LiveInterval &LI = LIS.getInterval(PrintReg.getDefReg());
// First display all regions in which the register is live-through and
// not used.
bool First = true;
for (const auto [I, Bounds] : enumerate(Regions)) {
if (Bounds.first == Bounds.second)
continue;
if (!PrintReg.Uses.contains(I) &&
LI.liveAt(LIS.getInstructionIndex(*Bounds.first)) &&
LI.liveAt(LIS.getInstructionIndex(*std::prev(Bounds.second))
.getRegSlot())) {
OS << (First ? " - " : ",") << I;
First = false;
}
}
OS << (First ? " --> " : " -> ");
// Then display regions in which the register is used.
auto It = PrintReg.Uses.begin();
OS << It->first;
while (++It != PrintReg.Uses.end())
OS << "," << It->first;
}
OS << "] ";
}
OS << printID(RegIdx) << ' ';
PrintReg.DefMI->print(OS, /*IsStandalone=*/true, /*SkipOpers=*/false,
/*SkipDebugLoc=*/false, /*AddNewLine=*/false);
OS << " @ ";
LIS.getInstructionIndex(*PrintReg.DefMI).print(OS);
});
}
Printable Rematerializer::printRegUsers(RegisterIdx RegIdx) const {
return Printable([&, RegIdx](raw_ostream &OS) {
for (const auto &[UseRegion, Users] : getReg(RegIdx).Uses) {
for (MachineInstr *MI : Users)
OS << " User " << printUser(MI, UseRegion) << '\n';
}
});
}
Printable Rematerializer::printUser(const MachineInstr *MI,
std::optional<unsigned> UseRegion) const {
return Printable([&, MI, UseRegion](raw_ostream &OS) {
RegisterIdx RegIdx = getDefRegIdx(*MI);
if (RegIdx != NoReg) {
OS << printID(RegIdx);
} else {
OS << "(-/-)[";
if (UseRegion)
OS << *UseRegion;
else
OS << '?';
OS << ']';
}
OS << ' ';
MI->print(OS, /*IsStandalone=*/true, /*SkipOpers=*/false,
/*SkipDebugLoc=*/false, /*AddNewLine=*/false);
OS << " @ ";
LIS.getInstructionIndex(*MI).print(OS);
});
}
Rollbacker::RollbackInfo::RollbackInfo(const Rematerializer &Remater,
RegisterIdx RegIdx) {
const Rematerializer::Reg &Reg = Remater.getReg(RegIdx);
DefReg = Reg.getDefReg();
DefRegion = Reg.DefRegion;
Dependencies = Reg.Dependencies;
InsertPos = std::next(Reg.DefMI->getIterator());
if (InsertPos != Reg.DefMI->getParent()->end())
NextRegIdx = Remater.getDefRegIdx(*InsertPos);
else
NextRegIdx = Rematerializer::NoReg;
}
void Rollbacker::rematerializerNoteRegCreated(const Rematerializer &Remater,
RegisterIdx RegIdx) {
if (RollingBack)
return;
Rematerializations[Remater.getOriginOf(RegIdx)].insert(RegIdx);
}
void Rollbacker::rematerializerNoteRegDeleted(const Rematerializer &Remater,
RegisterIdx RegIdx) {
if (RollingBack || Remater.isRematerializedRegister(RegIdx))
return;
DeadRegs.try_emplace(RegIdx, Remater, RegIdx);
}
void Rollbacker::rollback(Rematerializer &Remater) {
RollingBack = true;
// Re-create deleted registers.
for (auto &[RegIdx, Info] : DeadRegs) {
assert(!Remater.getReg(RegIdx).isAlive() && "register should be dead");
// The MI that was originally just after the MI defining the register we
// are trying to re-create may have been deleted. In such cases, we can
// re-create at that MI's own insert position (and apply the same logic
// recursively).
MachineBasicBlock::iterator InsertPos = Info.InsertPos;
RegisterIdx NextRegIdx = Info.NextRegIdx;
while (NextRegIdx != Rematerializer::NoReg) {
const auto *NextRegRollback = DeadRegs.find(NextRegIdx);
if (NextRegRollback == DeadRegs.end())
break;
InsertPos = NextRegRollback->second.InsertPos;
NextRegIdx = NextRegRollback->second.NextRegIdx;
}
Remater.recreateReg(RegIdx, Info.DefRegion, InsertPos, Info.DefReg,
std::move(Info.Dependencies));
}
// Rollback rematerializations.
for (const auto &[RegIdx, RematsOf] : Rematerializations) {
for (RegisterIdx RematRegIdx : RematsOf) {
// It is possible that rematerializations were deleted. Their users would
// have been transfered to some other rematerialization so we can safely
// ignore them. Original registers that were deleted were just re-created
// so we do not need to check for that.
if (Remater.getReg(RematRegIdx).isAlive())
Remater.transferAllUsers(RematRegIdx, RegIdx);
}
}
Remater.updateLiveIntervals();
DeadRegs.clear();
Rematerializations.clear();
RollingBack = false;
}
Reference in New Issue View Git Blame Copy Permalink