[AddressSanitizer] Add support for (constant-)masked loads and stores.

[filcab]

Summary

This patch adds support for instrumenting masked loads and stores under
ASan, if they have a constant mask.

isInterestingMemoryAccess now supports returning a mask to be applied to
the loads, and instrumentMop will use it to generate additional checks.

Added tests for v4i32 v8i32, and v4p0i32 (~v4i64) for both loads and
stores (as well as a test to verify we don't add checks to non-constant
masks).

[RKSimon]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

505–507

Add MaybeMask description to comment.

1003–1004

Add MaybeMask description to comment.

[kcc]

Could you please also add a test to compiler-rt/test/asan?
(you can do it as a separate Phab patch, if that's more convenient)

lib/Transforms/Instrumentation/AddressSanitizer.cpp

1191

could you please outline this into its own function, just as you did with doInstrumentAddress?

[filcab]

Update CreateGEP call to add 0 as the first parameter.

[filcab]

Forgot to outline part of a function.

[filcab]

Extract instrumentMaskedLoadOrStore()

[RKSimon]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

1055

Use isa<> instead as we don't use the return value from dyn_cast<> ?

1059

I don't think we can guarantee ConstantInt type here (UndefValue etc.?)

1152

I don't think we can guarantee ConstantInt type here (UndefValue etc.?)

[filcab]

Addressed rksimon's comments:
Basically try to discard undef in a safe way.

Brought back (fixed) test.

[vitalybuka]

LGTM

lib/Transforms/Instrumentation/AddressSanitizer.cpp

1064

You don't call this with nullptr, so for consistency with other output arguments i'd remove "if (MaybeMask)" check

[filcab]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

1003–1004

Deleted this one, since there's no reason to have those comments on both the prototype and impl.

1064

It's called in many places outside the function I changed, though.
In the future we will need to refactor this to support scatter/gathers, so I'd like to avoid doing too much now.

lib/Transforms/Instrumentation/AddressSanitizer.cpp

uint64_t getAllocaSizeInBytes(const AllocaInst &AI) const {
     uint64_t SizeInBytes =
         AI.getModule()->getDataLayout().getTypeAllocSize(Ty);
     return SizeInBytes * ArraySize;
   }
   /// Check if we want (and can) handle this alloca.
   bool isInterestingAlloca(const AllocaInst &AI);
 
   /// If it is an interesting memory access, return the PointerOperand
   /// and set IsWrite/Alignment. Otherwise return nullptr.
+  /// MaybeMask is an output parameter for the mask Value, if we're looking at a
+  /// masked load/store.
   Value *isInterestingMemoryAccess(Instruction *I, bool *IsWrite,
-                                   uint64_t *TypeSize, unsigned *Alignment);
+                                   uint64_t *TypeSize, unsigned *Alignment,
+                                   Value **MaybeMask = nullptr);
   void instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis, Instruction *I,
                      bool UseCalls, const DataLayout &DL);
   void instrumentPointerComparisonOrSubtraction(Instruction *I);
   void instrumentAddress(Instruction *OrigIns, Instruction *InsertBefore,
                          Value *Addr, uint32_t TypeSize, bool IsWrite,
                          Value *SizeArgument, bool UseCalls, uint32_t Exp);
   void instrumentUnusualSizeOrAlignment(Instruction *I, Value *Addr,
                                         uint32_t TypeSize, bool IsWrite,
bool IsInteresting =
        (!ClSkipPromotableAllocas || !isAllocaPromotable(&AI)) &&
        // inalloca allocas are not treated as static, and we don't want
        // dynamic alloca instrumentation for them as well.
        !AI.isUsedWithInAlloca());
 
   ProcessedAllocas[&AI] = IsInteresting;
   return IsInteresting;
 }
 
