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