[llvm] [CodeGen] Port `StackProtector` to new pass manager (PR #75334)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Jan 9 00:04:00 PST 2024
https://github.com/paperchalice updated https://github.com/llvm/llvm-project/pull/75334
>From af70d34e72b5045759c525a0aa29543d4c49f388 Mon Sep 17 00:00:00 2001
From: PaperChalice <29250197+paperchalice at users.noreply.github.com>
Date: Thu, 7 Dec 2023 16:20:03 +0800
Subject: [PATCH] [CodeGen] Port `StackProtector` to new pass manager
---
.../include/llvm/CodeGen/CodeGenPassBuilder.h | 3 +-
.../llvm/CodeGen/MachinePassRegistry.def | 3 +-
llvm/include/llvm/CodeGen/StackProtector.h | 93 +++++++----
llvm/lib/CodeGen/StackProtector.cpp | 154 ++++++++++++------
llvm/lib/Passes/PassBuilder.cpp | 1 +
llvm/lib/Passes/PassRegistry.def | 2 +
6 files changed, 176 insertions(+), 80 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h b/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
index 2100c30aad1180f..533d7774ba79c73 100644
--- a/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
@@ -40,6 +40,7 @@
#include "llvm/CodeGen/SelectOptimize.h"
#include "llvm/CodeGen/ShadowStackGCLowering.h"
#include "llvm/CodeGen/SjLjEHPrepare.h"
+#include "llvm/CodeGen/StackProtector.h"
#include "llvm/CodeGen/UnreachableBlockElim.h"
#include "llvm/CodeGen/WasmEHPrepare.h"
#include "llvm/CodeGen/WinEHPrepare.h"
@@ -748,7 +749,7 @@ void CodeGenPassBuilder<Derived>::addISelPrepare(AddIRPass &addPass) const {
// Add both the safe stack and the stack protection passes: each of them will
// only protect functions that have corresponding attributes.
addPass(SafeStackPass(&TM));
- addPass(StackProtectorPass());
+ addPass(StackProtectorPass(&TM));
if (Opt.PrintISelInput)
addPass(PrintFunctionPass(dbgs(),
diff --git a/llvm/include/llvm/CodeGen/MachinePassRegistry.def b/llvm/include/llvm/CodeGen/MachinePassRegistry.def
index b1b8ee8df29d167..2dc258c502e9c53 100644
--- a/llvm/include/llvm/CodeGen/MachinePassRegistry.def
+++ b/llvm/include/llvm/CodeGen/MachinePassRegistry.def
@@ -34,6 +34,7 @@ MODULE_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
#endif
FUNCTION_ANALYSIS("gc-function", GCFunctionAnalysis, ())
FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis, (PIC))
+FUNCTION_ANALYSIS("ssp-layout", SSPLayoutAnalysis, ())
FUNCTION_ANALYSIS("targetir", TargetIRAnalysis,
(std::move(TM.getTargetIRAnalysis())))
#undef FUNCTION_ANALYSIS
@@ -64,6 +65,7 @@ FUNCTION_PASS("safe-stack", SafeStackPass, (TM))
FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass, ())
FUNCTION_PASS("select-optimize", SelectOptimizePass, (TM))
FUNCTION_PASS("sjlj-eh-prepare", SjLjEHPreparePass, (TM))
+FUNCTION_PASS("stack-protector", StackProtectorPass, (TM))
FUNCTION_PASS("tlshoist", TLSVariableHoistPass, ())
FUNCTION_PASS("unreachableblockelim", UnreachableBlockElimPass, ())
FUNCTION_PASS("verify", VerifierPass, ())
@@ -134,7 +136,6 @@ MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis,
DUMMY_FUNCTION_PASS("atomic-expand", AtomicExpandPass, ())
DUMMY_FUNCTION_PASS("codegenprepare", CodeGenPreparePass, ())
DUMMY_FUNCTION_PASS("gc-lowering", GCLoweringPass, ())
-DUMMY_FUNCTION_PASS("stack-protector", StackProtectorPass, ())
#undef DUMMY_FUNCTION_PASS
#ifndef DUMMY_MACHINE_MODULE_PASS
diff --git a/llvm/include/llvm/CodeGen/StackProtector.h b/llvm/include/llvm/CodeGen/StackProtector.h
index 57cb7a1c85aea34..91d0b0d5e304bc8 100644
--- a/llvm/include/llvm/CodeGen/StackProtector.h
+++ b/llvm/include/llvm/CodeGen/StackProtector.h
@@ -19,6 +19,7 @@
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/IR/Instructions.h"
+#include "llvm/IR/PassManager.h"
#include "llvm/Pass.h"
#include "llvm/TargetParser/Triple.h"
@@ -30,25 +31,15 @@ class Module;
class TargetLoweringBase;
class TargetMachine;
-class StackProtector : public FunctionPass {
-private:
+class SSPLayoutInfo {
+ friend class StackProtectorPass;
+ friend class SSPLayoutAnalysis;
+ friend class StackProtector;
static constexpr unsigned DefaultSSPBufferSize = 8;
/// A mapping of AllocaInsts to their required SSP layout.
- using SSPLayoutMap = DenseMap<const AllocaInst *,
- MachineFrameInfo::SSPLayoutKind>;
-
- const TargetMachine *TM = nullptr;
-
- /// TLI - Keep a pointer of a TargetLowering to consult for determining
- /// target type sizes.
- const TargetLoweringBase *TLI = nullptr;
- Triple Trip;
-
- Function *F = nullptr;
- Module *M = nullptr;
-
- std::optional<DomTreeUpdater> DTU;
+ using SSPLayoutMap =
+ DenseMap<const AllocaInst *, MachineFrameInfo::SSPLayoutKind>;
/// Layout - Mapping of allocations to the required SSPLayoutKind.
/// StackProtector analysis will update this map when determining if an
@@ -59,23 +50,59 @@ class StackProtector : public FunctionPass {
/// protection when -fstack-protection is used.
unsigned SSPBufferSize = DefaultSSPBufferSize;
+ bool RequireStackProtector = false;
+
// A prologue is generated.
bool HasPrologue = false;
// IR checking code is generated.
bool HasIRCheck = false;
- /// InsertStackProtectors - Insert code into the prologue and epilogue of
- /// the function.
- ///
- /// - The prologue code loads and stores the stack guard onto the stack.
- /// - The epilogue checks the value stored in the prologue against the
- /// original value. It calls __stack_chk_fail if they differ.
- bool InsertStackProtectors();
+public:
+ // Return true if StackProtector is supposed to be handled by SelectionDAG.
+ bool shouldEmitSDCheck(const BasicBlock &BB) const;
+
+ void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
+};
+
+class SSPLayoutAnalysis : public AnalysisInfoMixin<SSPLayoutAnalysis> {
+ friend class AnalysisInfoMixin<SSPLayoutAnalysis>;
+ using SSPLayoutMap = SSPLayoutInfo::SSPLayoutMap;
+
+ static AnalysisKey Key;
+
+public:
+ using Result = SSPLayoutInfo;
+
+ Result run(Function &F, FunctionAnalysisManager &FAM);
- /// CreateFailBB - Create a basic block to jump to when the stack protector
- /// check fails.
- BasicBlock *CreateFailBB();
+ /// Check whether or not \p F needs a stack protector based upon the stack
+ /// protector level.
+ static bool requiresStackProtector(Function *F,
+ SSPLayoutMap *Layout = nullptr);
+};
+
+class StackProtectorPass : public PassInfoMixin<StackProtectorPass> {
+ const TargetMachine *TM;
+
+public:
+ explicit StackProtectorPass(const TargetMachine *TM) : TM(TM) {}
+ PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM);
+};
+
+class StackProtector : public FunctionPass {
+private:
+ /// A mapping of AllocaInsts to their required SSP layout.
+ using SSPLayoutMap = SSPLayoutInfo::SSPLayoutMap;
+
+ const TargetMachine *TM = nullptr;
+
+ Function *F = nullptr;
+ Module *M = nullptr;
+
+ std::optional<DomTreeUpdater> DTU;
+
+ SSPLayoutInfo LayoutInfo;
public:
static char ID; // Pass identification, replacement for typeid.
@@ -85,16 +112,22 @@ class StackProtector : public FunctionPass {
void getAnalysisUsage(AnalysisUsage &AU) const override;
// Return true if StackProtector is supposed to be handled by SelectionDAG.
- bool shouldEmitSDCheck(const BasicBlock &BB) const;
+ bool shouldEmitSDCheck(const BasicBlock &BB) const {
+ return LayoutInfo.shouldEmitSDCheck(BB);
+ }
bool runOnFunction(Function &Fn) override;
- void copyToMachineFrameInfo(MachineFrameInfo &MFI) const;
+ void copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
+ LayoutInfo.copyToMachineFrameInfo(MFI);
+ }
/// Check whether or not \p F needs a stack protector based upon the stack
/// protector level.
- static bool requiresStackProtector(Function *F, SSPLayoutMap *Layout = nullptr);
-
+ static bool requiresStackProtector(Function *F,
+ SSPLayoutMap *Layout = nullptr) {
+ return SSPLayoutAnalysis::requiresStackProtector(F, Layout);
+ }
};
} // end namespace llvm
diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 48dc7cb232e3051..fca822a485cafef 100644
--- a/llvm/lib/CodeGen/StackProtector.cpp
+++ b/llvm/lib/CodeGen/StackProtector.cpp
@@ -64,6 +64,90 @@ static cl::opt<bool> EnableSelectionDAGSP("enable-selectiondag-sp",
static cl::opt<bool> DisableCheckNoReturn("disable-check-noreturn-call",
cl::init(false), cl::Hidden);
+/// InsertStackProtectors - Insert code into the prologue and epilogue of the
+/// function.
+///
+/// - The prologue code loads and stores the stack guard onto the stack.
+/// - The epilogue checks the value stored in the prologue against the original
+/// value. It calls __stack_chk_fail if they differ.
+static bool InsertStackProtectors(const TargetMachine *TM, Function *F,
+ DomTreeUpdater *DTU, bool &HasPrologue,
+ bool &HasIRCheck);
+
+/// CreateFailBB - Create a basic block to jump to when the stack protector
+/// check fails.
+static BasicBlock *CreateFailBB(Function *F, const Triple &Trip);
+
+bool SSPLayoutInfo::shouldEmitSDCheck(const BasicBlock &BB) const {
+ return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
+}
+
+void SSPLayoutInfo::copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
+ if (Layout.empty())
+ return;
+
+ for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
+ if (MFI.isDeadObjectIndex(I))
+ continue;
+
+ const AllocaInst *AI = MFI.getObjectAllocation(I);
+ if (!AI)
+ continue;
+
+ SSPLayoutMap::const_iterator LI = Layout.find(AI);
+ if (LI == Layout.end())
+ continue;
+
+ MFI.setObjectSSPLayout(I, LI->second);
+ }
+}
+
+SSPLayoutInfo SSPLayoutAnalysis::run(Function &F,
+ FunctionAnalysisManager &FAM) {
+
+ SSPLayoutInfo Info;
+ Info.RequireStackProtector =
+ SSPLayoutAnalysis::requiresStackProtector(&F, &Info.Layout);
+ Info.SSPBufferSize = F.getFnAttributeAsParsedInteger(
+ "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
+ return Info;
+}
+
+AnalysisKey SSPLayoutAnalysis::Key;
+
+PreservedAnalyses StackProtectorPass::run(Function &F,
+ FunctionAnalysisManager &FAM) {
+ auto &Info = FAM.getResult<SSPLayoutAnalysis>(F);
+ auto *DT = FAM.getCachedResult<DominatorTreeAnalysis>(F);
+ DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
+
+ if (!Info.RequireStackProtector)
+ return PreservedAnalyses::all();
+
+ // TODO(etienneb): Functions with funclets are not correctly supported now.
+ // Do nothing if this is funclet-based personality.
+ if (F.hasPersonalityFn()) {
+ EHPersonality Personality = classifyEHPersonality(F.getPersonalityFn());
+ if (isFuncletEHPersonality(Personality))
+ return PreservedAnalyses::all();
+ }
+
+ ++NumFunProtected;
+ bool Changed = InsertStackProtectors(TM, &F, DT ? &DTU : nullptr,
+ Info.HasPrologue, Info.HasIRCheck);
+#ifdef EXPENSIVE_CHECKS
+ assert((!DT || DT->verify(DominatorTree::VerificationLevel::Full)) &&
+ "Failed to maintain validity of domtree!");
+#endif
+
+ if (!Changed)
+ return PreservedAnalyses::all();
+ PreservedAnalyses PA;
+ PA.preserve<SSPLayoutAnalysis>();
+ PA.preserve<DominatorTreeAnalysis>();
+ return PA;
+}
+
char StackProtector::ID = 0;
StackProtector::StackProtector() : FunctionPass(ID) {
@@ -90,14 +174,12 @@ bool StackProtector::runOnFunction(Function &Fn) {
if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
TM = &getAnalysis<TargetPassConfig>().getTM<TargetMachine>();
- Trip = TM->getTargetTriple();
- TLI = TM->getSubtargetImpl(Fn)->getTargetLowering();
- HasPrologue = false;
- HasIRCheck = false;
-
- SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
- "stack-protector-buffer-size", DefaultSSPBufferSize);
- if (!requiresStackProtector(F, &Layout))
+ LayoutInfo.HasPrologue = false;
+ LayoutInfo.HasIRCheck = false;
+
+ LayoutInfo.SSPBufferSize = Fn.getFnAttributeAsParsedInteger(
+ "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
+ if (!requiresStackProtector(F, &LayoutInfo.Layout))
return false;
// TODO(etienneb): Functions with funclets are not correctly supported now.
@@ -109,7 +191,9 @@ bool StackProtector::runOnFunction(Function &Fn) {
}
++NumFunProtected;
- bool Changed = InsertStackProtectors();
+ bool Changed =
+ InsertStackProtectors(TM, F, DTU ? &*DTU : nullptr,
+ LayoutInfo.HasPrologue, LayoutInfo.HasIRCheck);
#ifdef EXPENSIVE_CHECKS
assert((!DTU ||
DTU->getDomTree().verify(DominatorTree::VerificationLevel::Full)) &&
@@ -284,7 +368,8 @@ static const CallInst *findStackProtectorIntrinsic(Function &F) {
/// functions with aggregates that contain any buffer regardless of type and
/// size, and functions that contain stack-based variables that have had their
/// address taken.
-bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
+bool SSPLayoutAnalysis::requiresStackProtector(Function *F,
+ SSPLayoutMap *Layout) {
Module *M = F->getParent();
bool Strong = false;
bool NeedsProtector = false;
@@ -295,7 +380,7 @@ bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
SmallPtrSet<const PHINode *, 16> VisitedPHIs;
unsigned SSPBufferSize = F->getFnAttributeAsParsedInteger(
- "stack-protector-buffer-size", DefaultSSPBufferSize);
+ "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
if (F->hasFnAttribute(Attribute::SafeStack))
return false;
@@ -460,13 +545,12 @@ static bool CreatePrologue(Function *F, Module *M, Instruction *CheckLoc,
return SupportsSelectionDAGSP;
}
-/// InsertStackProtectors - Insert code into the prologue and epilogue of the
-/// function.
-///
-/// - The prologue code loads and stores the stack guard onto the stack.
-/// - The epilogue checks the value stored in the prologue against the original
-/// value. It calls __stack_chk_fail if they differ.
-bool StackProtector::InsertStackProtectors() {
+bool InsertStackProtectors(const TargetMachine *TM, Function *F,
+ DomTreeUpdater *DTU, bool &HasPrologue,
+ bool &HasIRCheck) {
+ auto *M = F->getParent();
+ auto *TLI = TM->getSubtargetImpl(*F)->getTargetLowering();
+
// If the target wants to XOR the frame pointer into the guard value, it's
// impossible to emit the check in IR, so the target *must* support stack
// protection in SDAG.
@@ -574,7 +658,7 @@ bool StackProtector::InsertStackProtectors() {
// merge pass will merge together all of the various BB into one including
// fail BB generated by the stack protector pseudo instruction.
if (!FailBB)
- FailBB = CreateFailBB();
+ FailBB = CreateFailBB(F, TM->getTargetTriple());
IRBuilder<> B(CheckLoc);
Value *Guard = getStackGuard(TLI, M, B);
@@ -589,8 +673,7 @@ bool StackProtector::InsertStackProtectors() {
SuccessProb.getNumerator());
SplitBlockAndInsertIfThen(Cmp, CheckLoc,
- /*Unreachable=*/false, Weights,
- DTU ? &*DTU : nullptr,
+ /*Unreachable=*/false, Weights, DTU,
/*LI=*/nullptr, /*ThenBlock=*/FailBB);
auto *BI = cast<BranchInst>(Cmp->getParent()->getTerminator());
@@ -608,9 +691,8 @@ bool StackProtector::InsertStackProtectors() {
return HasPrologue;
}
-/// CreateFailBB - Create a basic block to jump to when the stack protector
-/// check fails.
-BasicBlock *StackProtector::CreateFailBB() {
+BasicBlock *CreateFailBB(Function *F, const Triple &Trip) {
+ auto *M = F->getParent();
LLVMContext &Context = F->getContext();
BasicBlock *FailBB = BasicBlock::Create(Context, "CallStackCheckFailBlk", F);
IRBuilder<> B(FailBB);
@@ -633,27 +715,3 @@ BasicBlock *StackProtector::CreateFailBB() {
B.CreateUnreachable();
return FailBB;
}
-
-bool StackProtector::shouldEmitSDCheck(const BasicBlock &BB) const {
- return HasPrologue && !HasIRCheck && isa<ReturnInst>(BB.getTerminator());
-}
-
-void StackProtector::copyToMachineFrameInfo(MachineFrameInfo &MFI) const {
- if (Layout.empty())
- return;
-
- for (int I = 0, E = MFI.getObjectIndexEnd(); I != E; ++I) {
- if (MFI.isDeadObjectIndex(I))
- continue;
-
- const AllocaInst *AI = MFI.getObjectAllocation(I);
- if (!AI)
- continue;
-
- SSPLayoutMap::const_iterator LI = Layout.find(AI);
- if (LI == Layout.end())
- continue;
-
- MFI.setObjectSSPLayout(I, LI->second);
- }
-}
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 649451edc0e2c6c..98aea311901844e 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -88,6 +88,7 @@
#include "llvm/CodeGen/SelectOptimize.h"
#include "llvm/CodeGen/ShadowStackGCLowering.h"
#include "llvm/CodeGen/SjLjEHPrepare.h"
+#include "llvm/CodeGen/StackProtector.h"
#include "llvm/CodeGen/TypePromotion.h"
#include "llvm/CodeGen/WasmEHPrepare.h"
#include "llvm/CodeGen/WinEHPrepare.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index eaa7c3fc89241bc..658c30f2f05bddf 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -255,6 +255,7 @@ FUNCTION_ANALYSIS("should-not-run-function-passes",
ShouldNotRunFunctionPassesAnalysis())
FUNCTION_ANALYSIS("should-run-extra-vector-passes",
ShouldRunExtraVectorPasses())
+FUNCTION_ANALYSIS("ssp-layout", SSPLayoutAnalysis())
FUNCTION_ANALYSIS("stack-safety-local", StackSafetyAnalysis())
FUNCTION_ANALYSIS("targetir",
TM ? TM->getTargetIRAnalysis() : TargetIRAnalysis())
@@ -411,6 +412,7 @@ FUNCTION_PASS("sink", SinkingPass())
FUNCTION_PASS("sjlj-eh-prepare", SjLjEHPreparePass(TM))
FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
FUNCTION_PASS("slsr", StraightLineStrengthReducePass())
+FUNCTION_PASS("stack-protector", StackProtectorPass(TM))
FUNCTION_PASS("strip-gc-relocates", StripGCRelocates())
FUNCTION_PASS("structurizecfg", StructurizeCFGPass())
FUNCTION_PASS("tailcallelim", TailCallElimPass())
More information about the llvm-commits
mailing list