[Draft][MSAN] Optimize away poisoning allocas that are always written before load

[vitalybuka]

Summary

If we know that every path from an alloca leads to a store, we can optimize away poisoning its shadow (it'll be overwritten anyway).

I'm wondering if there's a better approach for finding all these def-uses. I investigated the DominatorTree, but I'm not sure how I would use that to check *all* the uses without blowing up the runtime for instrumenting an alloca to O(N^2) (N=# uses) or so.

[lebedev.ri]

Perhaps MemSSA-based DCE can be taught about it?

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3876

TODO: get rid of Const in this function name

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3899

Store may write only part of alloca.
How useful (binary size) the patch as is?

[eugenis]

Allocas are often used through bitcast or GEP, we should handle them as well.

[guiand]

I figured if it's using GEP then it's likely not going to be storing to the entire shadow. Bitcast is overwhelmingly common though (especially bitcast->lifetime.start, which I realize I need to handle).

[eugenis]

In D83595#2145447, @guiand wrote:

I figured if it's using GEP then it's likely not going to be storing to the entire shadow. Bitcast is overwhelmingly common though (especially bitcast->lifetime.start, which I realize I need to handle).

Right. From what I've seen, small aggregates on stack are quite common as well, and they are always initialized piecemeal. This can be seen as an improvement of this change, but it would be a pretty bit change to the algorithm, so consider including that from the start.

In D83595#2145347, @lebedev.ri wrote:

Perhaps MemSSA-based DCE can be taught about it?

That's an option, but it could be harder to pull off when poisoning is outlined. The code will look something like this:

%p = alloca i32
call __msan_poison_and_set_origin(%p, 4)
...
%s_p = inttoptr(xor(ptrtoint(%p), 0x50..00)))  ; shadow address for %p
store i32 zeroinitializer, %s_p
store i32 <value>, %p

To eliminate the dead call to __msan_poison, DCE would need to know the shadow mapping.

[guiand]

I've expanded this to be able to find chained uses like through a bitcast.

Mostly a proof-of-concept; this still needs a check to make sure the final store is over the whole size of the initial type. And, of course, the main issue is that this really slows down compilation. So I think using MemSSA might be necessary for performance reasons, unless there's some algorithm tweak I can make here that would reduce the complexity.

[fhahn]

