diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp --- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp @@ -194,6 +194,7 @@ #include #include #include +#include #include #include #include @@ -1027,6 +1028,7 @@ std::unique_ptr VAHelper; const TargetLibraryInfo *TLI; Instruction *FnPrologueEnd; + std::unique_ptr DT; // The following flags disable parts of MSan instrumentation based on // exclusion list contents and command-line options. @@ -1046,12 +1048,14 @@ SmallVector InstrumentationList; bool InstrumentLifetimeStart = ClHandleLifetimeIntrinsics; SmallSet AllocaSet; - SmallVector, 16> LifetimeStartList; + std::map> LifetimeStarts; + std::map> LifetimeEnds; SmallVector 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(F)) { bool SanitizeFunction = F.hasFnAttribute(Attribute::SanitizeMemory) && !ClDisableChecks; InsertChecks = SanitizeFunction; @@ -1296,9 +1300,11 @@ // 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 @@ -2701,7 +2707,40 @@ 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) { @@ -3257,6 +3296,9 @@ 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); diff --git a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll --- a/llvm/test/Instrumentation/MemorySanitizer/alloca.ll +++ b/llvm/test/Instrumentation/MemorySanitizer/alloca.ll @@ -240,10 +240,10 @@ ; 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)