[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

3633

TODO: get rid of Const in this function name

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3656

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

3948

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

3948

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

3948

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

3948

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

3948

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

3656

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

3848

No longer needed

3850

No longer needed

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3927

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

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3927

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/getaddrinfo.cpp

 
 void poison_stack_ahead() {
   char buf[100000];
+  __asm__ volatile("" ::"r"(buf));
   // With -O0 this poisons a large chunk of stack.
 }
 

compiler-rt/test/msan/qsort.cpp

 
 void poison_stack_and_param() {
   char x[10000];
+  __asm__ volatile("" ::"r"(x));
   int y;
   dummy(y, y, y, y, y);
 }

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallString.h"

     }
   }
 
+  // Bytes in an alloca already stored to
+  using StoredSet = BitVector;
+  // Mapping of Value modifying a pointer => Offset it's pointing to
+  using RedefOffsets = DenseMap<Value *, uint64_t>;
+  // Set of instructions consuming a given alloca
+  using UserSet = SmallPtrSet<Value *, 10>;
+
+  // 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(UserSet &Users, RedefOffsets &Defs,
+                         Instruction *CurrInst, StoredSet StoredBytes) {
+    while (CurrInst) {
+      if (Users.count(CurrInst)) {
+        if (StoreInst *Store = dyn_cast<StoreInst>(CurrInst)) {
+          const DataLayout &DL = F.getParent()->getDataLayout();
+          uint64_t StoreSize =
+              DL.getTypeStoreSize(Store->getValueOperand()->getType());
+          // The store could be a use if it's storing the alloca
+          // pointer somewhere. But we don't want that.
+          auto Def = Defs.find(Store->getPointerOperand());
+          if (Def != Defs.end()) {
+            StoredBytes.set(Def->second, Def->second + StoreSize);
+          } else {
+            return false;
+          }
+        } 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 if (CallInst *Call = dyn_cast<CallInst>(CurrInst)) {
+          // We want to consider instrumented memory writing functions
+          // as stores here.
+          if (Call->getCalledFunction() == MS.MemcpyFn.getCallee() ||
+              Call->getCalledFunction() == MS.MemsetFn.getCallee() ||
+              Call->getCalledFunction() == MS.MemmoveFn.getCallee()) {
+            Value *Dest = Call->getArgOperand(0);
+            Value *StoreSizeVal = Call->getArgOperand(2);
+
+            auto Def = Defs.find(Dest);
+            if (Def == Defs.end())
+              return false;
+
+            // Can only decide this covers the whole shadow if we can
+            // check the constant store size.
+            if (ConstantInt *StoreSizeConst =
+                    dyn_cast<ConstantInt>(StoreSizeVal)) {
+              uint64_t StoreSize = StoreSizeConst->getLimitedValue();
+              StoredBytes.set(Def->second, Def->second + StoreSize);
+            } else {
+              return false;
+            }
+          } else {
+            return false;
+          }
+        } else {
+          return false;
+        }
+      }
+
+      if (StoredBytes.all()) {
+        return true;
+      }
+
+      CurrInst = CurrInst->getNextNonDebugInstruction();
+      if (Defs.count(CurrInst)) {
+        // Append chained uses and defs to their relevant sets
+        populateRelevantAllocaUsers(CurrInst, Defs, Users);
+      }
+    }
+
+    // Reached end of basic block and still haven't stored everything
+    // TODO: Check successive BBs
+    return false;
+  };
+
+  void populateRelevantAllocaUsers(Instruction *I, RedefOffsets &Defs,
+                                   UserSet &Users, uint64_t StoreOffs = 0) {
+    Defs.insert({I, StoreOffs});
+
+    for (User *User : I->users()) {
+      if (isa<BitCastInst>(User)) {
+        Defs.insert({User, StoreOffs});
+        continue;
+      } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(User)) {
+        const DataLayout &DL = F.getParent()->getDataLayout();
+        APInt Offs(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
+        if (GEP->accumulateConstantOffset(DL, Offs)) {
+          Defs.insert({GEP, StoreOffs + Offs.getLimitedValue()});
+          continue;
+        }
+      }
+      Users.insert(User);
+    }
+  }
+
   void instrumentAlloca(AllocaInst &I, Instruction *InsPoint = nullptr) {
     if (!InsPoint)
       InsPoint = &I;

     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 {
+      RedefOffsets Defs;
+      UserSet Users;
+      StoredSet StoredBytes(TypeSize, false);
+
+      populateRelevantAllocaUsers(&I, Defs, Users);
+      if (firstUsesAreStore(Users, Defs, &I, StoredBytes))
+        return;
+    }
 
     if (MS.CompileKernel)
       poisonAllocaKmsan(I, IRB, Len);

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
+}
+
+define i32* @with_memset(i8 %val) sanitize_memory {
+  ; CHECK: define {{.*}} @with_memset
+  %p = alloca i32
+  %p8 = bitcast i32* %p to i8*
+  call void @llvm.memset(i8* %p8, i8 0, i64 4, i1 false)
+  ; CHECK-NOT: call void @__msan_poison_stack
+  ret i32* %p
+}
+
+%struct.multi = type { i32, i8, i8, i8, i8 }
+
+define i32* @piecemeal(i8 %val) sanitize_memory {
+  ; CHECK: define {{.*}} @piecemeal
+  %p = alloca %struct.multi
+  %p0 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 0
+  %p1 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 1
+  %p2 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 2
+  %p3 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 3
+  %p4 = getelementptr %struct.multi, %struct.multi* %p, i32 0, i32 4
+  store i32 0, i32* %p0
+  store i8 0, i8* %p1
+  store i8 0, i8* %p2
+  store i8 0, i8* %p3
+  store i8 0, i8* %p4
+  ; CHECK-NOT: call void @__msan_poison_stack
+  ret i32* %p0
+}
+
+declare void @llvm.memset(i8*, i8, i64, i1)
+declare void @llvm.lifetime.start(i64, i8* nocapture)
+declare void @llvm.lifetime.end(i64, i8* nocapture)

llvm/test/Instrumentation/MemorySanitizer/alloca.ll

 define void @static() sanitize_memory {
 entry:
   %x = alloca i32, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 

   br label %l
 l:
   %x = alloca i32, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 

 define void @array() sanitize_memory {
 entry:
   %x = alloca i32, i64 5, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   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
 }
 

 define void @unpoison_local() {
 entry:
   %x = alloca i32, i64 5, align 4
+  %y = getelementptr i32, i32* %x, i32 1
   ret void
 }
 

llvm/test/Instrumentation/MemorySanitizer/msan_x86_bts_asm.ll

   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]])