[sanitizer] Don't always instrument on every lifetime.start

[StephenFan]

Summary

In D129448, the solution to deal with bypassed variables makes multiple
lifetime start intrinsics for the same alloca object appears in multiple
different basic block. But MemorySanitizer assumes that the variable is
always uninitialized at the place of lifetime.start intrinsic and
instruments poisoning code for it. After D129448, this assumption is broken.

This patch fix the broken assumption by analysis the dominance relation
of lifetime start/end intrinsics.

[vitalybuka]

I assume we are going to abandon this?

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <map>
 #include <memory>
 #include <string>
 #include <tuple>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "msan"
 

 struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   Function &F;
   MemorySanitizer &MS;
   SmallVector<PHINode *, 16> ShadowPHINodes, OriginPHINodes;
   ValueMap<Value*, Value*> ShadowMap, OriginMap;
   std::unique_ptr<VarArgHelper> VAHelper;
   const TargetLibraryInfo *TLI;
   Instruction *FnPrologueEnd;
+  std::unique_ptr<DominatorTree> DT;
 
   // The following flags disable parts of MSan instrumentation based on
   // exclusion list contents and command-line options.
   bool InsertChecks;
   bool PropagateShadow;
   bool PoisonStack;
   bool PoisonUndef;
 
   struct ShadowOriginAndInsertPoint {
     Value *Shadow;
     Value *Origin;
     Instruction *OrigIns;
 
     ShadowOriginAndInsertPoint(Value *S, Value *O, Instruction *I)
       : Shadow(S), Origin(O), OrigIns(I) {}
   };
   SmallVector<ShadowOriginAndInsertPoint, 16> InstrumentationList;
   bool InstrumentLifetimeStart = ClHandleLifetimeIntrinsics;
   SmallSet<AllocaInst *, 16> AllocaSet;
-  SmallVector<std::pair<IntrinsicInst *, AllocaInst *>, 16> LifetimeStartList;
+  std::map<AllocaInst *, SmallVector<IntrinsicInst *, 1>> LifetimeStarts;
+  std::map<AllocaInst *, SmallVector<IntrinsicInst *, 1>> LifetimeEnds;
   SmallVector<StoreInst *, 16> StoreList;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
                          const TargetLibraryInfo &TLI)
-      : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI) {
+      : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI),
+        DT(std::make_unique<DominatorTree>(F)) {
     bool SanitizeFunction =
         F.hasFnAttribute(Attribute::SanitizeMemory) && !ClDisableChecks;
     InsertChecks = SanitizeFunction;
     PropagateShadow = SanitizeFunction;
     PoisonStack = SanitizeFunction && ClPoisonStack;
     PoisonUndef = SanitizeFunction && ClPoisonUndef;
 
     // In the presence of unreachable blocks, we may see Phi nodes with
for (PHINode *PN : ShadowPHINodes) {
       }
     }
 
     VAHelper->finalizeInstrumentation();
 
     // Poison llvm.lifetime.start intrinsics, if we haven't fallen back to
     // instrumenting only allocas.
     if (InstrumentLifetimeStart) {
-      for (auto Item : LifetimeStartList) {
-        instrumentAlloca(*Item.second, Item.first);
-        AllocaSet.erase(Item.second);
+      for (auto Item : LifetimeStarts) {
+        for (auto LifetimeStart : Item.second) {
+          instrumentAlloca(*Item.first, LifetimeStart);
+          AllocaSet.erase(Item.first);
+        }
       }
     }
     // Poison the allocas for which we didn't instrument the corresponding
     // lifetime intrinsics.
     for (AllocaInst *AI : AllocaSet)
       instrumentAlloca(*AI);
 
     bool InstrumentWithCalls = ClInstrumentationWithCallThreshold >= 0 &&
#endif
   }
 
   void handleLifetimeStart(IntrinsicInst &I) {
     if (!PoisonStack)
       return;
     AllocaInst *AI = llvm::findAllocaForValue(I.getArgOperand(1));
     if (!AI)
       InstrumentLifetimeStart = false;
-    LifetimeStartList.push_back(std::make_pair(&I, AI));
+
+    auto hasLifetimeEndInBetween = [this, AI](const IntrinsicInst *Former,
+                                              const IntrinsicInst *Latter) {
+      assert(DT->dominates(Former, Latter) && "Former doesn't dominate Latter");
+      for (auto *LifetimeEnd : LifetimeEnds[AI]) {
+        if (DT->dominates(Former, LifetimeEnd) &&
+            DT->dominates(LifetimeEnd, Latter))
+          return true;
+      }
+      return false;
+    };
+
+    for (auto &LifetimeStart : LifetimeStarts[AI]) {
+      // If there is no lifetime.end intrinsic between two lifetime.start
+      // intrinsics, the lifetime.start intrinsic that dominated by the other
+      // does not need to do instrumentation.
+      if (DT->dominates(&I, LifetimeStart) &&
+          !hasLifetimeEndInBetween(&I, LifetimeStart))
+        LifetimeStart = &I;
+      else if (DT->dominates(LifetimeStart, &I) &&
+               !hasLifetimeEndInBetween(LifetimeStart, &I))
+        return;
+    }
+
+    LifetimeStarts[AI].push_back(&I);
+  }
+
+  void handleLifetimeEnd(IntrinsicInst &I) {
+    if (!PoisonStack)
+      return;
+    AllocaInst *AI = llvm::findAllocaForValue(I.getArgOperand(1));
+    if (!AI)
+      InstrumentLifetimeStart = false;
+    LifetimeEnds[AI].push_back(&I);
   }
 
   void handleBswap(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value *Op = I.getArgOperand(0);
     Type *OpType = Op->getType();
     Function *BswapFunc = Intrinsic::getDeclaration(
       F.getParent(), Intrinsic::bswap, makeArrayRef(&OpType, 1));

   void visitIntrinsicInst(IntrinsicInst &I) {
     switch (I.getIntrinsicID()) {
     case Intrinsic::abs:
       handleAbsIntrinsic(I);
       break;
     case Intrinsic::lifetime_start:
       handleLifetimeStart(I);
       break;
+    case Intrinsic::lifetime_end:
+      handleLifetimeEnd(I);
+      break;
     case Intrinsic::launder_invariant_group:
     case Intrinsic::strip_invariant_group:
       handleInvariantGroup(I);
       break;
     case Intrinsic::bswap:
       handleBswap(I);
       break;
     case Intrinsic::masked_store:

llvm/test/Instrumentation/MemorySanitizer/alloca.ll

 ; CHECK: %x = alloca i8
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 1 {{.*}}, i8 -1, i64 1, i1 false)
 ; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 1)
 ; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 1,
 ; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 1,
 
 ; FIXME: It is invalid to instrument here Since variable x may be initialized by callee bar.
 ; CHECK-LABEL: l1:
-; INLINE: call void @llvm.memset.p0i8.i64(i8* align 1 {{.*}}, i8 -1, i64 1, i1 false)
-; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 1)
-; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 1,
-; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 1,
+; INLINE-NOT: call void @llvm.memset.p0i8.i64(i8* align 1 {{.*}}, i8 -1, i64 1, i1 false)
+; CALL-NOT: call void @__msan_poison_stack(i8* {{.*}}, i64 1)
+; ORIGIN-NOT: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 1,
+; KMSAN-NOT: call void @__msan_poison_alloca(i8* {{.*}}, i64 1,
 
 declare void @bar(i8* noundef, i32 noundef)
 declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
 declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)