Skip to content

[MIR] Support save/restore points with independent sets of registers #119358

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 23, 2025
Merged

[MIR] Support save/restore points with independent sets of registers #119358

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions llvm/include/llvm/CodeGen/MIRYamlMapping.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -634,19 +634,36 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::CalledGlobal)
namespace llvm {
namespace yaml {

// Struct representing one save/restore point in the 'savePoint'/'restorePoint'
// list
// Struct representing one save/restore point in the 'savePoint' /
// 'restorePoint' list. One point consists of machine basic block name and list
// of registers saved/restored in this basic block. In MIR it looks like:
// savePoint:
// - point: '%bb.1'
// registers:
// - '$rbx'
// - '$r12'
// ...
// restorePoint:
// - point: '%bb.1'
// registers:
// - '$rbx'
// - '$r12'
// If no register is saved/restored in the selected BB,
// field 'registers' is not specified.
struct SaveRestorePointEntry {
StringValue Point;
std::vector<StringValue> Registers;

bool operator==(const SaveRestorePointEntry &Other) const {
return Point == Other.Point;
return Point == Other.Point && Registers == Other.Registers;
}
};

template <> struct MappingTraits<SaveRestorePointEntry> {
static void mapping(IO &YamlIO, SaveRestorePointEntry &Entry) {
YamlIO.mapRequired("point", Entry.Point);
YamlIO.mapOptional("registers", Entry.Registers,
std::vector<StringValue>());
}
};

Expand Down
27 changes: 17 additions & 10 deletions llvm/include/llvm/CodeGen/MachineFrameInfo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,9 @@ class CalleeSavedInfo {
bool isSpilledToReg() const { return SpilledToReg; }
};

using SaveRestorePoints =
DenseMap<MachineBasicBlock *, std::vector<CalleeSavedInfo>>;

/// The MachineFrameInfo class represents an abstract stack frame until
/// prolog/epilog code is inserted. This class is key to allowing stack frame
/// representation optimizations, such as frame pointer elimination. It also
Expand Down Expand Up @@ -333,9 +336,9 @@ class MachineFrameInfo {
bool HasTailCall = false;

/// Not empty, if shrink-wrapping found a better place for the prologue.
SmallVector<MachineBasicBlock *, 4> SavePoints;
SaveRestorePoints SavePoints;
/// Not empty, if shrink-wrapping found a better place for the epilogue.
SmallVector<MachineBasicBlock *, 4> RestorePoints;
SaveRestorePoints RestorePoints;

/// Size of the UnsafeStack Frame
uint64_t UnsafeStackSize = 0;
Expand Down Expand Up @@ -825,17 +828,21 @@ class MachineFrameInfo {

void setCalleeSavedInfoValid(bool v) { CSIValid = v; }

ArrayRef<MachineBasicBlock *> getSavePoints() const { return SavePoints; }
void setSavePoints(ArrayRef<MachineBasicBlock *> NewSavePoints) {
SavePoints.assign(NewSavePoints.begin(), NewSavePoints.end());
}
ArrayRef<MachineBasicBlock *> getRestorePoints() const {
return RestorePoints;
const SaveRestorePoints &getRestorePoints() const { return RestorePoints; }

const SaveRestorePoints &getSavePoints() const { return SavePoints; }

void setSavePoints(SaveRestorePoints NewSavePoints) {
SavePoints = std::move(NewSavePoints);
}
void setRestorePoints(ArrayRef<MachineBasicBlock *> NewRestorePoints) {
RestorePoints.assign(NewRestorePoints.begin(), NewRestorePoints.end());

void setRestorePoints(SaveRestorePoints NewRestorePoints) {
RestorePoints = std::move(NewRestorePoints);
}

void clearSavePoints() { SavePoints.clear(); }
void clearRestorePoints() { RestorePoints.clear(); }

uint64_t getUnsafeStackSize() const { return UnsafeStackSize; }
void setUnsafeStackSize(uint64_t Size) { UnsafeStackSize = Size; }

Expand Down
20 changes: 14 additions & 6 deletions llvm/lib/CodeGen/MIRParser/MIRParser.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ class MIRParserImpl {
bool initializeSaveRestorePoints(
PerFunctionMIParsingState &PFS,
const std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
SmallVectorImpl<MachineBasicBlock *> &SaveRestorePoints);
llvm::SaveRestorePoints &SaveRestorePoints);

bool initializeCallSiteInfo(PerFunctionMIParsingState &PFS,
const yaml::MachineFunction &YamlMF);
Expand Down Expand Up @@ -872,11 +872,11 @@ bool MIRParserImpl::initializeFrameInfo(PerFunctionMIParsingState &PFS,
MFI.setHasTailCall(YamlMFI.HasTailCall);
MFI.setCalleeSavedInfoValid(YamlMFI.IsCalleeSavedInfoValid);
MFI.setLocalFrameSize(YamlMFI.LocalFrameSize);
SmallVector<MachineBasicBlock *, 4> SavePoints;
llvm::SaveRestorePoints SavePoints;
if (initializeSaveRestorePoints(PFS, YamlMFI.SavePoints, SavePoints))
return true;
MFI.setSavePoints(SavePoints);
SmallVector<MachineBasicBlock *, 4> RestorePoints;
llvm::SaveRestorePoints RestorePoints;
if (initializeSaveRestorePoints(PFS, YamlMFI.RestorePoints, RestorePoints))
return true;
MFI.setRestorePoints(RestorePoints);
Expand Down Expand Up @@ -1098,14 +1098,22 @@ bool MIRParserImpl::initializeConstantPool(PerFunctionMIParsingState &PFS,
bool MIRParserImpl::initializeSaveRestorePoints(
PerFunctionMIParsingState &PFS,
const std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
SmallVectorImpl<MachineBasicBlock *> &SaveRestorePoints) {
llvm::SaveRestorePoints &SaveRestorePoints) {
SMDiagnostic Error;
MachineBasicBlock *MBB = nullptr;
for (const yaml::SaveRestorePointEntry &Entry : YamlSRPoints) {
if (parseMBBReference(PFS, MBB, Entry.Point.Value))
return true;
SaveRestorePoints.push_back(MBB);
}

std::vector<CalleeSavedInfo> Registers;
for (auto &RegStr : Entry.Registers) {
Register Reg;
if (parseNamedRegisterReference(PFS, Reg, RegStr.Value, Error))
return error(Error, RegStr.SourceRange);
Registers.push_back(CalleeSavedInfo(Reg));
}
SaveRestorePoints.try_emplace(MBB, std::move(Registers));
}
return false;
}

Expand Down
39 changes: 31 additions & 8 deletions llvm/lib/CodeGen/MIRPrinter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,11 +149,13 @@ static void convertMCP(yaml::MachineFunction &MF,
static void convertMJTI(ModuleSlotTracker &MST, yaml::MachineJumpTable &YamlJTI,
const MachineJumpTableInfo &JTI);
static void convertMFI(ModuleSlotTracker &MST, yaml::MachineFrameInfo &YamlMFI,
const MachineFrameInfo &MFI);
const MachineFrameInfo &MFI,
const TargetRegisterInfo *TRI);
static void
convertSRPoints(ModuleSlotTracker &MST,
std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
ArrayRef<MachineBasicBlock *> SaveRestorePoints);
const llvm::SaveRestorePoints &SRPoints,
const TargetRegisterInfo *TRI);
static void convertStackObjects(yaml::MachineFunction &YMF,
const MachineFunction &MF,
ModuleSlotTracker &MST, MFPrintState &State);
Expand Down Expand Up @@ -204,7 +206,8 @@ static void printMF(raw_ostream &OS, const MachineModuleInfo &MMI,
convertMRI(YamlMF, MF, MF.getRegInfo(), MF.getSubtarget().getRegisterInfo());
MachineModuleSlotTracker &MST = State.MST;
MST.incorporateFunction(MF.getFunction());
convertMFI(MST, YamlMF.FrameInfo, MF.getFrameInfo());
convertMFI(MST, YamlMF.FrameInfo, MF.getFrameInfo(),
MF.getSubtarget().getRegisterInfo());
convertStackObjects(YamlMF, MF, MST, State);
convertEntryValueObjects(YamlMF, MF, MST);
convertCallSiteObjects(YamlMF, MF, MST);
Expand Down Expand Up @@ -339,7 +342,8 @@ static void convertMRI(yaml::MachineFunction &YamlMF, const MachineFunction &MF,
}

static void convertMFI(ModuleSlotTracker &MST, yaml::MachineFrameInfo &YamlMFI,
const MachineFrameInfo &MFI) {
const MachineFrameInfo &MFI,
const TargetRegisterInfo *TRI) {
YamlMFI.IsFrameAddressTaken = MFI.isFrameAddressTaken();
YamlMFI.IsReturnAddressTaken = MFI.isReturnAddressTaken();
YamlMFI.HasStackMap = MFI.hasStackMap();
Expand All @@ -360,9 +364,9 @@ static void convertMFI(ModuleSlotTracker &MST, yaml::MachineFrameInfo &YamlMFI,
YamlMFI.IsCalleeSavedInfoValid = MFI.isCalleeSavedInfoValid();
YamlMFI.LocalFrameSize = MFI.getLocalFrameSize();
if (!MFI.getSavePoints().empty())
convertSRPoints(MST, YamlMFI.SavePoints, MFI.getSavePoints());
convertSRPoints(MST, YamlMFI.SavePoints, MFI.getSavePoints(), TRI);
if (!MFI.getRestorePoints().empty())
convertSRPoints(MST, YamlMFI.RestorePoints, MFI.getRestorePoints());
convertSRPoints(MST, YamlMFI.RestorePoints, MFI.getRestorePoints(), TRI);
}

static void convertEntryValueObjects(yaml::MachineFunction &YMF,
Expand Down Expand Up @@ -619,16 +623,35 @@ static void convertMCP(yaml::MachineFunction &MF,
static void
convertSRPoints(ModuleSlotTracker &MST,
std::vector<yaml::SaveRestorePointEntry> &YamlSRPoints,
ArrayRef<MachineBasicBlock *> SRPoints) {
for (const auto &MBB : SRPoints) {
const llvm::SaveRestorePoints &SRPoints,
const TargetRegisterInfo *TRI) {
for (const auto &[MBB, CSInfos] : SRPoints) {
SmallString<16> Str;
yaml::SaveRestorePointEntry Entry;
raw_svector_ostream StrOS(Str);
StrOS << printMBBReference(*MBB);
Entry.Point = StrOS.str().str();
Str.clear();
for (const CalleeSavedInfo &Info : CSInfos) {
if (Info.getReg()) {
StrOS << printReg(Info.getReg(), TRI);
Entry.Registers.push_back(StrOS.str().str());
Str.clear();
}
}
// Sort here needed for stable output for lit tests
std::sort(Entry.Registers.begin(), Entry.Registers.end(),
[](const yaml::StringValue &Lhs, const yaml::StringValue &Rhs) {
return Lhs.Value < Rhs.Value;
});
YamlSRPoints.push_back(std::move(Entry));
}
// Sort here needed for stable output for lit tests
std::sort(YamlSRPoints.begin(), YamlSRPoints.end(),
[](const yaml::SaveRestorePointEntry &Lhs,
const yaml::SaveRestorePointEntry &Rhs) {
return Lhs.Point.Value < Rhs.Point.Value;
});
}

static void convertMJTI(ModuleSlotTracker &MST, yaml::MachineJumpTable &YamlJTI,
Expand Down
4 changes: 2 additions & 2 deletions llvm/lib/CodeGen/MachineFrameInfo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,14 +250,14 @@ void MachineFrameInfo::print(const MachineFunction &MF, raw_ostream &OS) const{
OS << "save points:\n";

for (auto &item : SavePoints)
OS << printMBBReference(*item) << "\n";
OS << printMBBReference(*item.first) << "\n";
} else
OS << "save points are empty\n";

if (!RestorePoints.empty()) {
OS << "restore points:\n";
for (auto &item : RestorePoints)
OS << printMBBReference(*item) << "\n";
OS << printMBBReference(*item.first) << "\n";
} else
OS << "restore points are empty\n";
}
Expand Down
34 changes: 26 additions & 8 deletions llvm/lib/CodeGen/PrologEpilogInserter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,8 +351,8 @@ bool PEIImpl::run(MachineFunction &MF) {
delete RS;
SaveBlocks.clear();
RestoreBlocks.clear();
MFI.setSavePoints({});
MFI.setRestorePoints({});
MFI.clearSavePoints();
MFI.clearRestorePoints();
return true;
}

Expand Down Expand Up @@ -431,10 +431,12 @@ void PEIImpl::calculateSaveRestoreBlocks(MachineFunction &MF) {
if (!MFI.getSavePoints().empty()) {
assert(MFI.getSavePoints().size() == 1 &&
"Multiple save points are not yet supported!");
SaveBlocks.push_back(MFI.getSavePoints().front());
const auto &SavePoint = *MFI.getSavePoints().begin();
SaveBlocks.push_back(SavePoint.first);
assert(MFI.getRestorePoints().size() == 1 &&
"Multiple restore points are not yet supported!");
MachineBasicBlock *RestoreBlock = MFI.getRestorePoints().front();
const auto &RestorePoint = *MFI.getRestorePoints().begin();
MachineBasicBlock *RestoreBlock = RestorePoint.first;
// If RestoreBlock does not have any successor and is not a return block
// then the end point is unreachable and we do not need to insert any
// epilogue.
Expand Down Expand Up @@ -563,8 +565,9 @@ static void updateLiveness(MachineFunction &MF) {

assert(MFI.getSavePoints().size() < 2 &&
"Multiple save points not yet supported!");
MachineBasicBlock *Save =
MFI.getSavePoints().empty() ? nullptr : MFI.getSavePoints().front();
MachineBasicBlock *Save = MFI.getSavePoints().empty()
? nullptr
: (*MFI.getSavePoints().begin()).first;

if (!Save)
Save = Entry;
Expand All @@ -577,8 +580,9 @@ static void updateLiveness(MachineFunction &MF) {

assert(MFI.getRestorePoints().size() < 2 &&
"Multiple restore points not yet supported!");
MachineBasicBlock *Restore =
MFI.getRestorePoints().empty() ? nullptr : MFI.getRestorePoints().front();
MachineBasicBlock *Restore = MFI.getRestorePoints().empty()
? nullptr
: (*MFI.getRestorePoints().begin()).first;
if (Restore)
// By construction Restore cannot be visited, otherwise it
// means there exists a path to Restore that does not go
Expand Down Expand Up @@ -687,6 +691,20 @@ void PEIImpl::spillCalleeSavedRegs(MachineFunction &MF) {
MFI.setCalleeSavedInfoValid(true);

std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();

// Fill SavePoints and RestorePoints with CalleeSavedRegisters
if (!MFI.getSavePoints().empty()) {
SaveRestorePoints SaveRestorePts;
for (const auto &SavePoint : MFI.getSavePoints())
SaveRestorePts.insert({SavePoint.first, CSI});
MFI.setSavePoints(std::move(SaveRestorePts));

SaveRestorePts.clear();
for (const auto &RestorePoint : MFI.getRestorePoints())
SaveRestorePts.insert({RestorePoint.first, CSI});
MFI.setRestorePoints(std::move(SaveRestorePts));
}

if (!CSI.empty()) {
if (!MFI.hasCalls())
NumLeafFuncWithSpills++;
Expand Down
12 changes: 6 additions & 6 deletions llvm/lib/CodeGen/ShrinkWrap.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -967,12 +967,12 @@ bool ShrinkWrapImpl::run(MachineFunction &MF) {
<< "\nRestore: " << printMBBReference(*Restore) << '\n');

MachineFrameInfo &MFI = MF.getFrameInfo();
SmallVector<MachineBasicBlock *, 4> SavePoints;
SmallVector<MachineBasicBlock *, 4> RestorePoints;
if (Save) {
SavePoints.push_back(Save);
RestorePoints.push_back(Restore);
}

// List of CalleeSavedInfo for registers will be added during prologepilog
// pass
SaveRestorePoints SavePoints({{Save, {}}});
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to confirm, is the invariant here that Save was always non-null along this path? If so, the change looks okay, otherwise you're adding a null block.

(To be clear, my LGTM is conditional on the answer to this question being always non-null.)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, Save is always not-null here, otherwise we we had to return under this condition (or earlier):

bool ArePointsInteresting() const { return Save != Entry && Save && Restore; }
...
 if (!ArePointsInteresting())
    return Changed;

SaveRestorePoints RestorePoints({{Restore, {}}});

MFI.setSavePoints(SavePoints);
MFI.setRestorePoints(RestorePoints);
++NumCandidates;
Expand Down
6 changes: 4 additions & 2 deletions llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,10 +213,12 @@ void SILowerSGPRSpills::calculateSaveRestoreBlocks(MachineFunction &MF) {
if (!MFI.getSavePoints().empty()) {
assert(MFI.getSavePoints().size() == 1 &&
"Multiple save points not yet supported!");
SaveBlocks.push_back(MFI.getSavePoints().front());
const auto &SavePoint = *MFI.getSavePoints().begin();
SaveBlocks.push_back(SavePoint.first);
assert(MFI.getRestorePoints().size() == 1 &&
"Multiple restore points not yet supported!");
MachineBasicBlock *RestoreBlock = MFI.getRestorePoints().front();
const auto &RestorePoint = *MFI.getRestorePoints().begin();
MachineBasicBlock *RestoreBlock = RestorePoint.first;
// If RestoreBlock does not have any successor and is not a return block
// then the end point is unreachable and we do not need to insert any
// epilogue.
Expand Down
3 changes: 1 addition & 2 deletions llvm/lib/Target/PowerPC/PPCFrameLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2081,9 +2081,8 @@ void PPCFrameLowering::processFunctionBeforeFrameFinalized(MachineFunction &MF,
if (!MFI.getSavePoints().empty() && MFI.hasTailCall()) {
assert(MFI.getRestorePoints().size() < 2 &&
"MFI can't contain multiple restore points!");
MachineBasicBlock *RestoreBlock = MFI.getRestorePoints().front();
for (MachineBasicBlock &MBB : MF) {
if (MBB.isReturnBlock() && (&MBB) != RestoreBlock)
if (MBB.isReturnBlock() && (!MFI.getRestorePoints().contains(&MBB)))
createTailCallBranchInstr(MBB);
}
}
Expand Down
4 changes: 4 additions & 0 deletions llvm/test/CodeGen/MIR/X86/frame-info-multiple-save-restore-points-parse.mir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,14 @@ liveins:
# CHECK: frameInfo:
# CHECK: savePoint:
# CHECK-NEXT: - point: '%bb.1'
# CHECK-NEXT: registers: []
# CHECK-NEXT: - point: '%bb.2'
# CHECK-NEXT: registers: []
# CHECK: restorePoint:
# CHECK-NEXT: - point: '%bb.2'
# CHECK-NEXT: registers: []
# CHECK-NEXT: - point: '%bb.3'
# CHECK-NEXT: registers: []
# CHECK: stack
frameInfo:
maxAlignment: 4
Expand Down
Loading