[MTE] Instrument use-after-scope for optnone functions.

[fmayer]

Summary

We always need the Dominator and PostDominatorTree for use-after-scope.

[fmayer]

Failed Tests (2):

LLVM :: CodeGen/AArch64/O0-pipeline.ll
LLVM :: CodeGen/AArch64/O3-pipeline.ll

[fmayer]

fix test failure.

do one change at a time.

[eugenis]

LGTM

TBH I'm on the fence about whether poison-at-lifetime is an optimization vs "core functionality", but lets go with this change. The important part is not to run the analysis when the function does not have the attribute (that's what the *-pipeline tests check).

[fmayer]

In D119993#3330790, @eugenis wrote:

LGTM

TBH I'm on the fence about whether poison-at-lifetime is an optimization vs "core functionality", but lets go with this change. The important part is not to run the analysis when the function does not have the attribute (that's what the *-pipeline tests check).

I think there should be a very high bar for disabling functionality with user-visible effects as "optimizations" because it adds lots of potential for confusion. In fact, from the perspective of the resulting code, the non-UAS code is potentially more optimized than the UAS one.

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

private:
   const StackSafetyGlobalInfo *SSI = nullptr;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     if (UseStackSafety)
       AU.addRequired<StackSafetyGlobalInfoWrapperPass>();
     if (MergeInit)
       AU.addRequired<AAResultsWrapperPass>();
+    AU.addRequired<PostDominatorTreeWrapperPass>();
+    AU.addRequired<DominatorTreeWrapperPass>();
   }
 };
 
 } // end anonymous namespace
 
 char AArch64StackTagging::ID = 0;
 
 INITIALIZE_PASS_BEGIN(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(StackSafetyGlobalInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
 INITIALIZE_PASS_END(AArch64StackTagging, DEBUG_TYPE, "AArch64 Stack Tagging",
                     false, false)
 
 FunctionPass *llvm::createAArch64StackTaggingPass(bool IsOptNone) {
   return new AArch64StackTagging(IsOptNone);
 }
 
 Instruction *AArch64StackTagging::collectInitializers(Instruction *StartInst,
memtag::StackInfoBuilder SIB(
       [this](const AllocaInst &AI) { return isInterestingAlloca(AI); });
   for (Instruction &I : instructions(F))
     SIB.visit(I);
   memtag::StackInfo &SInfo = SIB.get();
 
   if (SInfo.AllocasToInstrument.empty())
     return false;
 
-  std::unique_ptr<DominatorTree> DeleteDT;
-  DominatorTree *DT = nullptr;
-  if (auto *P = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
-    DT = &P->getDomTree();
-
-  if (DT == nullptr && (SInfo.AllocasToInstrument.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();
-  }
+  DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
+  PostDominatorTree &PDT =
+      getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
 
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
-  Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
+  Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, &DT);
 
   int NextTag = 0;
   for (auto &I : SInfo.AllocasToInstrument) {
     memtag::AllocaInfo &Info = I.second;
     assert(Info.AI && isInterestingAlloca(*Info.AI));
     TrackingVH<Instruction> OldAI = Info.AI;
     memtag::alignAndPadAlloca(Info, kTagGranuleSize);
     AllocaInst *AI = Info.AI;
     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(), Tag)});
     if (Info.AI->hasName())
       TagPCall->setName(Info.AI->getName() + ".tag");
     Info.AI->replaceAllUsesWith(TagPCall);
     TagPCall->setOperand(0, Info.AI);
 
     bool StandardLifetime =
         SInfo.UnrecognizedLifetimes.empty() &&
-        memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,
+        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 (SInfo.UnrecognizedLifetimes.empty() && StandardLifetime &&
         !SInfo.CallsReturnTwice) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       uint64_t Size =
           cast<ConstantInt>(Start->getArgOperand(0))->getZExtValue();
       Size = alignTo(Size, kTagGranuleSize);
       tagAlloca(AI, Start->getNextNode(), Start->getArgOperand(1), Size);
 
       auto TagEnd = [&](Instruction *Node) { untagAlloca(AI, Node, Size); };
-      if (!DT || !PDT ||
-          !memtag::forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd,
+      if (!memtag::forAllReachableExits(DT, PDT, Start, Info.LifetimeEnd,
                                         SInfo.RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }
     } else {
       uint64_t Size = Info.AI->getAllocationSizeInBits(*DL).getValue() / 8;
       Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
       tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);

llvm/test/CodeGen/AArch64/stack-tagging.ll

 ; CHECK:  call void @llvm.aarch64.settag(
 ; CHECK:  call void @use8(
 ; CHECK:  call void @llvm.aarch64.settag(
 ; CHECK:  call void @llvm.lifetime.end.p0i8(
 ; CHECK:  br label
 ; CHECK:  ret void
 
 
+define void @ScopeOptnone(i32 %b) sanitize_memtag optnone noinline {
+entry:
+  %x = alloca i32, align 4
+  %tobool = icmp eq i32 %b, 0
+  br i1 %tobool, label %if.end, label %if.then
+
+if.then:
+  %0 = bitcast i32* %x to i8*
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %0)
+  call void @use8(i8* %0) #3
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %0)
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+; CHECK-LABEL: define void @ScopeOptnone(
+; CHECK:  br i1
+; CHECK:  call void @llvm.lifetime.start.p0i8(
+; CHECK:  call void @llvm.aarch64.settag(
+; CHECK:  call void @use8(
+; CHECK:  call void @llvm.aarch64.settag(
+; CHECK:  call void @llvm.lifetime.end.p0i8(
+; CHECK:  br label
+; CHECK:  ret void
+
+
 ; Spooked by the multiple lifetime ranges, StackTagging remove all of them and sets tags on entry and exit.
 define void @BadScope(i32 %b) sanitize_memtag {
 entry:
   %x = alloca i32, align 4
   %tobool = icmp eq i32 %b, 0
   br i1 %tobool, label %if.end, label %if.then
 
 if.then: