MSAN: Allow emitting checks for struct types

[guiand]

Summary

This is a precursor to the general eager-checks patch.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1155

this can be improved now by moving the origin update code behind the shadow check

1409

extra "ShadowItem ="

1410

I don't think this is possible

1413

I'd prefer if this generated less pointless comparisons.

It would be nice if CreateICmpNE(i1 x, i1 false) automatically folded to "x". I think this is a reasonable change to make in IRBuilder, even though this is probably a very unlikely situation.

As an alternative, make convertShadowToScalar return i1 and remove icmp in the callers (call it something like convertShadowToBool). Check Combiner class how not to emit a "false" constant. As for the Maybe* callers, it's ok to not handle aggregates there.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1155

This check covers arrays as well, so I don't think I can get rid of it completely but I can change it to check for isArrayType(). I wrote some code to let convertShadowToScalar handle Array types as well, but I thought I would leave that to a separate patch.

[guiand]

Addressed comments

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1155

Ack.

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

10

Please check that the test passes when LLVM is build without assertions. AFAIR temp names are lost in that case.

19

There is still this unnecessary icmp.

[guiand]

Don't emit unnecessary icmp

[guiand]

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

10

As far as I can tell it's passing without assertions enabled, yeah.

[guiand]

Fixed msan_x86_bts_asm test after optimizing away icmp i1 %x, false

[eugenis]

Please upload with context.

[guiand]

Reuploaded with context

[eugenis]

LGTM with 2 comments

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1420–1423

Mention that the result may not be the same bit width as the original type in the comment (this is the usual property of shadow types).

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

20

could this line say

; CHECK: br i1 [[F1_OR]]

?

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1421

this code do not uses {} for one line blocks

1434

do you see a difference without this check?

1436

this comment is not very useful

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1434