If we know that every path from an alloca leads to a store, we can optimize away poisoning its shadow (it'll be overwritten anyway).

! In D83595#2145347, @lebedev.ri wrote:

Perhaps MemSSA-based DCE can be taught about it?

MemorySSA-based DSE has similar logic (check that a store is either overwritten or never read on all paths to a function exit), but it uses a MemorySSA traversal to do so. I don't think alloca are modeled directly in MemorySSA though, so the approach cannot be directly applied I think. It might be possible to add a pseudo-memorydef writing to the alloca directly after its definition. Then the MemorySSA based logic may apply.

[guiand]

Added ability to detect piecemeal initialization.

[guiand]

This actually has very significant effects on some, but not all, benchmarks.

Running grep I observed ~8% decrease in binary with this patch. But clang sees little effect: <1%.

On a few different benchmarks from a benchmark suite, this also decreased runtime overhead by a significant amount: 8% for sha512, and 13% for qsort.

[eugenis]

In D83595#2185554, @guiand wrote:

This actually has very significant effects on some, but not all, benchmarks.

Running grep I observed ~8% decrease in binary with this patch. But clang sees little effect: <1%.

On a few different benchmarks from a benchmark suite, this also decreased runtime overhead by a significant amount: 8% for sha512, and 13% for qsort.

This sounds worthwhile.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3968

I don't think this is entirely correct. Consider this:

entry { alloca 32; store 0..16 } -> A, B
A { store 16..32 } -> B
B { load }

traversing entry->A->B will mark B as done even though there is a path where control flow arrives at B with alloca not entirely initialized.

What happens to the benchmarks if this algorithm is limited to a single basic block?

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3968

I tried to write the code so that it would look to see if *all* paths are initialized before a load. So in your example, it would look at the path entry -> A -> B, and it would see that it's fully stored, and it would continue looking at the next path. For entry -> B, if (!firstUsesAreStore) return false would come into effect, and we won't optimize away the poisoning.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3968

Oh, I see where the issue is here: the if TraversedSet.count check. I needed some way to prevent cycles in the graph, so that was the first thing I reached for, but you're right, this approach breaks the algorithm.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3968

Right. And it does not even require a cycle.

This is a dataflow problem - assign labels to graph nodes showing which bytes are always initialized at entry, and update them based on predecessor's labels until the fixed point is reached. But in this case, I suspect, most of the optimization opportunities are within a single BB.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3968

I think this could be fixed by removing the target node from the set after the recursive call is finished. But I'll measure the impact of removing the inter-block stuff entirely to see how much value it adds.

[guiand]

I've cut it down to only within a basic block as @eugenis and @vitalybuka suggested.

[eugenis]

In general, this implementation looks pretty complex and easy to get wrong. I'd prefer something along the lines of AArch64StackTagging::collectInitializers - directly calculate the offset for each store/load instruction. It might do some extra work with unrelated memory instructions, but probably not too much.

How do you handle this case?

a = alloca
b = bitcast a
lifetime_start b
store b

When scanning from lifetime_start, this code will never encounter any direct use of a, and would miss the transitive use.

Missing diff context.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3895

prefer early exits / continue(s)

3915

This function is never called with StoreOffs != 0, which seems necessary to handle a chain of GEPs.

3920

Only if same BB.

[guiand]

Flattened some control flow, updated to properly use StoreOffs, and updated tests to cover chained GEPs

[guiand]

In D83595#2188021, @eugenis wrote:

In general, this implementation looks pretty complex and easy to get wrong. I'd prefer something along the lines of AArch64StackTagging::collectInitializers - directly calculate the offset for each store/load instruction. It might do some extra work with unrelated memory instructions, but probably not too much.

I'll take a look at collectInitializers. As for the current implementation -- yeah, I always figured there would be a better way. But I tried to be pretty conservative with how I implemented it, so while we might miss some stores, we should never "forget" to poison an alloca.

How do you handle this case?

a = alloca
b = bitcast a
lifetime_start b
store b

When scanning from lifetime_start, this code will never encounter any direct use of a, and would miss the transitive use.

The code currently scans from the alloca, rather than from the lifetime_start. This might make only searching in single BB pretty limiting, since afaict an alloca can be detached from its lifetime region.

[eugenis]

In D83595#2188282, @guiand wrote:

The code currently scans from the alloca, rather than from the lifetime_start. This might make only searching in single BB pretty limiting, since afaict an alloca can be detached from its lifetime region.

Oh right. That's actually pretty weak, it's very common for lifetime to start in a basic block other than entry.

[eugenis]

You want to start scanning at the same point where poisoning is going to be inserted.

[guiand]

It seems like collectInitializers leans heavily on the isPointerOffset function, which returns an offset if two pointers have a constant difference, nullopt if they don't. The problem here is that we can't distinguish isPointerOffset == nullopt happening because the offset is determined at runtime, or because the two pointers are completely unrelated.

It's a pretty big difference, because we don't want to poison a sequence like this:

%x = alloca [ i32, i32 }
%y = alloca i32
%z = load i32, i32* %y ; isPointerOffset == false
%x0 = getelementptr { i32, i32 }, { i32, i32 }* %x, i32 0, i32 0
%x1 = getelementptr { i32, i32 }, { i32, i32 }* %x, i32 0, i32 1
store i32 0, i32* %x0
store i32 0, i32* %x1

But we want to poison a sequence like this:

%x = alloca [ i32, i32 }
%y = getelementptr { i32, i32 }, { i32, i32 }* %x, i32 0, i32 %dynamic_offs
%z = load i32, i32* %y ; isPointerOffset == false
%x0 = getelementptr { i32, i32 }, { i32, i32 }* %x, i32 0, i32 0
%x1 = getelementptr { i32, i32 }, { i32, i32 }* %x, i32 0, i32 1
store i32 0, i32* %x0
store i32 0, i32* %x1

[eugenis]

Right, that's what the isNoModRef check is for.

[guiand]

Integrated with Alias Analyzer, uses simpler mechanism for walking through BB and determining stores to alloca

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3876

No longer needed

3878

No longer needed

3899

It depends a lot on the program. Grep saw like a 8% decrease in binary size, while clang saw 0.5% to 1% decrease.

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3951

not sure why StoredBytes is a parameter and not just a local var in firstUsesAreStore

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3951

True, that was leftover from a previous version of this patch.

[vitalybuka]

FYI @kda @kstoimenov

[vitalybuka]

stashing indefinitely

[vitalybuka]

Not any time soon.

compiler-rt/test/msan/Linux/fopencookie.cpp

static int cookie_seek(void *cookie, off64_t *offset, int whence) {
   return 0;
 }
 
 static int cookie_close(void *cookie) {
   assert((uintptr_t)cookie == kMagicCookie);
   return 0;
 }
 
-void PoisonStack() { char a[8192]; }
+void PoisonStack() {
+  char a[8192];
+  __asm__ volatile("" ::"r"(a));
+}
 
 void TestPoisonStack() {
   // Verify that PoisonStack has poisoned the stack - otherwise this test is not
   // testing anything.
   char a;
   assert(__msan_test_shadow(&a - 1000, 1) == 0);
 }
 

compiler-rt/test/msan/getaddrinfo.cpp

 // RUN: %clangxx_msan -O0 %s -o %t && %run %t
 
 #include <sys/types.h>
 #include <sys/socket.h>
 #include <netdb.h>
 #include <stdlib.h>
 
 void poison_stack_ahead() {
   char buf[100000];
+  __asm__ volatile("" ::"r"(x));
   // With -O0 this poisons a large chunk of stack.
 }
 
 int main(void) {
   poison_stack_ahead();
 
   struct addrinfo *ai;
 
   // This should trigger loading of libnss_dns and friends.
   // Those libraries are typically uninstrumented.They will call strlen() on a
   // stack-allocated buffer, which is very likely to be poisoned. Test that we
   // don't report this as an UMR.
   int res = getaddrinfo("not-in-etc-hosts", NULL, NULL, &ai);
   return 0;
 }

compiler-rt/test/msan/qsort.cpp

 constexpr size_t kSize2 = 7;
 
 bool seen2;
 
 void dummy(long a, long b, long c, long d, long e) {}
 
 void poison_stack_and_param() {
   char x[10000];
+  __asm__ volatile("" ::"r"(x));
 #ifndef TEST_MSAN_EAGER_CHECKS
   int y;
   dummy(y, y, y, y, y);
 #endif
 }
 
 __attribute__((always_inline)) int cmp(long a, long b) {
   if (a < b)

compiler-rt/test/msan/vararg.cpp

 __attribute__((noinline))
 void dummy(long a, long b, long c, long d, long e) {
   __asm__ __volatile__ (""::"r"(a), "r"(b), "r"(c), "r"(d), "r"(e));
 }
 
 __attribute__((noinline))
 void poison_stack_and_param() {
   char x[10000];
+  __asm__ volatile("" ::"r"(x));
 #ifndef TEST_MSAN_EAGER_CHECKS
   int y;
   dummy(y, y, y, y, y);
 #endif
 }
 
 __attribute__((noinline))
 int sum(int n, ...) {

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

if (PoisonStack) {
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len,
                       IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy())});
     } else {
       IRB.CreateCall(MS.MsanUnpoisonAllocaFn,
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
     }
   }
 
