summaryrefslogtreecommitdiffstats
path: root/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
diff options
context:
space:
mode:
Diffstat (limited to 'llvm/lib/Target/AArch64/AArch64StackTagging.cpp')
-rw-r--r--llvm/lib/Target/AArch64/AArch64StackTagging.cpp112
1 files changed, 29 insertions, 83 deletions
diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index 8b653e23be8f..59037440c2b9 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -290,14 +290,6 @@ public:
};
class AArch64StackTagging : public FunctionPass {
- struct AllocaInfo {
- AllocaInst *AI;
- TrackingVH<Instruction> OldAI; // Track through RAUW to replace debug uses.
- SmallVector<IntrinsicInst *, 2> LifetimeStart;
- SmallVector<IntrinsicInst *, 2> LifetimeEnd;
- SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;
- };
-
const bool MergeInit;
const bool UseStackSafety;
@@ -313,7 +305,7 @@ public:
}
bool isInterestingAlloca(const AllocaInst &AI);
- void alignAndPadAlloca(AllocaInfo &Info);
+ void alignAndPadAlloca(memtag::AllocaInfo &Info);
void tagAlloca(AllocaInst *AI, Instruction *InsertBefore, Value *Ptr,
uint64_t Size);
@@ -322,9 +314,9 @@ public:
Instruction *collectInitializers(Instruction *StartInst, Value *StartPtr,
uint64_t Size, InitializerBuilder &IB);
- Instruction *
- insertBaseTaggedPointer(const MapVector<AllocaInst *, AllocaInfo> &Allocas,
- const DominatorTree *DT);
+ Instruction *insertBaseTaggedPointer(
+ const MapVector<AllocaInst *, memtag::AllocaInfo> &Allocas,
+ const DominatorTree *DT);
bool runOnFunction(Function &F) override;
StringRef getPassName() const override { return "AArch64 Stack Tagging"; }
@@ -466,12 +458,12 @@ void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
}
Instruction *AArch64StackTagging::insertBaseTaggedPointer(
- const MapVector<AllocaInst *, AllocaInfo> &InterestingAllocas,
+ const MapVector<AllocaInst *, memtag::AllocaInfo> &AllocasToInstrument,
const DominatorTree *DT) {
BasicBlock *PrologueBB = nullptr;
// Try sinking IRG as deep as possible to avoid hurting shrink wrap.
- for (auto &I : InterestingAllocas) {
- const AllocaInfo &Info = I.second;
+ for (auto &I : AllocasToInstrument) {
+ const memtag::AllocaInfo &Info = I.second;
AllocaInst *AI = Info.AI;
if (!PrologueBB) {
PrologueBB = AI->getParent();
@@ -490,7 +482,7 @@ Instruction *AArch64StackTagging::insertBaseTaggedPointer(
return Base;
}
-void AArch64StackTagging::alignAndPadAlloca(AllocaInfo &Info) {
+void AArch64StackTagging::alignAndPadAlloca(memtag::AllocaInfo &Info) {
const Align NewAlignment =
max(MaybeAlign(Info.AI->getAlign()), kTagGranuleSize);
Info.AI->setAlignment(NewAlignment);
@@ -536,63 +528,17 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
if (MergeInit)
AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
- MapVector<AllocaInst *, AllocaInfo>
- InterestingAllocas; // need stable iteration order
- SmallVector<Instruction *, 8> RetVec;
- SmallVector<Instruction *, 4> UnrecognizedLifetimes;
-
- bool CallsReturnTwice = false;
- for (Instruction &I : instructions(F)) {
- if (CallInst *CI = dyn_cast<CallInst>(&I)) {
- if (CI->canReturnTwice()) {
- CallsReturnTwice = true;
- }
- }
- if (auto *AI = dyn_cast<AllocaInst>(&I)) {
- if (isInterestingAlloca(*AI)) {
- InterestingAllocas[AI].AI = AI;
- InterestingAllocas[AI].OldAI = AI;
- }
- continue;
- }
-
- if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I)) {
- for (Value *V : DVI->location_ops())
- if (auto *AI = dyn_cast_or_null<AllocaInst>(V))
- if (isInterestingAlloca(*AI) &&
- (InterestingAllocas[AI].DbgVariableIntrinsics.empty() ||
- InterestingAllocas[AI].DbgVariableIntrinsics.back() != DVI))
- InterestingAllocas[AI].DbgVariableIntrinsics.push_back(DVI);
- continue;
- }
-
- auto *II = dyn_cast<IntrinsicInst>(&I);
- if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
- II->getIntrinsicID() == Intrinsic::lifetime_end)) {
- AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
- if (!AI) {
- UnrecognizedLifetimes.push_back(&I);
- continue;
- }
- if (!isInterestingAlloca(*AI))
- continue;
- if (II->getIntrinsicID() == Intrinsic::lifetime_start)
- InterestingAllocas[AI].LifetimeStart.push_back(II);
- else
- InterestingAllocas[AI].LifetimeEnd.push_back(II);
- continue;
- }
-
- Instruction *ExitUntag = getUntagLocationIfFunctionExit(I);
- if (ExitUntag)
- RetVec.push_back(ExitUntag);
- }
+ memtag::StackInfoBuilder SIB(
+ [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
+ for (Instruction &I : instructions(F))
+ SIB.visit(I);
+ memtag::StackInfo &SInfo = SIB.get();
- if (InterestingAllocas.empty())
+ if (SInfo.AllocasToInstrument.empty())
return false;
- for (auto &I : InterestingAllocas) {
- AllocaInfo &Info = I.second;
+ for (auto &I : SInfo.AllocasToInstrument) {
+ memtag::AllocaInfo &Info = I.second;
assert(Info.AI && isInterestingAlloca(*Info.AI));
alignAndPadAlloca(Info);
}
@@ -602,7 +548,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
DT = &P->getDomTree();
- if (DT == nullptr && (InterestingAllocas.size() > 1 ||
+ if (DT == nullptr && (SInfo.AllocasToInstrument.size() > 1 ||
!F->hasFnAttribute(Attribute::OptimizeNone))) {
DeleteDT = std::make_unique<DominatorTree>(*F);
DT = DeleteDT.get();
@@ -621,11 +567,11 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
SetTagFunc =
Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
- Instruction *Base = insertBaseTaggedPointer(InterestingAllocas, DT);
+ Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
int NextTag = 0;
- for (auto &I : InterestingAllocas) {
- const AllocaInfo &Info = I.second;
+ for (auto &I : SInfo.AllocasToInstrument) {
+ const memtag::AllocaInfo &Info = I.second;
AllocaInst *AI = Info.AI;
int Tag = NextTag;
NextTag = (NextTag + 1) % 16;
@@ -642,15 +588,15 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
TagPCall->setOperand(0, Info.AI);
bool StandardLifetime =
- UnrecognizedLifetimes.empty() &&
- isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,
- ClMaxLifetimes);
+ SInfo.UnrecognizedLifetimes.empty() &&
+ memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,
+ ClMaxLifetimes);
// Calls to functions that may return twice (e.g. setjmp) confuse the
// postdominator analysis, and will leave us to keep memory tagged after
// function return. Work around this by always untagging at every return
// statement if return_twice functions are called.
- if (UnrecognizedLifetimes.empty() && StandardLifetime &&
- !CallsReturnTwice) {
+ if (SInfo.UnrecognizedLifetimes.empty() && StandardLifetime &&
+ !SInfo.CallsReturnTwice) {
IntrinsicInst *Start = Info.LifetimeStart[0];
uint64_t Size =
cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
@@ -659,8 +605,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
if (!DT || !PDT ||
- !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec,
- TagEnd)) {
+ !memtag::forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd,
+ SInfo.RetVec, TagEnd)) {
for (auto *End : Info.LifetimeEnd)
End->eraseFromParent();
}
@@ -668,7 +614,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);
- for (auto &RI : RetVec) {
+ for (auto &RI : SInfo.RetVec) {
untagAlloca(AI, RI, Size);
}
// We may have inserted tag/untag outside of any lifetime interval.
@@ -686,7 +632,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
// If we have instrumented at least one alloca, all unrecognized lifetime
// instrinsics have to go.
- for (auto &I : UnrecognizedLifetimes)
+ for (auto &I : SInfo.UnrecognizedLifetimes)
I->eraseFromParent();
return true;