Skip to content

[CFI][stackprobe] Shrink wrapper select safe prologue insertion block when inline stack probing is enabled #81676

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.

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

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 5 additions & 0 deletions llvm/include/llvm/CodeGen/TargetLowering.h
Original file line number Diff line number Diff line change
Expand Up @@ -2010,6 +2010,11 @@ class TargetLoweringBase {

virtual bool hasInlineStackProbe(const MachineFunction &MF) const { return false; }

virtual bool
isStackProbeInstrDefinedFlagRegLiveIn(const MachineBasicBlock *MBB) const {
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggest renaming to something like stackProbingMayClobberLiveInRegs.

return false;
}

virtual StringRef getStackProbeSymbolName(const MachineFunction &MF) const {
return "";
}
Expand Down
9 changes: 9 additions & 0 deletions llvm/lib/CodeGen/ShrinkWrap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -957,6 +957,15 @@ bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
if (!ArePointsInteresting())
return Changed;

const TargetLowering *TLI = MF.getSubtarget().getTargetLowering();
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggest placing this whole logic in performShrinkWrapping and/or postShrinkWrapping.
The test TLI->isStackProbeInstrDefinedFlagRegLiveIn(Save) can quickly discard a block as a candidate for prologue/epilogue, so by placing it inside those functions we can avoid iterating over all instructions in a candidate block.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

In performShrinkWrapping we need to calculate StackAddressUsed for each block, regardless of whether the test TLI->isStackProbeInstrDefinedFlagRegLiveIn() returns true or false, so it seems that we have to iterate through instructions.

The current approach - in which the test is placed after all the analysis has been done - is relatively simple and safe - we don't have to worry about missing this test in some code paths in performShrinkWrapping or postShrinkWrapping.

Copy link
Collaborator

Choose a reason for hiding this comment

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

If my reading of performShrinkWrapping is correct, the StackAddressUsed is conservatively set for blocks that are not suitable for a prologue, not necessarily meaning they computed an SP-based address.
Thus we can set this variable for blocks with NZCV live-in.

Anyway, looking a bit more, I now think the proper place to fix this issue is TargetFrameLowering::canUseAsPrologue, so I hopefully fixed it there and I also found another bug:

#81878
#81879

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thus we can set this variable for blocks with NZCV live-in

A block with NZCV live-in wouldn't necessarily be selected for prologue insertion, so StackAddressUsed still need to be calculated by iterating through and examining each instruction.

while (TLI->hasInlineStackProbe(MF) &&
Copy link
Collaborator

Choose a reason for hiding this comment

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

TLI->hasInlineStackProbe(MF) is a loop invariant, no need to be in the condition.

TLI->isStackProbeInstrDefinedFlagRegLiveIn(Save)) {
Save = FindIDom<>(*Save, Save->predecessors(), *MDT);
Restore = FindIDom<>(*Restore, Restore->successors(), *MPDT);
if (!ArePointsInteresting())
return Changed;
}

LLVM_DEBUG(dbgs() << "Final shrink wrap candidates:\nSave: "
<< printMBBReference(*Save) << ' '
<< "\nRestore: " << printMBBReference(*Restore) << '\n');
Expand Down
5 changes: 5 additions & 0 deletions llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -27349,3 +27349,8 @@ bool AArch64TargetLowering::hasInlineStackProbe(
return !Subtarget->isTargetWindows() &&
MF.getInfo<AArch64FunctionInfo>()->hasStackProbing();
}

bool AArch64TargetLowering::isStackProbeInstrDefinedFlagRegLiveIn(
const MachineBasicBlock *MBB) const {
return MBB->isLiveIn(AArch64::NZCV);
}
3 changes: 3 additions & 0 deletions llvm/lib/Target/AArch64/AArch64ISelLowering.h
Original file line number Diff line number Diff line change
Expand Up @@ -997,6 +997,9 @@ class AArch64TargetLowering : public TargetLowering {
/// True if stack clash protection is enabled for this functions.
bool hasInlineStackProbe(const MachineFunction &MF) const override;

bool isStackProbeInstrDefinedFlagRegLiveIn(
const MachineBasicBlock *MBB) const override;

private:
/// Keep a pointer to the AArch64Subtarget around so that we can
/// make the right decision when generating code for different targets.
Expand Down
5 changes: 5 additions & 0 deletions llvm/lib/Target/X86/X86ISelLowering.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -57882,6 +57882,11 @@ X86TargetLowering::getStackProbeSize(const MachineFunction &MF) const {
4096);
}

bool X86TargetLowering::isStackProbeInstrDefinedFlagRegLiveIn(
const MachineBasicBlock *MBB) const {
return MBB->isLiveIn(X86::EFLAGS);
}

Align X86TargetLowering::getPrefLoopAlignment(MachineLoop *ML) const {
if (ML && ML->isInnermost() &&
ExperimentalPrefInnermostLoopAlignment.getNumOccurrences())
Expand Down
3 changes: 3 additions & 0 deletions llvm/lib/Target/X86/X86ISelLowering.h
Original file line number Diff line number Diff line change
Expand Up @@ -1550,6 +1550,9 @@ namespace llvm {
bool hasInlineStackProbe(const MachineFunction &MF) const override;
StringRef getStackProbeSymbolName(const MachineFunction &MF) const override;

bool isStackProbeInstrDefinedFlagRegLiveIn(
const MachineBasicBlock *MBB) const override;

unsigned getStackProbeSize(const MachineFunction &MF) const;

bool hasVectorBlend() const override { return true; }
Expand Down