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

[fmayer]

Summary

this is to get the implementation closer to HWASan

[fmayer]

keep old behaviour

[fmayer]

rebase onto fix

[fmayer]

saving files is hard

[fmayer]

format

[eugenis]

Not sure about this. I kind of like that isInterestingAlloca is called only once for each alloca - this way it can run heavier analysis (it does not right now) and it is also allowed to change its mind (ex. after we add some instrumentation to the function).

[fmayer]

In D118960#3297625, @eugenis wrote:

Not sure about this. I kind of like that isInterestingAlloca is called only once for each alloca - this way it can run heavier analysis (it does not right now) and it is also allowed to change its mind (ex. after we add some instrumentation to the function).

That's fair. I was torn between the two approaches (the other way would be to change HWASan to resemble the current state of MTE stack tagging). I thought this has less room for error from forgetting to check for interestingness (because it's not in the list to begin with), and because as you said isInterestingAlloca doesn't do anything particularly expensive.

[eugenis]

LGTM

yeah nvm I'm ok with this

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

 
 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;
-    int Tag; // -1 for non-tagged allocations
   };
 
   const bool MergeInit;
   const bool UseStackSafety;
 
 public:
   static char ID; // Pass ID, replacement for typeid
 

 void AArch64StackTagging::untagAlloca(AllocaInst *AI, Instruction *InsertBefore,
                                       uint64_t Size) {
   IRBuilder<> IRB(InsertBefore);
   IRB.CreateCall(SetTagFunc, {IRB.CreatePointerCast(AI, IRB.getInt8PtrTy()),
                               ConstantInt::get(IRB.getInt64Ty(), Size)});
 }
 
 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;
     }
     PrologueBB = DT->findNearestCommonDominator(PrologueBB, AI->getParent());
   }
   assert(PrologueBB);
 
bool AArch64StackTagging::runOnFunction(Function &Fn) {
 
   if (UseStackSafety)
     SSI = &getAnalysis<StackSafetyGlobalInfoWrapperPass>().getResult();
   F = &Fn;
   DL = &Fn.getParent()->getDataLayout();
   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;
 
   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)) {
-      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;
     }
 
     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)
-        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);
     if (ExitUntag)
       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();
   }
 
   std::unique_ptr<PostDominatorTree> DeletePDT;
   PostDominatorTree *PDT = nullptr;
   if (auto *P = getAnalysisIfAvailable<PostDominatorTreeWrapperPass>())
     PDT = &P->getPostDomTree();
 
   if (PDT == nullptr && !F->hasFnAttribute(Attribute::OptimizeNone)) {
     DeletePDT = std::make_unique<PostDominatorTree>(*F);
     PDT = DeletePDT.get();
   }
 
   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);
     TagPCall->setOperand(0, Info.AI);
 
     bool StandardLifetime =
         UnrecognizedLifetimes.empty() &&
         isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,