[msan] Fix shadow computation for partially undefined constant fixed-length vectors

For each element of the vector, the shadow is initialized iff the
element is not undefined/poisoned. Previously, partially undefined
constant fixed-length vectors always had fully initialized shadows,
which leads to false negatives.

Updates the tests from #143823

Note: since this patch corrects a false negative, MSan can now detect more bugs (corollary: code/tests that passed MSan may start failing).

Author: thurstond

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -1989,6 +1989,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       }
       return Shadow;
     }
+    // Handle fully undefined values
+    // (partially undefined constant vectors are handled later)
     if (UndefValue *U = dyn_cast<UndefValue>(V)) {
       Value *AllOnes = (PropagateShadow && PoisonUndef) ? getPoisonedShadow(V)
                                                         : getCleanShadow(V);
@@ -2086,8 +2088,25 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       return ShadowPtr;
     }
 
-    // TODO: Partially undefined vectors are handled by the fall-through case
-    //       below (see partial-poison.ll); this causes false negatives.
+    // Check for partially undefined constant vectors
+    // TODO: scalable vectors (this is hard because we do not have IRBuilder)
+    if (   isa<FixedVectorType>(V->getType())
+        && isa<Constant>(V)
+        && (cast<Constant>(V))->containsUndefOrPoisonElement()
+        && PropagateShadow
+        && PoisonUndef) {
+      unsigned NumElems = (cast<FixedVectorType>(V->getType()))->getNumElements();
+      SmallVector<Constant *, 32> ShadowVector(NumElems);
+      for (unsigned i = 0; i != NumElems; ++i) {
+        Constant *Elem = (cast<Constant>(V))->getAggregateElement(i);
+        ShadowVector[i] = isa<UndefValue>(Elem) ? getPoisonedShadow(Elem) : getCleanShadow(Elem);
+      }
+
+      Value *ShadowConstant = ConstantVector::get(ShadowVector);
+      LLVM_DEBUG(dbgs() << "Partial undef constant vector: " << *V << " ==> " << *ShadowConstant << "\n");
+
+      return ShadowConstant;
+    }
 
     // For everything else the shadow is zero.
     return getCleanShadow(V);

llvm/test/Instrumentation/MemorySanitizer/partial-poison.ll

@@ -1,8 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
 ; RUN: opt < %s -S -passes='msan' 2>&1 | FileCheck %s
 ;
-; Test case to show that MSan computes shadows for partially poisoned vectors
-; as fully initialized, resulting in false negatives.
+; Regression test case for computing shadows of partially poisoned vectors
 
 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"
@@ -11,7 +10,7 @@ define <2 x i64> @left_poison(ptr %add.ptr) sanitize_memory {
 ; CHECK-LABEL: define <2 x i64> @left_poison(
 ; CHECK-SAME: ptr [[ADD_PTR:%.*]]) #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store <2 x i64> <i64 -1, i64 0>, ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> <i64 poison, i64 42>
 ;
   ret <2 x i64> <i64 poison, i64 42>
@@ -21,7 +20,7 @@ define <2 x i64> @right_poison(ptr %add.ptr) sanitize_memory {
 ; CHECK-LABEL: define <2 x i64> @right_poison(
 ; CHECK-SAME: ptr [[ADD_PTR:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store <2 x i64> <i64 0, i64 -1>, ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> <i64 42, i64 poison>
 ;
   ret <2 x i64> <i64 42, i64 poison>
@@ -51,7 +50,7 @@ define <2 x i64> @left_undef(ptr %add.ptr) sanitize_memory {
 ; CHECK-LABEL: define <2 x i64> @left_undef(
 ; CHECK-SAME: ptr [[ADD_PTR:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store <2 x i64> <i64 -1, i64 0>, ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> <i64 undef, i64 42>
 ;
   ret <2 x i64> <i64 undef, i64 42>
@@ -61,7 +60,7 @@ define <2 x i64> @right_undef(ptr %add.ptr) sanitize_memory {
 ; CHECK-LABEL: define <2 x i64> @right_undef(
 ; CHECK-SAME: ptr [[ADD_PTR:%.*]]) #[[ATTR0]] {
 ; CHECK-NEXT:    call void @llvm.donothing()
-; CHECK-NEXT:    store <2 x i64> zeroinitializer, ptr @__msan_retval_tls, align 8
+; CHECK-NEXT:    store <2 x i64> <i64 0, i64 -1>, ptr @__msan_retval_tls, align 8
 ; CHECK-NEXT:    ret <2 x i64> <i64 42, i64 undef>
 ;
   ret <2 x i64> <i64 42, i64 undef>