+  // 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 firstUsesAreStore(SmallPtrSet<BasicBlock *, 10> &TraversedSet,
+                         SmallPtrSet<Value *, 10> &Users,
+                         SmallPtrSet<Value *, 10> &Defs,
+                         Instruction *CurrInst) {
+    if (TraversedSet.size() == 100) {
+      // Have to cut it off somewhere
+      return false;
+    }
+    BasicBlock *BB = CurrInst->getParent();
+    TraversedSet.insert(BB);
+
+    while (CurrInst) {
+      if (Users.count(CurrInst)) {
+        if (StoreInst *Store = dyn_cast<StoreInst>(CurrInst)) {
+          // The store could be a use if it's storing the alloca
+          // pointer somewhere. But we don't want that.
+          return Defs.count(Store->getPointerOperand());
+        } else if (IntrinsicInst *Intrinsic =
+                       dyn_cast<IntrinsicInst>(CurrInst)) {
+          // Ignore lifetime intrinsics, count others as a use.
+          Intrinsic::ID ID = Intrinsic->getIntrinsicID();
+          if (ID != Intrinsic::lifetime_start && ID != Intrinsic::lifetime_end)
+            return false;
+        } else {
+          return false;
+        }
+      }
+      CurrInst = CurrInst->getNextNonDebugInstruction();
+      if (Defs.count(CurrInst)) {
+        // Append chained uses and defs to their relevant sets
+        populateRelevantAllocaUsers(CurrInst, Defs, Users);
+      }
+    }
+
+    // Traverse to the Node's successors if we haven't yet
+    for (auto Node = succ_begin(BB); Node != succ_end(BB); Node++) {
+      if (TraversedSet.count(*Node) == 0) {
+        Instruction *NextInst = Node->getFirstNonPHIOrDbgOrLifetime();
+        if (!firstUsesAreStore(TraversedSet, Users, Defs, NextInst))
+          return false;
+      }
+    }
+    return true;
+  };
+
+  void populateRelevantAllocaUsers(Instruction *I,
+                                   SmallPtrSet<Value *, 10> &Defs,
+                                   SmallPtrSet<Value *, 10> &Users) {
+    Defs.insert(I);
+
+    for (User *User : I->users()) {
+      if (isa<BitCastInst>(User)) {
+        Defs.insert(User);
+        continue;
+      } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(User)) {
+        if (GEP->getNumIndices() == 1 && GEP->hasAllZeroIndices()) {
+          Defs.insert(GEP);
+          continue;
+        }
+      }
+      Users.insert(User);
+    }
+  }
+
   void instrumentAlloca(AllocaInst &I, Instruction *InsPoint = nullptr) {
     if (!InsPoint)
       InsPoint = &I;
     IRBuilder<> IRB(InsPoint->getNextNode());
     const DataLayout &DL = F.getParent()->getDataLayout();
     uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
     Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
-    if (I.isArrayAllocation())
+    if (I.isArrayAllocation()) {
       Len = IRB.CreateMul(Len, I.getArraySize());
-
+    } else {
+      SmallPtrSet<Value *, 10> Users, Defs;
+      SmallPtrSet<BasicBlock *, 10> TraversedSet;
+      populateRelevantAllocaUsers(&I, Defs, Users);
+      if (firstUsesAreStore(TraversedSet, Users, Defs, &I))
+        return;
+    }
     if (MS.CompileKernel)
       poisonAllocaKmsan(I, IRB, Len);
     else
       poisonAllocaUserspace(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);

llvm/test/Instrumentation/MemorySanitizer/alloca-store.ll

+; Tests if MSAN can optimize away dead-code alloca poisons
+
+; RUN: opt < %s -msan-track-origins=2 -msan-check-access-address=0 -msan-poison-stack-with-call=1 -S -passes=msan 2>&1 | FileCheck  \
+; RUN: %s "--check-prefixes=CHECK"
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32* @simple(i32 %val) sanitize_memory {
+  ; CHECK: define {{.*}} @simple
+  %p = alloca i32
+  store i32 %val, i32* %p
+  ; CHECK-NOT: call void @__msan_poison_stack
+  ret i32* %p
+}
+
+define i32* @bitcast({ i32 } %val) sanitize_memory {
+  ; CHECK: define {{.*}} @bitcast
+  %p = alloca i32
+  %ps = bitcast i32* %p to { i32 }*
+  store { i32 } %val, { i32 }* %ps
+  ; CHECK-NOT: call void @__msan_poison_stack
+  ret i32* %p
+}
+
+define i32* @lifetime(i32 %val) sanitize_memory {
+  ; CHECK: define {{.*}} @lifetime
+  %p = alloca i32
+  %p8 = bitcast i32* %p to i8*
+  call void @llvm.lifetime.start(i64 4, i8* %p8)
+  store i32 %val, i32* %p
+  ; CHECK-NOT: call void @__msan_poison_stack
+  call void @llvm.lifetime.end(i64 4, i8* %p8)
+  ret i32* %p
+}
+
+declare void @llvm.lifetime.start(i64, i8* nocapture)
+declare void @llvm.lifetime.end(i64, i8* nocapture)

llvm/test/Instrumentation/MemorySanitizer/alloca.ll

 ; RUN: opt < %s -msan -msan-kernel=1 -S | FileCheck %s --check-prefixes=CHECK,KMSAN
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 define void @static() sanitize_memory {
 entry:
   %x = alloca i32, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 
 ; CHECK-LABEL: define void @static(
 ; 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
 
 
 define void @dynamic() sanitize_memory {
 entry:
   br label %l
 l:
   %x = alloca i32, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 
 ; CHECK-LABEL: define void @dynamic(
 ; 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
 
 define void @array() sanitize_memory {
 entry:
   %x = alloca i32, i64 5, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 
 ; CHECK-LABEL: define void @array(
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 20, i1 false)
 ; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 20)
 ; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 20,
 ; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 20,
 ; CHECK: ret void
 
 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
 }
 
 ; CHECK-LABEL: define void @array_non_const(
 ; 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: ret void
 
 ; Check that the local is unpoisoned in the absence of sanitize_memory
 define void @unpoison_local() {
 entry:
   %x = alloca i32, i64 5, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 
 ; 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)

llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll

 
 if.then:                                          ; preds = %entry
   ret i32 0
 
 if.else:                                          ; preds = %entry
   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]])
 
 ; Hooks for inputs usually go before the assembly statement. But here we have none,
 ; because %nr is passed by value. However we check %nr for being initialized.
 ; CHECK-CONS: [[NRC3:%.*]] = bitcast i64* %nr to i8*
 ; CHECK-CONS: call { i8*, i32* } @__msan_metadata_ptr_for_load_8(i8* [[NRC3]])