[llvm] 6759cdd - [NFC] [MTE] Use helpers for stack tagging.
Florian Mayer via llvm-commits
llvm-commits at lists.llvm.org
Fri Feb 11 16:14:33 PST 2022
Author: Florian Mayer
Date: 2022-02-11T16:01:46-08:00
New Revision: 6759cdd82918e186ded1eda1a86b47e13469a811
URL: https://github.com/llvm/llvm-project/commit/6759cdd82918e186ded1eda1a86b47e13469a811
DIFF: https://github.com/llvm/llvm-project/commit/6759cdd82918e186ded1eda1a86b47e13469a811.diff
LOG: [NFC] [MTE] Use helpers for stack tagging.
Reviewed By: eugenis
Differential Revision: https://reviews.llvm.org/D119503
Added:
Modified:
llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
llvm/lib/Target/AArch64/AArch64StackTagging.cpp
llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
index d1548ecb58c06..b466393012778 100644
--- a/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
+++ b/llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h
@@ -19,6 +19,7 @@
#include "llvm/IR/Instruction.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Module.h"
+#include "llvm/IR/ValueHandle.h"
namespace llvm {
namespace memtag {
@@ -74,6 +75,7 @@ Instruction *getUntagLocationIfFunctionExit(Instruction &Inst);
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;
diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
index 0336a2f1d0ee0..59037440c2b99 100644
--- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
+++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp
@@ -290,14 +290,6 @@ class InitializerBuilder {
};
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 @@ class AArch64StackTagging : public FunctionPass {
}
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 @@ class AArch64StackTagging : public FunctionPass {
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 = memtag::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() &&
+ 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();
@@ -660,7 +606,7 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
if (!DT || !PDT ||
!memtag::forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd,
- RetVec, TagEnd)) {
+ 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;
diff --git a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
index bd4e4efd25aff..132f44f88534e 100644
--- a/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
+++ b/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp
@@ -65,8 +65,10 @@ void StackInfoBuilder::visit(Instruction &Inst) {
}
}
if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
- if (IsInterestingAlloca(*AI))
- Info.AllocasToInstrument.insert({AI, {}});
+ if (IsInterestingAlloca(*AI)) {
+ Info.AllocasToInstrument[AI].AI = AI;
+ Info.AllocasToInstrument[AI].OldAI = AI;
+ }
return;
}
auto *II = dyn_cast<IntrinsicInst>(&Inst);
More information about the llvm-commits
mailing list