-/// If I is an interesting memory access, return the PointerOperand
-/// and set IsWrite/Alignment. Otherwise return nullptr.
 Value *AddressSanitizer::isInterestingMemoryAccess(Instruction *I,
                                                    bool *IsWrite,
                                                    uint64_t *TypeSize,
-                                                   unsigned *Alignment) {
+                                                   unsigned *Alignment,
+                                                   Value **MaybeMask) {
   // Skip memory accesses inserted by another instrumentation.
   if (I->getMetadata("nosanitize")) return nullptr;
 
   // Do not instrument the load fetching the dynamic shadow address.
   if (LocalDynamicShadow == I)
     return nullptr;
 
   Value *PtrOperand = nullptr;
if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
     *Alignment = 0;
     PtrOperand = RMW->getPointerOperand();
   } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(I)) {
     if (!ClInstrumentAtomics) return nullptr;
     *IsWrite = true;
     *TypeSize = DL.getTypeStoreSizeInBits(XCHG->getCompareOperand()->getType());
     *Alignment = 0;
     PtrOperand = XCHG->getPointerOperand();
+  } else if (auto CI = dyn_cast<CallInst>(I)) {
+    auto *F = dyn_cast<Function>(CI->getCalledValue());
+    if (F && (F->getName().startswith("llvm.masked.load.") ||
+              F->getName().startswith("llvm.masked.store."))) {
+      unsigned OpOffset = 0;
+      if (F->getName().startswith("llvm.masked.store.")) {
+        // Masked store has an initial operand for the value.
+        OpOffset = 1;
+        *IsWrite = true;
+      } else {
+        *IsWrite = false;
+      }
+      // Only instrument if the mask is constant for now.
+      if (dyn_cast<ConstantVector>(CI->getOperand(2 + OpOffset))) {
+        auto BasePtr = CI->getOperand(0 + OpOffset);
+        auto Ty = cast<PointerType>(BasePtr->getType())->getElementType();
+        *TypeSize = DL.getTypeStoreSizeInBits(Ty);
+        *Alignment = (unsigned)cast<ConstantInt>(CI->getOperand(1 + OpOffset))
+                         ->getZExtValue();
+        if (MaybeMask)
+          *MaybeMask = CI->getOperand(2 + OpOffset);
+        PtrOperand = BasePtr;
+      }
+    }
   }
 
   // Do not instrument acesses from different address spaces; we cannot deal
   // with them.
   if (PtrOperand) {
     Type *PtrTy = cast<PointerType>(PtrOperand->getType()->getScalarType());
     if (PtrTy->getPointerAddressSpace() != 0)
       return nullptr;
void AddressSanitizer::instrumentPointerComparisonOrSubtraction(
   Value *Param[2] = {I->getOperand(0), I->getOperand(1)};
   for (Value *&i : Param) {
     if (i->getType()->isPointerTy())
       i = IRB.CreatePointerCast(i, IntptrTy);
   }
   IRB.CreateCall(F, Param);
 }
 
+static void doInstrumentAddress(AddressSanitizer *Pass, Instruction *I,
+                                Value *Addr, unsigned Alignment,
+                                unsigned Granularity, uint32_t TypeSize,
+                                bool IsWrite, Value *SizeArgument,
+                                bool UseCalls, uint32_t Exp) {
+  // Instrument a 1-, 2-, 4-, 8-, or 16- byte access with one check
+  // if the data is properly aligned.
+  if ((TypeSize == 8 || TypeSize == 16 || TypeSize == 32 || TypeSize == 64 ||
+       TypeSize == 128) &&
+      (Alignment >= Granularity || Alignment == 0 || Alignment >= TypeSize / 8))
+    return Pass->instrumentAddress(I, I, Addr, TypeSize, IsWrite, nullptr,
+                                   UseCalls, Exp);
+  Pass->instrumentUnusualSizeOrAlignment(I, Addr, TypeSize, IsWrite, nullptr,
+                                         UseCalls, Exp);
+}
+
 void AddressSanitizer::instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
                                      Instruction *I, bool UseCalls,
                                      const DataLayout &DL) {
   bool IsWrite = false;
   unsigned Alignment = 0;
   uint64_t TypeSize = 0;
-  Value *Addr = isInterestingMemoryAccess(I, &IsWrite, &TypeSize, &Alignment);
+  Value *MaybeMask = nullptr;
+  Value *Addr =
+      isInterestingMemoryAccess(I, &IsWrite, &TypeSize, &Alignment, &MaybeMask);
   assert(Addr);
 
   // Optimization experiments.
   // The experiments can be used to evaluate potential optimizations that remove
   // instrumentation (assess false negatives). Instead of completely removing
   // some instrumentation, you set Exp to a non-zero value (mask of optimization
   // experiments that want to remove instrumentation of this instruction).
   // If Exp is non-zero, this pass will emit special calls into runtime
   // (e.g. __asan_report_exp_load1 instead of __asan_report_load1). These calls
   // make runtime terminate the program in a special way (with a different
   // exit status). Then you run the new compiler on a buggy corpus, collect
   // the special terminations (ideally, you don't see them at all -- no false
void AddressSanitizer::instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
   }
 
   if (IsWrite)
     NumInstrumentedWrites++;
   else
     NumInstrumentedReads++;
 
   unsigned Granularity = 1 << Mapping.Scale;
-  // Instrument a 1-, 2-, 4-, 8-, or 16- byte access with one check
-  // if the data is properly aligned.
-  if ((TypeSize == 8 || TypeSize == 16 || TypeSize == 32 || TypeSize == 64 ||
-       TypeSize == 128) &&
-      (Alignment >= Granularity || Alignment == 0 || Alignment >= TypeSize / 8))
-    return instrumentAddress(I, I, Addr, TypeSize, IsWrite, nullptr, UseCalls,
-                             Exp);
-  instrumentUnusualSizeOrAlignment(I, Addr, TypeSize, IsWrite, nullptr,
-                                   UseCalls, Exp);
+  if (MaybeMask) {
+    auto Mask = cast<ConstantVector>(MaybeMask);
+    auto *VTy = cast<PointerType>(Addr->getType())->getElementType();
+    uint64_t ElemTypeSize = DL.getTypeStoreSizeInBits(VTy->getScalarType());
+    unsigned Num = VTy->getVectorNumElements();
+    for (unsigned Idx = 0; Idx < Num; ++Idx) {
+      auto Masked = cast<ConstantInt>(Mask->getOperand(Idx));
+      if (Masked->isAllOnesValue()) {
+        IRBuilder<> IRB(I);
+        auto InstrumentedAddress =
+            IRB.CreateGEP(Addr, {ConstantInt::get(IntptrTy, 0),
+                                 ConstantInt::get(IntptrTy, Idx)});
+        doInstrumentAddress(this, I, InstrumentedAddress, Alignment,
+                            Granularity, ElemTypeSize, IsWrite, nullptr,
       }
+    }
+  } else {
+    doInstrumentAddress(this, I, Addr, Alignment, Granularity, TypeSize,
+                        IsWrite, nullptr, UseCalls, Exp);
+  }
+}
 
 Instruction *AddressSanitizer::generateCrashCode(Instruction *InsertBefore,
                                                  Value *Addr, bool IsWrite,
                                                  size_t AccessSizeIndex,
                                                  Value *SizeArgument,
                                                  uint32_t Exp) {
   IRBuilder<> IRB(InsertBefore);
   Value *ExpVal = Exp == 0 ? nullptr : ConstantInt::get(IRB.getInt32Ty(), Exp);