[MSAN] Allow inserting array checks

[guiand]

Summary

Flattens arrays by ORing together all their elements.

[eugenis]

Do you want to use this in eager checks when passing an array as a function argument (does that even happen in C++?) ?

llvm/test/Instrumentation/MemorySanitizer/check-array.ll

11

just pass a [2xi24]* through the function arguments (and call it something other than main)

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1160

Please reformat

[guiand]

I don't believe this happens in C++ but wanted to add it for completeness's sake (or if some other LLVM based language might use it).

[eugenis]

In D84446#2170458, @guiand wrote:

I don't believe this happens in C++ but wanted to add it for completeness's sake (or if some other LLVM based language might use it).

Right, but I don't see you doing it in this patch.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1781–1784

@eugenis With eager checks enabled, any noundef array arguments or return values will attempt to call this code. The behavior without this patch is for this assertion to fire. The patch enables these checks to work correctly.

If you'd like, I can make this clear in the test I added by using -msan-eager-checks.

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1781–1784

Got it, MSan is simply broken without this patch - it crashes on valid IR.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1781–1784

Yep, that's why I made the test case use checking the main return value -- it's independent of eager-checks

[vitalybuka]

LGTM with clang-format

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

void paintOrigin(IRBuilder<> &IRB, Value *Origin, Value *OriginPtr,
     }
   }
 
   void storeOrigin(IRBuilder<> &IRB, Value *Addr, Value *Shadow, Value *Origin,
                    Value *OriginPtr, Align Alignment, bool AsCall) {
     const DataLayout &DL = F.getParent()->getDataLayout();
     const Align OriginAlignment = std::max(kMinOriginAlignment, Alignment);
     unsigned StoreSize = DL.getTypeStoreSize(Shadow->getType());
-    if (Shadow->getType()->isArrayTy()) {
-      paintOrigin(IRB, updateOrigin(Origin, IRB), OriginPtr, StoreSize,
-                  OriginAlignment);
-    } else {
-      Value *ConvertedShadow = convertShadowToScalar(Shadow, IRB);
-      if (auto *ConstantShadow = dyn_cast<Constant>(ConvertedShadow)) {
-        if (ClCheckConstantShadow && !ConstantShadow->isZeroValue())
-          paintOrigin(IRB, updateOrigin(Origin, IRB), OriginPtr, StoreSize,
-                      OriginAlignment);
-        return;
-      }
 
-      unsigned TypeSizeInBits =
-          DL.getTypeSizeInBits(ConvertedShadow->getType());
-      unsigned SizeIndex = TypeSizeToSizeIndex(TypeSizeInBits);
-      if (AsCall && SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
-        FunctionCallee Fn = MS.MaybeStoreOriginFn[SizeIndex];
-        Value *ConvertedShadow2 = IRB.CreateZExt(
-            ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
-        IRB.CreateCall(Fn, {ConvertedShadow2,
-                            IRB.CreatePointerCast(Addr, IRB.getInt8PtrTy()),
-                            Origin});
-      } else {
-        Value *Cmp = convertToBool(ConvertedShadow, IRB, "_mscmp");
-        Instruction *CheckTerm = SplitBlockAndInsertIfThen(
-            Cmp, &*IRB.GetInsertPoint(), false, MS.OriginStoreWeights);
-        IRBuilder<> IRBNew(CheckTerm);
-        paintOrigin(IRBNew, updateOrigin(Origin, IRBNew), OriginPtr, StoreSize,
-                    OriginAlignment);
-      }
-    }
-  }
+    unsigned TypeSizeInBits = DL.getTypeSizeInBits(ConvertedShadow->getType());
+      Value *ConvertedShadow2 =
+          IRB.CreateZExt(ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
+      IRB.CreateCall(Fn,
+                     {ConvertedShadow2,
+                      IRB.CreatePointerCast(Addr, IRB.getInt8PtrTy()), Origin});
 
   void materializeStores(bool InstrumentWithCalls) {
     for (StoreInst *SI : StoreList) {
       IRBuilder<> IRB(SI);
       Value *Val = SI->getValueOperand();
       Value *Addr = SI->getPointerOperand();
       Value *Shadow = SI->isAtomic() ? getCleanShadow(Val) : getShadow(Val);
       Value *ShadowPtr, *OriginPtr;
for (unsigned Idx = 0; Idx < Struct->getNumElements(); Idx++) {
         Aggregator = IRB.CreateOr(Aggregator, ShadowBool);
       else
         Aggregator = ShadowBool;
     }
 
     return Aggregator;
   }
 
+  // Extract combined shadow of array elements
+  Value *collapseArrayShadow(ArrayType *Array, Value *Shadow,
+                             IRBuilder<> &IRB) {
+    if (!Array->getNumElements())
+      return IRB.getIntN(/* width */ 1, /* value */ 0);
+
+    Value *FirstItem = IRB.CreateExtractValue(Shadow, 0);
+    Value *Aggregator = convertShadowToScalar(FirstItem, IRB);
+
+    for (unsigned Idx = 1; Idx < Array->getNumElements(); Idx++) {
+      Value *ShadowItem = IRB.CreateExtractValue(Shadow, Idx);
+      Value *ShadowInner = convertShadowToScalar(ShadowItem, IRB);
+      Aggregator = IRB.CreateOr(Aggregator, ShadowInner);
+    }
+    return Aggregator;
+  }
+
   /// Convert a shadow value to it's flattened variant. The resulting
   /// shadow may not necessarily have the same bit width as the input
   /// value, but it will always be comparable to zero.
   Value *convertShadowToScalar(Value *V, IRBuilder<> &IRB) {
     if (StructType *Struct = dyn_cast<StructType>(V->getType()))
       return collapseStructShadow(Struct, V, IRB);
+    if (ArrayType *Array = dyn_cast<ArrayType>(V->getType()))
+      return collapseArrayShadow(Array, V, IRB);
     Type *Ty = V->getType();
     Type *NoVecTy = getShadowTyNoVec(Ty);
     if (Ty == NoVecTy) return V;
     return IRB.CreateBitCast(V, NoVecTy);
   }
 
   // Convert a scalar value to an i1 by comparing with 0
   Value *convertToBool(Value *V, IRBuilder<> &IRB, const Twine &name = "") {
struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   ///
   /// This location will be later instrumented with a check that will print a
   /// UMR warning in runtime if the shadow value is not 0.
   void insertShadowCheck(Value *Shadow, Value *Origin, Instruction *OrigIns) {
     assert(Shadow);
     if (!InsertChecks) return;
 #ifndef NDEBUG
     Type *ShadowTy = Shadow->getType();
-    assert(
-        (isa<IntegerType>(ShadowTy) || isa<VectorType>(ShadowTy) ||
-         isa<StructType>(ShadowTy)) &&
-        "Can only insert checks for integer, vector, and struct shadow types");
+    assert((isa<IntegerType>(ShadowTy) || isa<VectorType>(ShadowTy) ||
+            isa<StructType>(ShadowTy) || isa<ArrayType>(ShadowTy)) &&
+           "Can only insert checks for integer, vector, and aggregate shadow "
+           "types");
 #endif
     InstrumentationList.push_back(
         ShadowOriginAndInsertPoint(Shadow, Origin, OrigIns));
   }
 
   /// Remember the place where a shadow check should be inserted.
   ///
   /// This location will be later instrumented with a check that will print a

llvm/test/Instrumentation/MemorySanitizer/check-array.ll

+; RUN: opt < %s -msan-eager-checks -msan-check-access-address=0 -msan-track-origins=1 -S -passes='module(msan-module),function(msan)' 2>&1 | \
+; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK,CHECK-ORIGINS %s
+
+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 noundef [2 x i24] @check_array([2 x i24]* %p) sanitize_memory {
+; CHECK: @check_array([2 x i24]* [[P:%.*]])
+; CHECK: [[O:%.*]] = load [2 x i24], [2 x i24]* [[P]]
+  %o = load [2 x i24], [2 x i24]* %p
+; CHECK: [[FIELD0:%.+]] = extractvalue [2 x i24] %_msld, 0
+; CHECK: [[FIELD1:%.+]] = extractvalue [2 x i24] %_msld, 1
+; CHECK: [[F1_OR:%.+]] = or i24 [[FIELD0]], [[FIELD1]]
+; CHECK: %_mscmp = icmp ne i24 [[F1_OR]], 0
+; CHECK: br i1 %_mscmp
+; CHECK: call void @__msan_warning
+; CHECK: ret [2 x i24] [[O]]
+  ret [2 x i24] %o
+}