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 @@ -143,6 +143,7 @@ #include "llvm/ADT/APInt.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DepthFirstIterator.h" +#include "llvm/ADT/SmallSet.h" #include "llvm/ADT/SmallString.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringExtras.h" @@ -1023,6 +1024,7 @@ : Shadow(S), Origin(O), OrigIns(I) {} }; SmallVector InstrumentationList; + SmallSet AllocaSet; SmallVector StoreList; MemorySanitizerVisitor(Function &F, MemorySanitizer &MS, @@ -1279,6 +1281,10 @@ VAHelper->finalizeInstrumentation(); + // Poison allocas that weren't instrumented at llvm.lifetime.start. + for (AllocaInst *AI : AllocaSet) + delayedVisitAllocaInst(*AI); + bool InstrumentWithCalls = ClInstrumentationWithCallThreshold >= 0 && InstrumentationList.size() + StoreList.size() > (unsigned)ClInstrumentationWithCallThreshold; @@ -2536,6 +2542,27 @@ return false; } + void handleLifetimeStart(IntrinsicInst &I) { + if (!PoisonStack) + return; + ConstantInt *Len = dyn_cast(I.getArgOperand(0)); + if (Len->isMinusOne()) + return; + + IRBuilder<> IRB(I.getNextNode()); + DenseMap AllocaForValue; + AllocaInst *AI = + llvm::findAllocaForValue(I.getArgOperand(1), AllocaForValue); + if (!AI) + return; + if (AllocaSet.count(AI)) + AllocaSet.erase(AI); + if (MS.CompileKernel) + instrumentAllocaKmsan(*AI, IRB, Len); + else + instrumentAllocaUserspace(*AI, IRB, Len); + } + void handleBswap(IntrinsicInst &I) { IRBuilder<> IRB(&I); Value *Op = I.getArgOperand(0); @@ -2951,6 +2978,9 @@ void visitIntrinsicInst(IntrinsicInst &I) { switch (I.getIntrinsicID()) { + case Intrinsic::lifetime_start: + handleLifetimeStart(I); + break; case Intrinsic::bswap: handleBswap(I); break; @@ -3413,9 +3443,7 @@ } } - void visitAllocaInst(AllocaInst &I) { - setShadow(&I, getCleanShadow(&I)); - setOrigin(&I, getCleanOrigin()); + void delayedVisitAllocaInst(AllocaInst &I) { IRBuilder<> IRB(I.getNextNode()); const DataLayout &DL = F.getParent()->getDataLayout(); uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType()); @@ -3429,6 +3457,14 @@ instrumentAllocaUserspace(I, IRB, Len); } + void visitAllocaInst(AllocaInst &I) { + setShadow(&I, getCleanShadow(&I)); + setOrigin(&I, getCleanOrigin()); + // We'll get to this alloca later unless it's poisoned at the corresponding + // llvm.lifetime.start. + AllocaSet.insert(&I); + } + void visitSelectInst(SelectInst& I) { IRBuilder<> IRB(&I); // a = select b, c, d 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 @@ -89,3 +89,67 @@ ; KMSAN: call void @__msan_unpoison_alloca(i8* {{.*}}, i64 20) ; CHECK: ret void +; Check that every llvm.lifetime.start() causes poisoning of locals. +define void @lifetime_start() sanitize_memory { +entry: + %x = alloca i32, align 4 + %c = bitcast i32* %x to i8* + br label %another_bb + +another_bb: + call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %c) + store i32 7, i32* %x + call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %c) + call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %c) + store i32 8, i32* %x + call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %c) + ret void +} + +; CHECK-LABEL: define void @lifetime_start( +; CHECK-LABEL: entry: +; CHECK: %x = alloca i32 +; CHECK-LABEL: another_bb: + +; CHECK: call void @llvm.lifetime.start +; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false) +; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4) +; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4, +; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4, + +; CHECK: call void @llvm.lifetime.start +; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false) +; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4) +; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4, +; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4, + +; CHECK: ret void + +; If an object is variable sized, ignore llvm.lifetime.start. +define void @lifetime_start_var(i64 %cnt) sanitize_memory { +entry: + %x = alloca i32, i64 %cnt, align 4 + %c = bitcast i32* %x to i8* + call void @llvm.lifetime.start.p0i8(i64 -1, i8* nonnull %c) + call void @llvm.lifetime.end.p0i8(i64 -1, i8* nonnull %c) + ret void +} + +; CHECK-LABEL: define void @lifetime_start_var( +; CHECK-LABEL: entry: +; CHECK: %x = alloca i32, i64 %cnt +; CHECK: %[[A:.*]] = mul i64 4, %cnt +; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false) +; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 %[[A]]) +; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 %[[A]], +; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 %[[A]], +; CHECK: call void @llvm.lifetime.start +; CHECK: call void @llvm.lifetime.end + +; CHECK: ret void + + + + +declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) +declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)