If you look at msan_x86_bts_asm.ll, this check eliminates an unnecessary icmp in some cases.

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()->isAggregateType()) {
       paintOrigin(IRB, updateOrigin(Origin, IRB), OriginPtr, StoreSize,
                   OriginAlignment);
     } else {
-      Value *ConvertedShadow = convertToShadowTyNoVec(Shadow, IRB);
+      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 =
void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     // This may invalidate some of the following checks and needs to be done
     // at the very end.
   }
 
   void materializeOneCheck(Instruction *OrigIns, Value *Shadow, Value *Origin,
                            bool AsCall) {
     IRBuilder<> IRB(OrigIns);
     LLVM_DEBUG(dbgs() << "  SHAD0 : " << *Shadow << "\n");
-    Value *ConvertedShadow = convertToShadowTyNoVec(Shadow, IRB);
+    Value *ConvertedShadow = convertShadowToScalar(Shadow, IRB);
     LLVM_DEBUG(dbgs() << "  SHAD1 : " << *ConvertedShadow << "\n");
 
     if (auto *ConstantShadow = dyn_cast<Constant>(ConvertedShadow)) {
       if (ClCheckConstantShadow && !ConstantShadow->isZeroValue()) {
         insertWarningFn(IRB, Origin);
       }
       return;
     }
Type *getShadowTy(Type *OrigTy) {
     if (ArrayType *AT = dyn_cast<ArrayType>(OrigTy)) {
       return ArrayType::get(getShadowTy(AT->getElementType()),
                             AT->getNumElements());
     }
     if (StructType *ST = dyn_cast<StructType>(OrigTy)) {
       SmallVector<Type*, 4> Elements;
       for (unsigned i = 0, n = ST->getNumElements(); i < n; i++)
         Elements.push_back(getShadowTy(ST->getElementType(i)));
       StructType *Res = StructType::get(*MS.C, Elements, ST->isPacked());
       LLVM_DEBUG(dbgs() << "getShadowTy: " << *ST << " ===> " << *Res << "\n");
       return Res;
     }
     uint32_t TypeSize = DL.getTypeSizeInBits(OrigTy);
     return IntegerType::get(*MS.C, TypeSize);
   }
 
   /// Flatten a vector type.
   Type *getShadowTyNoVec(Type *ty) {
     if (VectorType *vt = dyn_cast<VectorType>(ty))
       return IntegerType::get(*MS.C,
                               vt->getPrimitiveSizeInBits().getFixedSize());
     return ty;
   }
 
+  /// Extract combined shadow of struct elements as a bool
+  Value *collapseStructShadow(StructType *Struct, Value *Shadow,
+                              IRBuilder<> &IRB) {
+    Value *Aggregator = IRB.getIntN(/* width */ 1, /* value */ 0);
+    for (unsigned i = 0; i < Struct->getNumElements(); i++) {
+      Value *ShadowItem = ShadowItem = IRB.CreateExtractValue(Shadow, i);
+      assert(ShadowItem && "Failed to create ExtractValue instruction!");
+      Value *ShadowInner = convertShadowToScalar(ShadowItem, IRB);
+      Value *ShadowBool =
+          IRB.CreateICmpNE(ShadowInner, getCleanShadow(ShadowInner));
+      Aggregator = IRB.CreateOr(Aggregator, ShadowBool);
+    }
+    return Aggregator;
+  }
+
   /// Convert a shadow value to it's flattened variant.
-  Value *convertToShadowTyNoVec(Value *V, IRBuilder<> &IRB) {
+  Value *convertShadowToScalar(Value *V, IRBuilder<> &IRB) {
+    if (StructType *Struct = dyn_cast<StructType>(V->getType())) {
+      return collapseStructShadow(Struct, V, IRB);
+    }
     Type *Ty = V->getType();
     Type *NoVecTy = getShadowTyNoVec(Ty);
     if (Ty == NoVecTy) return V;
     return IRB.CreateBitCast(V, NoVecTy);
   }
 
   /// Compute the integer shadow offset that corresponds to a given
   /// application address.
   ///
   /// Offset = (Addr & ~AndMask) ^ XorMask
   Value *getShadowPtrOffset(Value *Addr, IRBuilder<> &IRB) {
     Value *OffsetLong = IRB.CreatePointerCast(Addr, MS.IntptrTy);
 
     uint64_t AndMask = MS.MapParams->AndMask;
     if (AndMask)
       OffsetLong =
           IRB.CreateAnd(OffsetLong, ConstantInt::get(MS.IntptrTy, ~AndMask));
 
     uint64_t XorMask = MS.MapParams->XorMask;
     if (XorMask)
       OffsetLong =
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)) &&
-           "Can only insert checks for integer and vector shadow types");
+    assert(
+        (isa<IntegerType>(ShadowTy) || isa<VectorType>(ShadowTy) ||
+         isa<StructType>(ShadowTy)) &&
+        "Can only insert checks for integer, vector, and struct 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
Combiner &Add(Value *OpShadow, Value *OpOrigin) {
       if (MSV->MS.TrackOrigins) {
         assert(OpOrigin);
         if (!Origin) {
           Origin = OpOrigin;
         } else {
           Constant *ConstOrigin = dyn_cast<Constant>(OpOrigin);
           // No point in adding something that might result in 0 origin value.
           if (!ConstOrigin || !ConstOrigin->isNullValue()) {
-            Value *FlatShadow = MSV->convertToShadowTyNoVec(OpShadow, IRB);
+            Value *FlatShadow = MSV->convertShadowToScalar(OpShadow, IRB);
             Value *Cond =
                 IRB.CreateICmpNE(FlatShadow, MSV->getCleanShadow(FlatShadow));
             Origin = IRB.CreateSelect(Cond, OpOrigin, Origin);
           }
         }
       }
       return *this;
     }

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

+; RUN: opt < %s -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"
+
+; CHECK-LABEL: @main
+define { i32, i8 } @main() sanitize_memory {
+; CHECK: %p = inttoptr i64 0 to { i32, i8 }*
+  %p = inttoptr i64 0 to { i32, i8 } *
+; CHECK: %o = load { i32, i8 }, { i32, i8 }* %p
+  %o = load { i32, i8 }, { i32, i8 } *%p
+; CHECK: [[FIELD0:%.+]] = extractvalue { i32, i8 } %_msld, 0
+; CHECK: [[F0_POISONED:%.+]] = icmp ne i32 [[FIELD0]]
+; CHECK: [[F0_OR:%.+]] = or i1 false, [[F0_POISONED]]
+; CHECK: [[FIELD1:%.+]] = extractvalue { i32, i8 } %_msld, 1
+; CHECK: [[F1_POISONED:%.+]] = icmp ne i8 [[FIELD1]]
+; CHECK: [[F1_OR:%.+]] = or i1 [[F0_OR]], [[F1_POISONED]]
+; CHECK: %_mscmp = icmp ne i1 [[F1_OR]], false
+; CHECK: br i1 %_mscmp
+; CHECK: call void @__msan_warning
+; CHECK: ret { i32, i8 } %o
+  ret { i32, i8 } %o
+}