MSan: handle llvm.lifetime.start intrinsic

[glider]

Summary

When a variable goes into scope several times within a single function
or when two variables from different scopes share a stack slot it may
be incorrect to poison such scoped locals at the beginning of the
function.
In the former case it may lead to false negatives (see
https://github.com/google/sanitizers/issues/590), in the latter - to
incorrect reports (because only one origin remains on the stack).

If Clang emits lifetime intrinsics with known size for such scoped
variables we insert code poisoning them after each call to
llvm.lifetime.start(). This is done on a best-effort basis.
If for a certain intrinsic we fail to find a corresponding alloca, we
fall back to poisoning allocas for the whole function, as it's now
impossible to tell which alloca was missed.

The new instrumentation may slow down hot loops containing local
variables with lifetime intrinsics, so we allow disabling it with
-mllvm -msan-handle-lifetime-intrinsics=false.

[glider]

Removed errs()

[eugenis]

Please test the case when alloca can not be found.
See https://bugs.llvm.org/show_bug.cgi?id=41481 for the reference.

With this change msan will start detecting new bugs, ex.:
for (int i = 0; ...; ++i) {

int x;
if (i == 0) x=0;
read(x);

}

It can also slow down some programs a lot.

Do you plan to remove SanitizeMemory condition from shouldEmitLifetimeMarkers ?

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2556

IMHO failure to find the underlying alloca must result in poisoning of all allocas in the prologue - exactly because we don't know which one we've missed.
Alternatively, since we know the size, we could instrument this lifetime.start anyway using an invalid origin id.

3446

Let's call this function InstrumentAlloca, and rename instrumentAllocaUserspace to poisonAllocaUserspace.

[glider]

In D60617#1464971, @eugenis wrote:

Please test the case when alloca can not be found.
See https://bugs.llvm.org/show_bug.cgi?id=41481 for the reference.

Ok, will do. The repro in the bug doesn't produce IR containing a select statement for me, but the idea is clear.

With this change msan will start detecting new bugs, ex.:
for (int i = 0; ...; ++i) {

int x;
if (i == 0) x=0;
read(x);

}

It can also slow down some programs a lot.

Agreed. This also makes less sense without origins - do you think we shall hide it behind track-origins>0 ?

Do you plan to remove SanitizeMemory condition from shouldEmitLifetimeMarkers ?

Sorry, I don't get it. We won't see any lifetime markers with that condition removed, right?
What's the point of this fix then?

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2556

I like your first proposal better.
Using an invalid origin is impractical, as we won't be able to even find the function this poisoned local belongs to.

3446

same for instrumentAllocaKmsan?

[eugenis]

In D60617#1467093, @glider wrote:

In D60617#1464971, @eugenis wrote:

Please test the case when alloca can not be found.
See https://bugs.llvm.org/show_bug.cgi?id=41481 for the reference.

Ok, will do. The repro in the bug doesn't produce IR containing a select statement for me, but the idea is clear.

With this change msan will start detecting new bugs, ex.:
for (int i = 0; ...; ++i) {

int x;
if (i == 0) x=0;
read(x);

}

It can also slow down some programs a lot.

Agreed. This also makes less sense without origins - do you think we shall hide it behind track-origins>0 ?

No, I think it is a good change. But it would be nice to have an llvm flag to back out just in case.

Do you plan to remove SanitizeMemory condition from shouldEmitLifetimeMarkers ?

Sorry, I don't get it. We won't see any lifetime markers with that condition removed, right?
What's the point of this fix then?

We will see a lot more lifetime markers.
I'm talking about this:

// Disable lifetime markers in msan builds.
// FIXME: Remove this when msan works with lifetime markers.
if (LangOpts.Sanitize.has(SanitizerKind::Memory))
  return false;

http://llvm.org/viewvc/llvm-project?view=revision&revision=235613

[glider]

Added a test containing an intrinsic without a known alloca.
Handled this case by falling back to instrumenting only allocas.
Added a flag to disable llvm.lifetime.start handling.

[glider]

s/instrumentAllocaKmsan/poisonAllocaKmsan
s/instrumentAllocaUserspace/poisonAllocaUserspace

[glider]

Updated the patch, PTAL.

Re: shouldEmitLifetimeMarkers() - thanks, I wasn't aware we disable lifetime markers with -O0.
I don't need that for the kernel (which is built with -O2), but I can sure revert that CL if you think it's ok for the userspace.

FWIW those bits are in clang/lib/CodeGen/CodeGenFunction.cpp now, and removing just them doesn't break any tests.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1030

This does not have to be so complicated.
Get rid of FallbackAllocaSet; instead remove instructions from AllocaSet as they are being poisoned at lifetime calls, then poison remaining alloca instructions.

This way if untraceable lifetime was seen, you can simply skip the lifetime poisoning step.

Please rename InstrumentOnlyAllocas to something that does not suggest the we are instrumenting _only_ allocas (and nothing else).

1300

Move this logic to a separate function.

2550

What's wrong with variable length allocas?

[glider]

Fixed comments.

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1030

Agreed, we don't need FallbackAllocaSet, we can just poison allocas from LifetimeStartList as if they were in AllocaSet (see the updated patch).

1291

Reused instrumentAlloca() here instead of calling KMSAN/userspace-specific functions.

2550

You're right, fixed this.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1290

What if InstrumentLifetimeStart == false, but there is an alloca with multiple lifetime.start ? It looks like you'll poison it several times in a row.

[glider]

Fixed the case when an alloca was initialized twice.

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1290

Fixed this by adding alloca to AllocaSet instead of calling instrumentAlloca(*Item.second, nullptr).

[glider]

ptal

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2556

What if AllocaInst is visited after all of its lifetime.start calls? I'm not 100% that this is impossible, and then this code will poison both allocainst and lifetime.start.

IMHO, the visitor should simply collect allocas and lifetimes, and then the finalization code can erase allocas as it instruments their lifetime.starts.

[glider]

Simplified the logic per Evgeniy's request.
No test changes.

[glider]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

2556

You're right, we'd better not over-complicate the collection logic.
I however doubt a lifetime intrinsic can occur before an alloca, because it's using that alloca's result.

[glider]

PTAL

[eugenis]

LGTM
Thank you!

Please remove the condition in clang as well in a separate patch. This thing that you've just fixed is the only reason for MSan condition in shouldEmitLifetimeMarkers() to exist, so I'm afraid it's on you now to clean that up.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 ///  KernelMemorySanitizer only supports X86_64 at the moment.
 ///
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #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"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
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;
+  SmallSet<AllocaInst *, 16> AllocaSet;
   SmallVector<StoreInst *, 16> StoreList;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
                          const TargetLibraryInfo &TLI)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI) {
     bool SanitizeFunction = F.hasFnAttribute(Attribute::SanitizeMemory);
     InsertChecks = SanitizeFunction;
     PropagateShadow = SanitizeFunction;
     PoisonStack = SanitizeFunction && ClPoisonStack;
     PoisonUndef = SanitizeFunction && ClPoisonUndef;
     // FIXME: Consider using SpecialCaseList to specify a list of functions that
for (PHINode *PN : ShadowPHINodes) {
       for (size_t v = 0; v < NumValues; v++) {
         PNS->addIncoming(getShadow(PN, v), PN->getIncomingBlock(v));
         if (PNO) PNO->addIncoming(getOrigin(PN, v), PN->getIncomingBlock(v));
       }
     }
 
     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;
 
     // Insert shadow value checks.
     materializeChecks(InstrumentWithCalls);
 
     // Delayed instrumentation of StoreInst.
     // This may not add new address checks.
     materializeStores(InstrumentWithCalls);
 
     return true;
   }
 
   /// Compute the shadow type that corresponds to a given Value.
   Type *getShadowTy(Value *V) {
     return getShadowTy(V->getType());
   }
 
   /// Compute the shadow type that corresponds to a given Type.
   Type *getShadowTy(Type *OrigTy) {
bool handleUnknownIntrinsic(IntrinsicInst &I) {
     if (I.doesNotAccessMemory())
       if (maybeHandleSimpleNomemIntrinsic(I))
         return true;
 
     // FIXME: detect and handle SSE maskstore/maskload
     return false;
   }
 
+  void handleLifetimeStart(IntrinsicInst &I) {
+    if (!PoisonStack)
+      return;
+    ConstantInt *Len = dyn_cast<ConstantInt>(I.getArgOperand(0));
+    if (Len->isMinusOne())
+      return;
+
+    IRBuilder<> IRB(I.getNextNode());
+    DenseMap<Value *, AllocaInst *> 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);
     Type *OpType = Op->getType();
     Function *BswapFunc = Intrinsic::getDeclaration(
       F.getParent(), Intrinsic::bswap, makeArrayRef(&OpType, 1));
     setShadow(&I, IRB.CreateCall(BswapFunc, getShadow(Op)));
     setOrigin(&I, getOrigin(Op));
Value *S = IRB.CreateCall(I.getCalledFunction(),
                               {getShadow(&I, 0), I.getOperand(1)});
     S = IRB.CreateOr(SMask, S);
     setShadow(&I, S);
     setOriginForNaryOp(I);
   }
 
   void visitIntrinsicInst(IntrinsicInst &I) {
     switch (I.getIntrinsicID()) {
+    case Intrinsic::lifetime_start:
+      handleLifetimeStart(I);
+      break;
     case Intrinsic::bswap:
       handleBswap(I);
       break;
     case Intrinsic::masked_store:
       handleMaskedStore(I);
       break;
     case Intrinsic::masked_load:
       handleMaskedLoad(I);
if (PoisonStack) {
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len,
                       IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy())});
     } else {
       IRB.CreateCall(MS.MsanUnpoisonAllocaFn,
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
     }
   }
 
-  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());
     Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
     if (I.isArrayAllocation())
       Len = IRB.CreateMul(Len, I.getArraySize());
 
     if (MS.CompileKernel)
       instrumentAllocaKmsan(I, IRB, Len);
     else
       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
     Value *B = I.getCondition();
     Value *C = I.getTrueValue();
     Value *D = I.getFalseValue();
     Value *Sb = getShadow(B);
     Value *Sc = getShadow(C);

llvm/test/Instrumentation/MemorySanitizer/alloca.ll

 
 ; CHECK-LABEL: define void @unpoison_local(
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 0, i64 20, i1 false)
 ; CALL: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 0, i64 20, i1 false)
 ; ORIGIN-NOT: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 20,
 ; 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)