[mte] [nfc] do not keep Tag in AllocaInfo

this is to get the implementation closer to HWASan

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D118960
This commit is contained in:
Florian Mayer 2022-02-03 17:10:28 -08:00
parent 034adaf5be
commit 02dbed4b7f

View File

@ -296,7 +296,6 @@ class AArch64StackTagging : public FunctionPass {
SmallVector<IntrinsicInst *, 2> LifetimeStart; SmallVector<IntrinsicInst *, 2> LifetimeStart;
SmallVector<IntrinsicInst *, 2> LifetimeEnd; SmallVector<IntrinsicInst *, 2> LifetimeEnd;
SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics; SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;
int Tag; // -1 for non-tagged allocations
}; };
const bool MergeInit; const bool MergeInit;
@ -467,15 +466,13 @@ void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
} }
Instruction *AArch64StackTagging::insertBaseTaggedPointer( Instruction *AArch64StackTagging::insertBaseTaggedPointer(
const MapVector<AllocaInst *, AllocaInfo> &Allocas, const MapVector<AllocaInst *, AllocaInfo> &InterestingAllocas,
const DominatorTree *DT) { const DominatorTree *DT) {
BasicBlock *PrologueBB = nullptr; BasicBlock *PrologueBB = nullptr;
// Try sinking IRG as deep as possible to avoid hurting shrink wrap. // Try sinking IRG as deep as possible to avoid hurting shrink wrap.
for (auto &I : Allocas) { for (auto &I : InterestingAllocas) {
const AllocaInfo &Info = I.second; const AllocaInfo &Info = I.second;
AllocaInst *AI = Info.AI; AllocaInst *AI = Info.AI;
if (Info.Tag < 0)
continue;
if (!PrologueBB) { if (!PrologueBB) {
PrologueBB = AI->getParent(); PrologueBB = AI->getParent();
continue; continue;
@ -539,7 +536,8 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
if (MergeInit) if (MergeInit)
AA = &getAnalysis<AAResultsWrapperPass>().getAAResults(); AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
MapVector<AllocaInst *, AllocaInfo> Allocas; // need stable iteration order MapVector<AllocaInst *, AllocaInfo>
InterestingAllocas; // need stable iteration order
SmallVector<Instruction *, 8> RetVec; SmallVector<Instruction *, 8> RetVec;
SmallVector<Instruction *, 4> UnrecognizedLifetimes; SmallVector<Instruction *, 4> UnrecognizedLifetimes;
@ -551,17 +549,20 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
} }
} }
if (auto *AI = dyn_cast<AllocaInst>(&I)) { if (auto *AI = dyn_cast<AllocaInst>(&I)) {
Allocas[AI].AI = AI; if (isInterestingAlloca(*AI)) {
Allocas[AI].OldAI = AI; InterestingAllocas[AI].AI = AI;
InterestingAllocas[AI].OldAI = AI;
}
continue; continue;
} }
if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I)) { if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I)) {
for (Value *V : DVI->location_ops()) for (Value *V : DVI->location_ops())
if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) if (auto *AI = dyn_cast_or_null<AllocaInst>(V))
if (Allocas[AI].DbgVariableIntrinsics.empty() || if (isInterestingAlloca(*AI) &&
Allocas[AI].DbgVariableIntrinsics.back() != DVI) (InterestingAllocas[AI].DbgVariableIntrinsics.empty() ||
Allocas[AI].DbgVariableIntrinsics.push_back(DVI); InterestingAllocas[AI].DbgVariableIntrinsics.back() != DVI))
InterestingAllocas[AI].DbgVariableIntrinsics.push_back(DVI);
continue; continue;
} }
@ -573,10 +574,13 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
UnrecognizedLifetimes.push_back(&I); UnrecognizedLifetimes.push_back(&I);
continue; continue;
} }
if (!isInterestingAlloca(*AI))
continue;
if (II->getIntrinsicID() == Intrinsic::lifetime_start) if (II->getIntrinsicID() == Intrinsic::lifetime_start)
Allocas[AI].LifetimeStart.push_back(II); InterestingAllocas[AI].LifetimeStart.push_back(II);
else else
Allocas[AI].LifetimeEnd.push_back(II); InterestingAllocas[AI].LifetimeEnd.push_back(II);
continue;
} }
Instruction *ExitUntag = getUntagLocationIfFunctionExit(I); Instruction *ExitUntag = getUntagLocationIfFunctionExit(I);
@ -584,35 +588,21 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
RetVec.push_back(ExitUntag); RetVec.push_back(ExitUntag);
} }
if (Allocas.empty()) if (InterestingAllocas.empty())
return false; return false;
int NextTag = 0; for (auto &I : InterestingAllocas) {
int NumInterestingAllocas = 0;
for (auto &I : Allocas) {
AllocaInfo &Info = I.second; AllocaInfo &Info = I.second;
assert(Info.AI); assert(Info.AI && isInterestingAlloca(*Info.AI));
if (!isInterestingAlloca(*Info.AI)) {
Info.Tag = -1;
continue;
}
alignAndPadAlloca(Info); alignAndPadAlloca(Info);
NumInterestingAllocas++;
Info.Tag = NextTag;
NextTag = (NextTag + 1) % 16;
} }
if (NumInterestingAllocas == 0)
return false;
std::unique_ptr<DominatorTree> DeleteDT; std::unique_ptr<DominatorTree> DeleteDT;
DominatorTree *DT = nullptr; DominatorTree *DT = nullptr;
if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>()) if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
DT = &P->getDomTree(); DT = &P->getDomTree();
if (DT == nullptr && (NumInterestingAllocas > 1 || if (DT == nullptr && (InterestingAllocas.size() > 1 ||
!F->hasFnAttribute(Attribute::OptimizeNone))) { !F->hasFnAttribute(Attribute::OptimizeNone))) {
DeleteDT = std::make_unique<DominatorTree>(*F); DeleteDT = std::make_unique<DominatorTree>(*F);
DT = DeleteDT.get(); DT = DeleteDT.get();
@ -631,21 +621,21 @@ bool AArch64StackTagging::runOnFunction(Function &Fn) {
SetTagFunc = SetTagFunc =
Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag); Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
Instruction *Base = insertBaseTaggedPointer(Allocas, DT); Instruction *Base = insertBaseTaggedPointer(InterestingAllocas, DT);
for (auto &I : Allocas) { int NextTag = 0;
for (auto &I : InterestingAllocas) {
const AllocaInfo &Info = I.second; const AllocaInfo &Info = I.second;
AllocaInst *AI = Info.AI; AllocaInst *AI = Info.AI;
if (Info.Tag < 0) int Tag = NextTag;
continue; NextTag = (NextTag + 1) % 16;
// Replace alloca with tagp(alloca). // Replace alloca with tagp(alloca).
IRBuilder<> IRB(Info.AI->getNextNode()); IRBuilder<> IRB(Info.AI->getNextNode());
Function *TagP = Intrinsic::getDeclaration( Function *TagP = Intrinsic::getDeclaration(
F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()}); F->getParent(), Intrinsic::aarch64_tagp, {Info.AI->getType()});
Instruction *TagPCall = Instruction *TagPCall =
IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base, IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base,
ConstantInt::get(IRB.getInt64Ty(), Info.Tag)}); ConstantInt::get(IRB.getInt64Ty(), Tag)});
if (Info.AI->hasName()) if (Info.AI->hasName())
TagPCall->setName(Info.AI->getName() + ".tag"); TagPCall->setName(Info.AI->getName() + ".tag");
Info.AI->replaceAllUsesWith(TagPCall); Info.AI->replaceAllUsesWith(TagPCall);