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 @@ -3627,6 +3627,49 @@ } } + // Determine whether all paths from this alloca lead to storing into it. + // If true, we can omit poisoning the alloca because it'll just be + // overwritten anyway. + bool firstUsesAreConstStore(SmallPtrSet &TraversedSet, + const SmallPtrSet &Users, + AllocaInst &Alloca, BasicBlock *BB) { + if (TraversedSet.size() == 100) { + // Have to cut it off somewhere + return false; + } + TraversedSet.insert(BB); + + Instruction *CurrInst; + // In the Alloca's basic block, we shouldn't start looking from + // the begining of the block because we'll just find the alloca + if (BB == Alloca.getParent()) { + CurrInst = Alloca.getNextNonDebugInstruction(); + } else { + CurrInst = BB->getFirstNonPHIOrDbgOrLifetime(); + } + + while (CurrInst) { + if (Users.find(CurrInst) != Users.end()) { + if (StoreInst *Store = dyn_cast(CurrInst)) { + // The store could be a use if it's storing the alloca + // pointer somewhere. But we don't want that. + return Store->getPointerOperand() == &Alloca; + } else { + return false; + } + } + CurrInst = CurrInst->getNextNonDebugInstruction(); + } + + // Traverse to the Node's successors if we haven't yet + for (auto Node = succ_begin(BB); Node != succ_end(BB); Node++) { + if (TraversedSet.find(*Node) == TraversedSet.end()) + if (!firstUsesAreConstStore(TraversedSet, Users, Alloca, *Node)) + return false; + } + return true; + }; + void instrumentAlloca(AllocaInst &I, Instruction *InsPoint = nullptr) { if (!InsPoint) InsPoint = &I; @@ -3637,6 +3680,15 @@ if (I.isArrayAllocation()) Len = IRB.CreateMul(Len, I.getArraySize()); + SmallPtrSet Users; + for (User *User : I.users()) { + Users.insert(User); + } + + SmallPtrSet TraversedSet; + if (firstUsesAreConstStore(TraversedSet, Users, I, I.getParent())) + return; + if (MS.CompileKernel) poisonAllocaKmsan(I, IRB, Len); else 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 @@ -30,6 +30,7 @@ define void @static() sanitize_memory { entry: %x = alloca i32, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -47,6 +48,7 @@ br label %l l: %x = alloca i32, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -61,6 +63,7 @@ define void @array() sanitize_memory { entry: %x = alloca i32, i64 5, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -75,6 +78,7 @@ define void @array_non_const(i64 %cnt) sanitize_memory { entry: %x = alloca i32, i64 %cnt, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } @@ -91,6 +95,7 @@ define void @unpoison_local() { entry: %x = alloca i32, i64 5, align 4 + %y = getelementptr i32, i32* %x, i32 1 ret void } diff --git a/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll b/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll --- a/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll +++ b/llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll @@ -58,9 +58,9 @@ ret i32 1 } -; %nr is first poisoned, then unpoisoned (written to). Need to optimize this in the future. -; CHECK: [[NRC1:%.*]] = bitcast i64* %nr to i8* -; CHECK: call void @__msan_poison_alloca(i8* [[NRC1]]{{.*}}) +; %nr is first poisoned, then unpoisoned (written to). No need to poison. +; CHECK-NOT: [[NRC1:%.*]] = bitcast i64* %nr to i8* +; CHECK-NOT: call void @__msan_poison_alloca(i8* [[NRC1]]{{.*}}) ; CHECK: [[NRC2:%.*]] = bitcast i64* %nr to i8* ; CHECK: call { i8*, i32* } @__msan_metadata_ptr_for_store_8(i8* [[NRC2]])