[MTE] [HWASan] Use LoopInfo for reachability queries.

[fmayer]

Summary

The reachability queries default to "reachable" after exploring too many
basic blocks. LoopInfo helps it skip over the whole loop.

[eugenis]

Could you add a test?

[fmayer]

add test

[fmayer]

In D127917#3602467, @eugenis wrote:

Could you add a test?

Done.

[eugenis]

LGTM

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

 //
 //===----------------------------------------------------------------------===//
 #ifndef LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 #define LLVM_TRANSFORMS_UTILS_MEMORYTAGGINGSUPPORT_H
 
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Support/Alignment.h"
 
 namespace llvm {
 class DominatorTree;
 class DbgVariableIntrinsic;
 class IntrinsicInst;
 class PostDominatorTree;
 class AllocaInst;
 class Instruction;
 namespace memtag {
 // For an alloca valid between lifetime markers Start and Ends, call the
 // Callback for all possible exits out of the lifetime in the containing
 // function, which can return from the instructions in RetVec.
 //
 // Returns whether Ends covered all possible exits. If they did not,
 // the caller should remove Ends to ensure that work done at the other
 // exits does not happen outside of the lifetime.
 bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
-                          const Instruction *Start,
+                          const LoopInfo &LI, const Instruction *Start,
                           const SmallVectorImpl<IntrinsicInst *> &Ends,
                           const SmallVectorImpl<Instruction *> &RetVec,
                           llvm::function_ref<void(Instruction *)> Callback);
 
 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
                         const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
-                        const DominatorTree *DT, size_t MaxLifetimes);
+                        const DominatorTree *DT, const LoopInfo *LI,
+                        size_t MaxLifetimes);
 
 Instruction *getUntagLocationIfFunctionExit(Instruction &Inst);
 
 struct AllocaInfo {
   AllocaInst *AI;
   SmallVector<IntrinsicInst *, 2> LifetimeStart;
   SmallVector<IntrinsicInst *, 2> LifetimeEnd;
   SmallVector<DbgVariableIntrinsic *, 2> DbgVariableIntrinsics;

llvm/lib/Target/AArch64/AArch64StackTagging.cpp

 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 #include <cassert>
 #include <iterator>
+#include <memory>
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "aarch64-stack-tagging"
 
 static cl::opt<bool> ClMergeInit(
     "stack-tagging-merge-init", cl::Hidden, cl::init(true),
bool AArch64StackTagging::runOnFunction(Function &Fn) {
   if (auto *P = getAnalysisIfAvailable<PostDominatorTreeWrapperPass>())
     PDT = &P->getPostDomTree();
 
   if (PDT == nullptr) {
     DeletePDT = std::make_unique<PostDominatorTree>(*F);
     PDT = DeletePDT.get();
   }
 
+  std::unique_ptr<LoopInfo> DeleteLI;
+  LoopInfo *LI = nullptr;
+  if (auto *LIWP = getAnalysisIfAvailable<LoopInfoWrapperPass>()) {
+    LI = &LIWP->getLoopInfo();
+  } else {
+    DeleteLI = std::make_unique<LoopInfo>(*DT);
+    LI = DeleteLI.get();
+  }
+
   SetTagFunc =
       Intrinsic::getDeclaration(F->getParent(), Intrinsic::aarch64_settag);
 
   Instruction *Base = insertBaseTaggedPointer(SInfo.AllocasToInstrument, DT);
 
   int NextTag = 0;
   for (auto &I : SInfo.AllocasToInstrument) {
     memtag::AllocaInfo &Info = I.second;
for (auto &I : SInfo.AllocasToInstrument) {
     TagPCall->setOperand(0, Info.AI);
 
     // 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.
     bool StandardLifetime =
         SInfo.UnrecognizedLifetimes.empty() &&
-        memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT,
+        memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, DT, LI,
                                    ClMaxLifetimes) &&
         !SInfo.CallsReturnTwice;
     if (StandardLifetime) {
       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,
+          !memtag::forAllReachableExits(*DT, *PDT, *LI, Start, Info.LifetimeEnd,
                                         SInfo.RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }
     } else {
       uint64_t Size = *Info.AI->getAllocationSizeInBits(*DL) / 8;
       Value *Ptr = IRB.CreatePointerCast(TagPCall, IRB.getInt8PtrTy());
       tagAlloca(AI, &*IRB.GetInsertPoint(), Ptr, Size);

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

public:
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
   bool isInterestingAlloca(const AllocaInst &AI);
   void tagAlloca(IRBuilder<> &IRB, AllocaInst *AI, Value *Tag, size_t Size);
   Value *tagPointer(IRBuilder<> &IRB, Type *Ty, Value *PtrLong, Value *Tag);
   Value *untagPointer(IRBuilder<> &IRB, Value *PtrLong);
   bool instrumentStack(memtag::StackInfo &Info, Value *StackTag,
-                       const DominatorTree &DT, const PostDominatorTree &PDT);
+                       const DominatorTree &DT, const PostDominatorTree &PDT,
+                       const LoopInfo &LI);
   Value *readRegister(IRBuilder<> &IRB, StringRef Name);
   bool instrumentLandingPads(SmallVectorImpl<Instruction *> &RetVec);
   Value *getNextTagWithCall(IRBuilder<> &IRB);
   Value *getStackBaseTag(IRBuilder<> &IRB);
   Value *getAllocaTag(IRBuilder<> &IRB, Value *StackTag, AllocaInst *AI,
                       unsigned AllocaNo);
   Value *getUARTag(IRBuilder<> &IRB, Value *StackTag);
 

 static bool isLifetimeIntrinsic(Value *V) {
   auto *II = dyn_cast<IntrinsicInst>(V);
   return II && II->isLifetimeStartOrEnd();
 }
 
 bool HWAddressSanitizer::instrumentStack(memtag::StackInfo &SInfo,
                                          Value *StackTag,
                                          const DominatorTree &DT,
-                                         const PostDominatorTree &PDT) {
+                                         const PostDominatorTree &PDT,
+                                         const LoopInfo &LI) {
   // Ideally, we want to calculate tagged stack base pointer, and rewrite all
   // alloca addresses using that. Unfortunately, offsets are not known yet
   // (unless we use ASan-style mega-alloca). Instead we keep the base tag in a
   // temp, shift-OR it into each alloca address and xor with the retag mask.
   // This generates one extra instruction per alloca use.
   unsigned int I = 0;
 
   for (auto &KV : SInfo.AllocasToInstrument) {
for (auto &KV : SInfo.AllocasToInstrument) {
     };
     // 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.
     bool StandardLifetime =
         SInfo.UnrecognizedLifetimes.empty() &&
         memtag::isStandardLifetime(Info.LifetimeStart, Info.LifetimeEnd, &DT,
-                                   ClMaxLifetimes) &&
+                                   &LI, ClMaxLifetimes) &&
         !SInfo.CallsReturnTwice;
     if (DetectUseAfterScope && StandardLifetime) {
       IntrinsicInst *Start = Info.LifetimeStart[0];
       IRB.SetInsertPoint(Start->getNextNode());
       tagAlloca(IRB, AI, Tag, Size);
-      if (!memtag::forAllReachableExits(DT, PDT, Start, Info.LifetimeEnd,
+      if (!memtag::forAllReachableExits(DT, PDT, LI, Start, Info.LifetimeEnd,
                                         SInfo.RetVec, TagEnd)) {
         for (auto *End : Info.LifetimeEnd)
           End->eraseFromParent();
       }
     } else {
       tagAlloca(IRB, AI, Tag, Size);
       for (auto *RI : SInfo.RetVec)
         TagEnd(RI);
bool HWAddressSanitizer::sanitizeFunction(Function &F,
   emitPrologue(EntryIRB,
                /*WithFrameRecord*/ ClRecordStackHistory &&
                    Mapping.WithFrameRecord &&
                    !SInfo.AllocasToInstrument.empty());
 
   if (!SInfo.AllocasToInstrument.empty()) {
     const DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(F);
     const PostDominatorTree &PDT = FAM.getResult<PostDominatorTreeAnalysis>(F);
+    const LoopInfo &LI = FAM.getResult<LoopAnalysis>(F);
     Value *StackTag =
         ClGenerateTagsWithCalls ? nullptr : getStackBaseTag(EntryIRB);
-    instrumentStack(SInfo, StackTag, DT, PDT);
+    instrumentStack(SInfo, StackTag, DT, PDT, LI);
   }
 
   // If we split the entry block, move any allocas that were originally in the
   // entry block back into the entry block so that they aren't treated as
   // dynamic allocas.
   if (EntryIRB.GetInsertBlock() != &F.getEntryBlock()) {
     InsertPt = &*F.getEntryBlock().begin();
     for (Instruction &I :

llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp

 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/IntrinsicInst.h"
 
 namespace llvm {
 namespace memtag {
 namespace {
 bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
-                                 const DominatorTree *DT, size_t MaxLifetimes) {
+                                 const DominatorTree *DT, const LoopInfo *LI,
+                                 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, LI))
         return true;
     }
   }
   return false;
 }
 } // namespace
 
 bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
-                          const Instruction *Start,
+                          const LoopInfo &LI, const Instruction *Start,
                           const SmallVectorImpl<IntrinsicInst *> &Ends,
                           const SmallVectorImpl<Instruction *> &RetVec,
                           llvm::function_ref<void(Instruction *)> Callback) {
   if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) {
     Callback(Ends[0]);
     return true;
   }
   SmallPtrSet<BasicBlock *, 2> EndBlocks;
   for (auto *End : Ends) {
     EndBlocks.insert(End->getParent());
   }
   SmallVector<Instruction *, 8> ReachableRetVec;
   unsigned NumCoveredExits = 0;
   for (auto *RI : RetVec) {
-    if (!isPotentiallyReachable(Start, RI, nullptr, &DT))
+    if (!isPotentiallyReachable(Start, RI, nullptr, &DT, &LI))
       continue;
     ReachableRetVec.push_back(RI);
     // If there is an end in the same basic block as the return, we know for
     // sure that the return is covered. Otherwise, we can check whether there
     // is a way to reach the RI from the start of the lifetime without passing
     // through an end.
     if (EndBlocks.count(RI->getParent()) > 0 ||
-        !isPotentiallyReachable(Start, RI, &EndBlocks, &DT)) {
+        !isPotentiallyReachable(Start, RI, &EndBlocks, &DT, &LI)) {
       ++NumCoveredExits;
     }
   }
   // If there's a mix of covered and non-covered exits, just put the untag
   // on exits, so we avoid the redundancy of untagging twice.
   if (NumCoveredExits == ReachableRetVec.size()) {
     for (auto *End : Ends)
       Callback(End);
   } else {
     for (auto *RI : ReachableRetVec)
       Callback(RI);
     // We may have inserted untag outside of the lifetime interval.
     // 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, const LoopInfo *LI,
+                        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)));
+           !maybeReachableFromEachOther(LifetimeEnd, DT, LI, MaxLifetimes)));
 }
 
 Instruction *getUntagLocationIfFunctionExit(Instruction &Inst) {
   if (isa<ReturnInst>(Inst)) {
     if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall())
       return CI;
     return &Inst;
   }

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

+; We set a low dom-tree-reachability-max-bbs-to-explore to check whether the
+; loop analysis is working. Without skipping over the loop, we would need more
+; than 4 BB to reach end from entry.
+
+; RUN: opt -S -dom-tree-reachability-max-bbs-to-explore=4 -aarch64-stack-tagging %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"
+
+define dso_local void @foo(i1 %x, i32 %n) sanitize_memtag {
+entry:
+  %c = alloca [1024 x i8], align 1
+  call void @llvm.lifetime.start.p0(i64 1024, ptr nonnull %c)
+  %cmp2.not = icmp eq i32 %n, 0
+  br i1 %x, label %entry2, label %noloop
+
+entry2:
+  br i1 %cmp2.not, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+; CHECK-LABEL: for.cond.cleanup:
+; CHECK: call{{.*}}settag
+; CHECK: call{{.*}}lifetime.end
+  call void @llvm.lifetime.end.p0(i64 1024, ptr nonnull %c)
+  call void @bar(ptr noundef nonnull inttoptr (i64 120 to ptr))
+  br label %end
+
+for.body:                                         ; preds = %entry, %for.body
+  %i.03 = phi i32 [ %inc, %for.body2 ], [ 0, %entry2 ]
+  call void @bar(ptr noundef nonnull %c) #3
+  br label %for.body2
+
+for.body2:
+  %inc = add nuw nsw i32 %i.03, 1
+  %cmp = icmp ult i32 %inc, %n
+  br i1 %cmp, label %for.body, label %for.cond.cleanup, !llvm.loop !13
+
+noloop:
+; CHECK-LABEL: noloop:
+; CHECK: call{{.*}}settag
+; CHECK: call{{.*}}lifetime.end
+  call void @llvm.lifetime.end.p0(i64 1024, ptr nonnull %c)
+  br label %end
+
+end:
+; CHECK-LABEL: end:
+; CHECK-NOT: call{{.*}}settag
+  ret void
+}
+
+; Function Attrs: argmemonly mustprogress nocallback nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture) #0
+declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture) #0
+
+declare dso_local void @bar(ptr noundef)
+
+attributes #0 = { argmemonly mustprogress nocallback nofree nosync nounwind willreturn }
+
+!13 = distinct !{!13, !14}
+!14 = !{!"llvm.loop.mustprogress"}