[llvm] [CodeGen] Port `StackProtector` to new pass manager (PR #75334)
via llvm-commits
llvm-commits at lists.llvm.org
Wed Dec 13 04:58:40 PST 2023
https://github.com/paperchalice updated https://github.com/llvm/llvm-project/pull/75334
>From 4b273ab8b5201a3635a44fcdbdeb9670f445c656 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 | 88 +++++++++----
llvm/lib/CodeGen/StackProtector.cpp | 123 +++++++++++++++---
llvm/lib/Passes/PassBuilder.cpp | 1 +
llvm/lib/Passes/PassRegistry.def | 2 +
6 files changed, 175 insertions(+), 45 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h b/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
index 3d8646291bb04c..7d543ec727afdb 100644
--- a/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
@@ -38,6 +38,7 @@
#include "llvm/CodeGen/SafeStack.h"
#include "llvm/CodeGen/SelectOptimize.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"
@@ -736,7 +737,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 ba81dbadf4bd60..3a85ec88db7d41 100644
--- a/llvm/include/llvm/CodeGen/MachinePassRegistry.def
+++ b/llvm/include/llvm/CodeGen/MachinePassRegistry.def
@@ -32,6 +32,7 @@ MODULE_PASS("jmc-instrumenter", JMCInstrumenterPass, ())
#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
@@ -62,6 +63,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, ())
@@ -133,7 +135,6 @@ DUMMY_FUNCTION_PASS("atomic-expand", AtomicExpandPass, ())
DUMMY_FUNCTION_PASS("codegenprepare", CodeGenPreparePass, ())
DUMMY_FUNCTION_PASS("gc-lowering", GCLoweringPass, ())
DUMMY_FUNCTION_PASS("shadow-stack-gc-lowering", ShadowStackGCLoweringPass, ())
-DUMMY_FUNCTION_PASS("stack-protector", StackProtectorPass, ())
#undef DUMMY_FUNCTION_PASS
#ifndef DUMMY_MODULE_PASS
diff --git a/llvm/include/llvm/CodeGen/StackProtector.h b/llvm/include/llvm/CodeGen/StackProtector.h
index 57cb7a1c85aea3..fda7f3e067e182 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,20 +31,71 @@ 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>;
+ using SSPLayoutMap =
+ DenseMap<const AllocaInst *, MachineFrameInfo::SSPLayoutKind>;
- const TargetMachine *TM = nullptr;
+ /// Layout - Mapping of allocations to the required SSPLayoutKind.
+ /// StackProtector analysis will update this map when determining if an
+ /// AllocaInst triggers a stack protector.
+ SSPLayoutMap Layout;
+
+ /// The minimum size of buffers that will receive stack smashing
+ /// protection when -fstack-protection is used.
+ unsigned SSPBufferSize = DefaultSSPBufferSize;
+
+ bool RequireStackProtector = false;
- /// TLI - Keep a pointer of a TargetLowering to consult for determining
- /// target type sizes.
- const TargetLoweringBase *TLI = nullptr;
- Triple Trip;
+ // A prologue is generated.
+ bool HasPrologue = false;
+
+ // IR checking code is generated.
+ bool HasIRCheck = false;
+
+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);
+
+ /// 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;
@@ -57,7 +109,7 @@ class StackProtector : public FunctionPass {
/// The minimum size of buffers that will receive stack smashing
/// protection when -fstack-protection is used.
- unsigned SSPBufferSize = DefaultSSPBufferSize;
+ unsigned SSPBufferSize = SSPLayoutInfo::DefaultSSPBufferSize;
// A prologue is generated.
bool HasPrologue = false;
@@ -65,18 +117,6 @@ class StackProtector : public FunctionPass {
// 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();
-
- /// CreateFailBB - Create a basic block to jump to when the stack protector
- /// check fails.
- BasicBlock *CreateFailBB();
-
public:
static char ID; // Pass identification, replacement for typeid.
@@ -93,8 +133,8 @@ class StackProtector : public FunctionPass {
/// 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);
};
} // end namespace llvm
diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 48dc7cb232e305..fe9a8945e278da 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,13 +174,11 @@ 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);
+ "stack-protector-buffer-size", SSPLayoutInfo::DefaultSSPBufferSize);
if (!requiresStackProtector(F, &Layout))
return false;
@@ -109,7 +191,8 @@ bool StackProtector::runOnFunction(Function &Fn) {
}
++NumFunProtected;
- bool Changed = InsertStackProtectors();
+ bool Changed = InsertStackProtectors(TM, F, DTU ? &*DTU : nullptr,
+ HasPrologue, HasIRCheck);
#ifdef EXPENSIVE_CHECKS
assert((!DTU ||
DTU->getDomTree().verify(DominatorTree::VerificationLevel::Full)) &&
@@ -284,7 +367,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 +379,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;
@@ -410,6 +494,10 @@ bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
return NeedsProtector;
}
+bool StackProtector::requiresStackProtector(Function *F, SSPLayoutMap *Layout) {
+ return SSPLayoutAnalysis::requiresStackProtector(F, Layout);
+}
+
/// Create a stack guard loading and populate whether SelectionDAG SSP is
/// supported.
static Value *getStackGuard(const TargetLoweringBase *TLI, Module *M,
@@ -460,13 +548,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 +661,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 +676,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 +694,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);
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 95b9fb7ad73500..a4dada78e0279b 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -86,6 +86,7 @@
#include "llvm/CodeGen/SafeStack.h"
#include "llvm/CodeGen/SelectOptimize.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 d8fc7cd8a231ff..1dc5406351fc6d 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -253,6 +253,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())
@@ -409,6 +410,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