[llvm] 7a20015 - [StackProtector] Rewrite dominator tree update handling
Roman Lebedev via llvm-commits
llvm-commits at lists.llvm.org
Sun Dec 11 17:53:36 PST 2022
Author: Roman Lebedev
Date: 2022-12-12T04:53:11+03:00
New Revision: 7a2001509b7f787b93ee5bcc41c149973feaaa74
URL: https://github.com/llvm/llvm-project/commit/7a2001509b7f787b93ee5bcc41c149973feaaa74
DIFF: https://github.com/llvm/llvm-project/commit/7a2001509b7f787b93ee5bcc41c149973feaaa74.diff
LOG: [StackProtector] Rewrite dominator tree update handling
Added:
Modified:
llvm/include/llvm/CodeGen/StackProtector.h
llvm/lib/CodeGen/StackProtector.cpp
llvm/test/CodeGen/X86/stack-protector-no-return.ll
Removed:
################################################################################
diff --git a/llvm/include/llvm/CodeGen/StackProtector.h b/llvm/include/llvm/CodeGen/StackProtector.h
index b96c0c74fabc0..7eefc756c4f14 100644
--- a/llvm/include/llvm/CodeGen/StackProtector.h
+++ b/llvm/include/llvm/CodeGen/StackProtector.h
@@ -18,6 +18,7 @@
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/IR/Instructions.h"
#include "llvm/Pass.h"
@@ -49,7 +50,7 @@ class StackProtector : public FunctionPass {
Function *F;
Module *M;
- DominatorTree *DT;
+ std::optional<DomTreeUpdater> DTU;
/// Layout - Mapping of allocations to the required SSPLayoutKind.
/// StackProtector analysis will update this map when determining if an
diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp
index 9a1063ed7f33b..c7731a7e50fe8 100644
--- a/llvm/lib/CodeGen/StackProtector.cpp
+++ b/llvm/lib/CodeGen/StackProtector.cpp
@@ -46,6 +46,7 @@
#include "llvm/Support/CommandLine.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include <optional>
#include <utility>
@@ -83,9 +84,8 @@ void StackProtector::getAnalysisUsage(AnalysisUsage &AU) const {
bool StackProtector::runOnFunction(Function &Fn) {
F = &Fn;
M = F->getParent();
- DominatorTreeWrapperPass *DTWP =
- getAnalysisIfAvailable<DominatorTreeWrapperPass>();
- DT = DTWP ? &DTWP->getDomTree() : nullptr;
+ 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();
@@ -109,7 +109,14 @@ bool StackProtector::runOnFunction(Function &Fn) {
}
++NumFunProtected;
- return InsertStackProtectors();
+ bool Changed = InsertStackProtectors();
+#ifdef EXPENSIVE_CHECKS
+ assert((!DTU ||
+ DTU->getDomTree().verify(DominatorTree::VerificationLevel::Full)) &&
+ "Failed to maintain validity of domtree!");
+#endif
+ DTU.reset();
+ return Changed;
}
/// \param [out] IsLarge is set to true if a protectable array is found and
@@ -442,7 +449,6 @@ bool StackProtector::InsertStackProtectors() {
TLI->useStackGuardXorFP() ||
(EnableSelectionDAGSP && !TM->Options.EnableFastISel);
AllocaInst *AI = nullptr; // Place on stack that stores the stack guard.
- bool RecalculateDT = false;
BasicBlock *FailBB = nullptr;
for (BasicBlock &BB : llvm::make_early_inc_range(*F)) {
@@ -532,8 +538,8 @@ bool StackProtector::InsertStackProtectors() {
// ...
// %1 = <stack guard>
// %2 = load StackGuardSlot
- // %3 = cmp i1 %1, %2
- // br i1 %3, label %SP_return, label %CallStackCheckFailBlk
+ // %3 = icmp ne i1 %1, %2
+ // br i1 %3, label %CallStackCheckFailBlk, label %SP_return
//
// SP_return:
// ret ...
@@ -548,62 +554,33 @@ bool StackProtector::InsertStackProtectors() {
if (!FailBB)
FailBB = CreateFailBB();
- // Split the basic block before the return instruction.
- BasicBlock *NewBB =
- BB.splitBasicBlock(CheckLoc->getIterator(), "SP_return");
-
- // Remove default branch instruction to the new BB.
- BB.getTerminator()->eraseFromParent();
-
- // Move the newly created basic block to the point right after the old
- // basic block so that it's in the "fall through" position.
- NewBB->moveAfter(&BB);
-
- // Generate the stack protector instructions in the old basic block.
- IRBuilder<> B(&BB);
+ IRBuilder<> B(CheckLoc);
Value *Guard = getStackGuard(TLI, M, B);
LoadInst *LI2 = B.CreateLoad(B.getInt8PtrTy(), AI, true);
- Value *Cmp = B.CreateICmpEQ(Guard, LI2);
+ auto *Cmp = cast<ICmpInst>(B.CreateICmpNE(Guard, LI2));
auto SuccessProb =
BranchProbabilityInfo::getBranchProbStackProtector(true);
auto FailureProb =
BranchProbabilityInfo::getBranchProbStackProtector(false);
MDNode *Weights = MDBuilder(F->getContext())
- .createBranchWeights(SuccessProb.getNumerator(),
- FailureProb.getNumerator());
- B.CreateCondBr(Cmp, NewBB, FailBB, Weights);
+ .createBranchWeights(FailureProb.getNumerator(),
+ SuccessProb.getNumerator());
- // Update the dominator tree if we need to.
- if (DT && DT->isReachableFromEntry(&BB))
- RecalculateDT = true;
+ SplitBlockAndInsertIfThen(Cmp, CheckLoc,
+ /*Unreachable=*/false, Weights,
+ DTU ? &*DTU : nullptr,
+ /*LI=*/nullptr, /*ThenBlock=*/FailBB);
+
+ auto *BI = cast<BranchInst>(Cmp->getParent()->getTerminator());
+ BasicBlock *NewBB = BI->getSuccessor(1);
+ NewBB->setName("SP_return");
+ NewBB->moveAfter(&BB);
+
+ Cmp->setPredicate(Cmp->getInversePredicate());
+ BI->swapSuccessors();
}
}
- // TODO: Refine me, use faster way to update DT.
- // Now we have spilt the BB, some like:
- // ===================================
- // BB:
- // RetOrNoReturnCall
- // ==>
- // BB:
- // CondBr
- // NewBB:
- // RetOrNoReturnCall
- // FailBB: (*)
- // HandleStackCheckFail
- // ===================================
- // The faster way should cover:
- // For NewBB, it should success the old BB's dominatees.
- // 1) return: it didn't have dominatee
- // 2) no-return call: there may has dominatees.
- //
- // For FailBB, it may be created before, So
- // 1) if it has 1 Predecessors, add it into DT.
- // 2) if it has 2 Predecessors, it should has no dominator, remove it from DT.
- // 3) if it has 3 or more Predecessors, DT has removed it, do nothing.
- if (RecalculateDT)
- DT->recalculate(*F);
-
// Return if we didn't modify any basic blocks. i.e., there are no return
// statements in the function.
return HasPrologue;
diff --git a/llvm/test/CodeGen/X86/stack-protector-no-return.ll b/llvm/test/CodeGen/X86/stack-protector-no-return.ll
index 8f391c17285a3..5d535471057c4 100644
--- a/llvm/test/CodeGen/X86/stack-protector-no-return.ll
+++ b/llvm/test/CodeGen/X86/stack-protector-no-return.ll
@@ -32,7 +32,7 @@ define void @_Z7catchesv() #0 personality i8* null {
; CHECK-NEXT: movq %fs:40, %rax
; CHECK-NEXT: cmpq (%rsp), %rax
; CHECK-NEXT: jne .LBB0_6
-; CHECK-NEXT: # %bb.5: # %SP_return2
+; CHECK-NEXT: # %bb.5: # %SP_return3
; CHECK-NEXT: popq %rax
; CHECK-NEXT: .cfi_def_cfa_offset 8
; CHECK-NEXT: retq
More information about the llvm-commits
mailing list