[mte] support more complicated lifetimes (e.g. for exceptions).

[fmayer]

newline at eof

[fmayer]

unnecessary comment

[eugenis]

LGTM, nice!

llvm/include/llvm/Transforms/Utils/MemoryTaggingSupport.h

if (NumCoveredExits == ReachableRetVec.size()) {
     // Signal the caller to remove the lifetime end call for this alloca.
     return false;
   }
   return true;
 }
 
 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
                         const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
-                        const DominatorTree &DT, size_t MaxLifetimes);
+                        const DominatorTree *DT, size_t MaxLifetimes);
 } // namespace llvm
 
 #endif

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

 
 static cl::opt<unsigned> ClScanLimit("stack-tagging-merge-init-scan-limit",
                                      cl::init(40), cl::Hidden);
 
 static cl::opt<unsigned>
     ClMergeInitSizeLimit("stack-tagging-merge-init-size-limit", cl::init(272),
                          cl::Hidden);
 
+static cl::opt<size_t> ClMaxLifetimes(
+    "stack-tagging-max-lifetimes-for-alloca", cl::Hidden, cl::init(3),
+    cl::ReallyHidden,
+    cl::desc("How many lifetime ends to handle for a single alloca."),
+    cl::Optional);
+
 static const Align kTagGranuleSize = Align(16);
 
 namespace {
 
 class InitializerBuilder {
   uint64_t Size;
   const DataLayout *DL;
   Value *BasePtr;
for (auto &I : Allocas) {
     Instruction *TagPCall =
         IRB.CreateCall(TagP, {Constant::getNullValue(Info.AI->getType()), Base,
                               ConstantInt::get(IRB.getInt64Ty(), Info.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,
+                           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 (UnrecognizedLifetimes.empty() && Info.LifetimeStart.size() == 1 &&
-        Info.LifetimeEnd.size() == 1 && !CallsReturnTwice) {
+    if (UnrecognizedLifetimes.empty() && StandardLifetime &&
+        !CallsReturnTwice) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
-      IntrinsicInst *End = Info.LifetimeEnd[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 ||
           !forAllReachableExits(*DT, *PDT, Start, Info.LifetimeEnd, RetVec,
-                                TagEnd))
+                                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);
       for (auto &RI : RetVec) {
         untagAlloca(AI, RI, Size);
       }
       // We may have inserted tag/untag outside of any lifetime interval.

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

for (auto &KV : AllocasToInstrument) {
     size_t AlignedSize = alignTo(Size, Mapping.getObjectAlignment());
     auto TagEnd = [&](Instruction *Node) {
       IRB.SetInsertPoint(Node);
       Value *UARTag = getUARTag(IRB, StackTag);
       tagAlloca(IRB, AI, UARTag, AlignedSize);
     };
     bool StandardLifetime =
         UnrecognizedLifetimes.empty() &&
-        isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, GetDT(),
+        isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, &GetDT(),
                            ClMaxLifetimes);
     if (ShouldDetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IRB.SetInsertPoint(Start->getNextNode());
       tagAlloca(IRB, AI, Tag, Size);
       if (!forAllReachableExits(GetDT(), GetPDT(), Start, Info.LifetimeEnd,
                                 RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 
 namespace llvm {
 namespace {
 bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
-                                 const DominatorTree &DT, size_t MaxLifetimes) {
+                                 const DominatorTree *DT, size_t MaxLifetimes) {
   // If we have too many lifetime ends, give up, as the algorithm below is N^2.
   if (Insts.size() > MaxLifetimes)
     return true;
   for (size_t I = 0; I < Insts.size(); ++I) {
     for (size_t J = 0; J < Insts.size(); ++J) {
       if (I == J)
         continue;
-      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, &DT))
+      if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, DT))
         return true;
     }
   }
   return false;
 }
 } // namespace
 
 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
                         const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
-                        const DominatorTree &DT, size_t MaxLifetimes) {
+                        const DominatorTree *DT, size_t MaxLifetimes) {
   // An alloca that has exactly one start and end in every possible execution.
   // If it has multiple ends, they have to be unreachable from each other, so
   // at most one of them is actually used for each execution of the function.
   return LifetimeStart.size() == 1 &&
          (LifetimeEnd.size() == 1 ||
           (LifetimeEnd.size() > 0 &&
            !maybeReachableFromEachOther(LifetimeEnd, DT, MaxLifetimes)));
 }
 } // namespace llvm

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

+; RUN: opt -S -aarch64-stack-tagging -stack-tagging-use-stack-safety=0 %s -o - | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-arm-unknown-eabi"
+
+define  void @f(i1 %cond) local_unnamed_addr sanitize_memtag {
+start:
+; CHECK-LABEL: start:
+  %a = alloca i8, i32 48, align 8
+  call void @llvm.lifetime.start.p0i8(i64 48, i8* nonnull %a)
+; CHECK: call void @llvm.aarch64.settag(i8* %a.tag, i64 48)
+  br i1 %cond, label %next0, label %next1
+
+next0:
+; CHECK-LABEL: next0:
+; CHECK: call void @llvm.aarch64.settag
+  call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a)
+  br label %exit0
+
+exit0:
+; CHECK-LABEL: exit0:
+; CHECK-NOT: call void @llvm.aarch64.settag
+  ret void
+
+next1:
+; CHECK-LABEL: next1:
+; CHECK: call void @llvm.aarch64.settag
+  call void @llvm.lifetime.end.p0i8(i64 40, i8* nonnull %a)
+  br label %exit1
+
+exit1:
+; CHECK-LABEL: exit1:
+; CHECK-NOT: call void @llvm.aarch64.settag
+  ret void
+}
+
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)