Files
Lucas Ramirez 64de40e5cb [CodeGen][AMDGPU] Prepare rematerializer for multi-def remat support (NFC) (#197579)
This makes some NFCs to the rematerializer before adding initial support
for rematerializing registers with multiple defs. The main change is
that, in the representation of (un)rematerializable register
dependencies, we drop references to machine operand indices which lose
meaning in the multi-def case.

Other minor changes listed below.

- Removal of `DefRegion` argument to `Rematerializer::recreateReg`.
Registers are always re-created in their original region so there is no
need to set their region again.
- Removal of `InsertPos` unused argument to
`Rematerializer::postRematerialization`.
- Refactor of how AMDGPU's scheduler checks whether a given register is
rematerializable.
2026-07-14 12:50:11 +02:00

903 lines
34 KiB
C++

//=====-- 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 "llvm/Support/ErrorHandling.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<RegisterIdx, 2> NewDeps;
// Copy all dependencies because recursive rematerialization of dependencies
// may invalidate references to the backing vector of registers.
SmallVector<RegisterIdx, 2> OldDeps(getReg(RootIdx).Dependencies);
for (RegisterIdx DepRegIdx : 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(DepRegIdx);
RegisterIdx NewDepRegIdx;
if (RematIdx == DRI.DependencyMap.end())
NewDepRegIdx = rematerializeToPos(DepRegIdx, UseRegion, InsertPos, DRI);
else
NewDepRegIdx = RematIdx->second;
NewDeps.push_back(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 (RegisterIdx &DepRegIdx : Regs[UserRegIdx].Dependencies) {
if (DepRegIdx == FromRegIdx) {
DepRegIdx = 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 (const auto &[UnrematReg, Mask] : getUnrematableDeps(RegIdx)) {
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 unrematerializable register "
<< printReg(UnrematReg, &TRI, 0, &MRI) << " with lanemask "
<< Mask << '\n');
}
}
LISUpdates.clear();
}
bool Rematerializer::isMOIdenticalAtUses(MachineOperand &MO,
ArrayRef<SlotIndex> Uses) const {
unsigned SubIdx = MO.getSubReg();
LaneBitmask Mask = SubIdx ? TRI.getSubRegIndexLaneMask(SubIdx)
: MRI.getMaxLaneMaskForVReg(MO.getReg());
return isRegIdenticalAtUses(
MO.getReg(), Mask,
LIS.getInstructionIndex(*MO.getParent()).getRegSlot(true), Uses);
}
bool Rematerializer::isRegIdenticalAtUses(Register Reg, LaneBitmask Mask,
SlotIndex RefSlot,
ArrayRef<SlotIndex> Uses) const {
if (Uses.empty())
return true;
const LiveInterval &LI = LIS.getInterval(Reg);
const VNInfo *DefVN = LI.getVNInfoAt(RefSlot);
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};
SmallVector<RegisterIdx, 8> DeleteOrder{RootIdx};
do {
// A deleted register's dependencies may be deletable too.
const Reg &DeleteReg = getReg(DepDAG.pop_back_val());
for (RegisterIdx DepRegIdx : DeleteReg.Dependencies) {
// All dependencies lose a user (the deleted register).
Reg &DepReg = Regs[DepRegIdx];
DepReg.eraseUser(DeleteReg.DefMI, DeleteReg.DefRegion);
if (DepReg.Uses.empty()) {
DeleteOrder.push_back(DepRegIdx);
DepDAG.push_back(DepRegIdx);
}
}
} while (!DepDAG.empty());
for (RegisterIdx RegIdx : 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) {
noteRegWillBeDeleted(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();
UnrematableDeps.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() == UnrematableDeps.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, rematerializable or not. If the same
// rematerializable register is used multiple times we just need to consider
// it once.
SmallSetVector<RegisterIdx, 2> RematDeps;
SmallMapVector<Register, LaneBitmask, 2> UnrematDeps;
for (const MachineOperand &MO : DefMI.all_uses()) {
Register DepReg = getRegDependency(MO);
if (!DepReg)
continue;
unsigned DepRegIdx = DepReg.virtRegIndex();
if (!SeenRegs[DepRegIdx])
addRegIfRematerializable(DepRegIdx, MIRegion, SeenRegs);
if (auto DepIt = RegToIdx.find(DepReg); DepIt != RegToIdx.end()) {
RematDeps.insert(DepIt->second);
} else {
LaneBitmask &CurrentMask =
UnrematDeps.try_emplace(DepReg, LaneBitmask::getNone()).first->second;
LaneBitmask Mask = MO.getSubReg()
? TRI.getSubRegIndexLaneMask(MO.getSubReg())
: MRI.getMaxLaneMaskForVReg(DepReg);
CurrentMask |= Mask;
}
}
// The register is rematerializable.
RematReg.Dependencies = RematDeps.takeVector();
RegToIdx.insert({DefReg, Regs.size()});
Regs.push_back(RematReg);
UnrematableDeps.push_back(UnrematDeps.takeVector());
}
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<RegisterIdx> &&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);
noteRegCreated(NewRegIdx);
LLVM_DEBUG(dbgs() << "** Rematerialized " << printID(RegIdx) << " as "
<< printRematReg(NewRegIdx) << '\n');
return NewRegIdx;
}
void Rematerializer::recreateReg(RegisterIdx RegIdx,
MachineBasicBlock::iterator InsertPos,
Register DefReg) {
assert(RegToIdx.contains(DefReg) && "unknown defined register");
assert(RegToIdx.at(DefReg) == RegIdx && "incorrect defined register");
assert(!getReg(RegIdx).isAlive() && "register is still alive");
Reg &OriginReg = Regs[RegIdx];
// 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)).isAlive() && "expected alive origin");
ModelRegIdx = getOriginOf(RegIdx);
}
const MachineInstr &ModelDefMI = *getReg(ModelRegIdx).DefMI;
TII.reMaterialize(*RegionMBB[OriginReg.DefRegion], InsertPos, DefReg, 0,
ModelDefMI);
OriginReg.DefMI = &*std::prev(InsertPos);
postRematerialization(ModelRegIdx, RegIdx);
LLVM_DEBUG(dbgs() << "** Recreated " << printID(RegIdx) << " as "
<< printRematReg(RegIdx) << '\n');
}
void Rematerializer::postRematerialization(RegisterIdx ModelRegIdx,
RegisterIdx RematRegIdx) {
// 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 &[OldDepRegIdx, NewDepRegIdx] : ZipedDeps) {
LLVM_DEBUG(dbgs() << " Dependency: " << printID(OldDepRegIdx) << " -> "
<< printID(NewDepRegIdx) << '\n');
Reg &NewDepReg = Regs[NewDepRegIdx];
if (OldDepRegIdx != NewDepRegIdx) {
Reg &OldDepReg = Regs[OldDepRegIdx];
RematReg.DefMI->substituteRegister(OldDepReg.getDefReg(),
NewDepReg.getDefReg(), 0, TRI);
LISUpdates.insert(OldDepRegIdx);
}
NewDepReg.addUser(RematReg.DefMI, RematReg.DefRegion);
LISUpdates.insert(NewDepRegIdx);
}
}
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 (RegisterIdx DepRegIdx : getReg(RegIdx).Dependencies)
WalkTree(DepRegIdx, 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.isAlive()) {
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);
OS << printID(RegIdx);
if (!SkipRegions) {
OS << " [" << PrintReg.DefRegion;
if (!PrintReg.Uses.empty()) {
assert(PrintReg.isAlive() && "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 (PrintReg.Uses.contains(I))
continue;
// The register must be live at the live-ins and live-outs of the
// region.
MachineBasicBlock::iterator LiveIn =
skipDebugInstructionsForward(Bounds.first, Bounds.second);
if (LiveIn == Bounds.second) {
// The region has no non-debug instructions, it's hard to assess
// whether the register is live across it without an index.
continue;
}
// LiveIn is inside the range and a non-debug instruction so we know
// this will also point to a non-debug instruction within the region.
MachineBasicBlock::iterator LiveOut = skipDebugInstructionsBackward(
std::prev(Bounds.second), Bounds.first);
if (LI.liveAt(LIS.getInstructionIndex(*LiveIn)) &&
LI.liveAt(LIS.getInstructionIndex(*LiveOut).getDeadSlot())) {
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 << "] ";
}
if (PrintReg.isAlive()) {
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);
});
}
void Rollbacker::rematerializerNoteRegCreated(const Rematerializer &Remater,
RegisterIdx RegIdx) {
if (RollingBack)
return;
assert(Remater.isRematerializedRegister(RegIdx) && "only remats are created");
Rematerializations[Remater.getOriginOf(RegIdx)].insert(RegIdx);
}
void Rollbacker::rematerializerNoteRegWillBeDeleted(
const Rematerializer &Remater, RegisterIdx RegIdx) {
if (RollingBack)
return;
// Find a valid re-creation position after the register's definition.
MachineInstr *DefMI = Remater.getReg(RegIdx).DefMI;
MachineBasicBlock *ParentMBB = DefMI->getParent();
MachineBasicBlock::iterator ValidPos = std::next(DefMI->getIterator());
while (ValidPos != ParentMBB->end() && isRollbackableMI(*ValidPos, Remater))
ValidPos = std::next(ValidPos);
if (Remater.isRematerializedRegister(RegIdx)) {
// Rematerializations will not be re-created. Previously deleted registers
// that reference this register's defining instruction as their re-creation
// position should instead be re-created at a valid position after the
// deleted MI.
invalidatePosition(DefMI, ValidPos);
return;
}
// Original registers can be re-created. Add a re-creation position for the
// definition of the rematerializable register.
DeadRegs.push_back(DeadReg(RegIdx, Remater));
const InsertBeforePos InsertPos = makePos(ValidPos, ParentMBB);
PosToIdx[InsertPos].insert(Positions.size());
Positions.push_back(InsertPos);
}
void Rollbacker::rollback(Rematerializer &Remater) {
RollingBack = true;
// As we re-create registers, map deleted definitions to re-created ones. This
// allows to replace invalid re-creation positions that reference deleted
// definitions to valid new positions while restoring original MI order.
DenseMap<MachineInstr *, MachineInstr *> Replacements;
unsigned PositionIndex = Positions.size();
// Re-create deleted registers in reverse order of deletion. Related registers
// are deleted in reverse def-use order so this ensures we re-create registers
// in def-use order. This also ensures that re-creation positions that became
// invalid due to later MI deletions can be corrected as we go.
for (const DeadReg &Reg : reverse(DeadRegs)) {
assert(!Remater.getReg(Reg.Idx).isAlive() && "register should be dead");
// Determine re-creation position for the register's definition.
MachineBasicBlock::iterator InsertPosition;
InsertBeforePos Pos = Positions[--PositionIndex];
if (auto *MBB = dyn_cast<MachineBasicBlock *>(Pos)) {
InsertPosition = MBB->end();
} else {
auto *MI = cast<MachineInstr *>(Pos);
InsertPosition = Replacements.lookup_or(MI, MI)->getIterator();
}
Remater.recreateReg(Reg.Idx, InsertPosition, Reg.DefReg);
const Rematerializer::Reg &RecreateReg = Remater.getReg(Reg.Idx);
if (!Replacements.insert({Reg.DefMI, RecreateReg.DefMI}).second)
llvm_unreachable("duplicate deleted MI");
}
// 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();
Positions.clear();
PosToIdx.clear();
Rematerializations.clear();
RollingBack = false;
}
bool Rollbacker::isRollbackableMI(const MachineInstr &MI,
const Rematerializer &Remater) const {
RegisterIdx RegIdx = Remater.getDefRegIdx(MI);
if (RegIdx == Rematerializer::NoReg ||
!Remater.isRematerializedRegister(RegIdx))
return false;
// It is possible that the MI defines a rematerializable register that was not
// recorded if the rollbacker was attached to the rematerializer after the
// rematerialization happened. In such cases the MI won't be rolled back.
auto RematsOf = Rematerializations.find(Remater.getOriginOf(RegIdx));
if (RematsOf == Rematerializations.end())
return false;
return RematsOf->getSecond().contains(RegIdx);
}
void Rollbacker::invalidatePosition(MachineInstr *MI,
MachineBasicBlock::iterator It) {
const InsertBeforePos MIPos = InsertBeforePos(MI),
NewPos = makePos(It, MI->getParent());
auto MIIndices = PosToIdx.find(MIPos);
if (MIIndices == PosToIdx.end())
return;
assert(!MIIndices->getSecond().empty() && "no index hold position");
for (unsigned I : MIIndices->getSecond())
Positions[I] = NewPos;
PosToIdx.try_emplace(NewPos).first->getSecond().insert_range(
MIIndices->getSecond());
PosToIdx.erase(MIPos);
}