diff --git a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp --- a/llvm/lib/Target/AArch64/AArch64StackTagging.cpp +++ b/llvm/lib/Target/AArch64/AArch64StackTagging.cpp @@ -296,7 +296,6 @@ SmallVector LifetimeStart; SmallVector LifetimeEnd; SmallVector DbgVariableIntrinsics; - int Tag; // -1 for non-tagged allocations }; const bool MergeInit; @@ -467,15 +466,13 @@ } Instruction *AArch64StackTagging::insertBaseTaggedPointer( - const MapVector &Allocas, + const MapVector &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,11 +536,12 @@ if (MergeInit) AA = &getAnalysis().getAAResults(); - MapVector Allocas; // need stable iteration order + MapVector InterestingAllocas; // need stable iteration order SmallVector RetVec; SmallVector UnrecognizedLifetimes; bool CallsReturnTwice = false; + bool HasUninterestingAllocas = false; for (Instruction &I : instructions(F)) { if (CallInst *CI = dyn_cast(&I)) { if (CI->canReturnTwice()) { @@ -551,17 +549,22 @@ } } if (auto *AI = dyn_cast(&I)) { - Allocas[AI].AI = AI; - Allocas[AI].OldAI = AI; + if (isInterestingAlloca(*AI)) { + InterestingAllocas[AI].AI = AI; + InterestingAllocas[AI].OldAI = AI; + } else { + HasUninterestingAllocas = true; + } continue; } if (auto *DVI = dyn_cast(&I)) { for (Value *V : DVI->location_ops()) if (auto *AI = dyn_cast_or_null(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 +576,13 @@ 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 +590,21 @@ RetVec.push_back(ExitUntag); } - if (Allocas.empty()) - return false; + if (InterestingAllocas.empty()) + return HasUninterestingAllocas; - 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 true; - std::unique_ptr DeleteDT; DominatorTree *DT = nullptr; if (auto *P = getAnalysisIfAvailable()) DT = &P->getDomTree(); - if (DT == nullptr && (NumInterestingAllocas > 1 || + if (DT == nullptr && (InterestingAllocas.size() > 1 || !F->hasFnAttribute(Attribute::OptimizeNone))) { DeleteDT = std::make_unique(*F); DT = DeleteDT.get(); @@ -631,21 +623,21 @@ 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);