[msan] Change logic of ClInstrumentationWithCallThreshold

[vitalybuka]

Summary

According to logs, ClInstrumentationWithCallThreshold is workaround
for slow backend with large number of basic blocks.
However, I can't reproduce that one, but I see significant slowdown
after ClCheckConstantShadow. Without ClInstrumentationWithCallThreshold
compiler is able to eliminate many of the branches.

So maybe we should drop ClInstrumentationWithCallThreshold completly.

For now I just change the logic to ignore constant shadow so it will
not trigger callback fallback too early.

[vitalybuka]

update

[kstoimenov]

This seems to have changed the logic for emitting callbacks to such that it would emit inlined code and once the threshold is reached it would start outlining them. Before it would either always emit outlined or always inlined for each function. Is that correct and if so it is intended?

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1153

s/folloup/follow-up/

1154

To avoid negative integer underflow.

[vitalybuka]

This comment was removed by vitalybuka.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1154

Total count is useful for debugging
maybe this way?

[vitalybuka]

In D133880#3790453, @kstoimenov wrote:

This seems to have changed the logic for emitting callbacks to such that it would emit inlined code and once the threshold is reached it would start outlining them. Before it would either always emit outlined or always inlined for each function. Is that correct and if so it is intended?

Yes. This is just workaround for huge number of basic blocks. So I don't think this particular part is important.
Alternative is to first pass handle constants, then InstrumentWithCalls, and run another pass as before, using InstrumentWithCalls.
However it looks like causing more convoluted code without a reason.

[vitalybuka]

fix compilation

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

ShadowOriginAndInsertPoint(Value *S, Value *O, Instruction *I)
         : Shadow(S), Origin(O), OrigIns(I) {}
   };
   SmallVector<ShadowOriginAndInsertPoint, 16> InstrumentationList;
   DenseMap<const DILocation *, int> LazyWarningDebugLocationCount;
   bool InstrumentLifetimeStart = ClHandleLifetimeIntrinsics;
   SmallSetVector<AllocaInst *, 16> AllocaSet;
   SmallVector<std::pair<IntrinsicInst *, AllocaInst *>, 16> LifetimeStartList;
   SmallVector<StoreInst *, 16> StoreList;
+  int64_t SplittableBlocksCount = 0;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
                          const TargetLibraryInfo &TLI)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI) {
     bool SanitizeFunction =
         F.hasFnAttribute(Attribute::SanitizeMemory) && !ClDisableChecks;
     InsertChecks = SanitizeFunction;
     PropagateShadow = SanitizeFunction;
if (MS.CompileKernel) {
       insertKmsanPrologue(IRB);
     }
 
     LLVM_DEBUG(if (!InsertChecks) dbgs()
                << "MemorySanitizer is not inserting checks into '"
                << F.getName() << "'\n");
   }
 
+  bool instrumentWithCalls(Value *V) {
+    // Constants likely will be eliminated by follow-up passes.
+    if (isa<Constant>(V))
+      return false;
+
+    ++SplittableBlocksCount;
+    return ClInstrumentationWithCallThreshold >= 0 &&
+           SplittableBlocksCount > ClInstrumentationWithCallThreshold;
+  }
+
   bool isInPrologue(Instruction &I) {
     return I.getParent() == FnPrologueEnd->getParent() &&
            (&I == FnPrologueEnd || I.comesBefore(FnPrologueEnd));
   }
 
   // Creates a new origin and records the stack trace. In general we can call
   // this function for any origin manipulation we like. However it will cost
   // runtime resources. So use this wisely only if it can provide additional
for (unsigned i = Ofs; i < (Size + kOriginSize - 1) / kOriginSize; ++i) {
       Value *GEP =
           i ? IRB.CreateConstGEP1_32(MS.OriginTy, OriginPtr, i) : OriginPtr;
       IRB.CreateAlignedStore(Origin, GEP, CurrentAlignment);
       CurrentAlignment = kMinOriginAlignment;
     }
   }
 
   void storeOrigin(IRBuilder<> &IRB, Value *Addr, Value *Shadow, Value *Origin,
-                   Value *OriginPtr, Align Alignment, bool AsCall) {
+                   Value *OriginPtr, Align Alignment) {
     const DataLayout &DL = F.getParent()->getDataLayout();
     const Align OriginAlignment = std::max(kMinOriginAlignment, Alignment);
     unsigned StoreSize = DL.getTypeStoreSize(Shadow->getType());
     Value *ConvertedShadow = convertShadowToScalar(Shadow, IRB);
     if (auto *ConstantShadow = dyn_cast<Constant>(ConvertedShadow)) {
       if (!ClCheckConstantShadow || ConstantShadow->isZeroValue()) {
         // Origin is not needed: value is initialized or const shadow is
         // ignored.
         return;
       }
       if (llvm::isKnownNonZero(ConvertedShadow, DL)) {
         // Copy origin as the value is definitely uninitialized.
         paintOrigin(IRB, updateOrigin(Origin, IRB), OriginPtr, StoreSize,
                     OriginAlignment);
         return;
       }
       // Fallback to runtime check, which still can be optimized out later.
     }
 
     unsigned TypeSizeInBits = DL.getTypeSizeInBits(ConvertedShadow->getType());
     unsigned SizeIndex = TypeSizeToSizeIndex(TypeSizeInBits);
-    if (AsCall && SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
+    if (instrumentWithCalls(ConvertedShadow) &&
+        SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
       FunctionCallee Fn = MS.MaybeStoreOriginFn[SizeIndex];
       Value *ConvertedShadow2 =
           IRB.CreateZExt(ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
       CallBase *CB = IRB.CreateCall(
           Fn, {ConvertedShadow2,
                IRB.CreatePointerCast(Addr, IRB.getInt8PtrTy()), Origin});
       CB->addParamAttr(0, Attribute::ZExt);
       CB->addParamAttr(2, Attribute::ZExt);
     } 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);
     }
   }
 
-  void materializeStores(bool InstrumentWithCalls) {
+  void materializeStores() {
     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;
       Type *ShadowTy = Shadow->getType();
       const Align Alignment = SI->getAlign();
       const Align OriginAlignment = std::max(kMinOriginAlignment, Alignment);
       std::tie(ShadowPtr, OriginPtr) =
           getShadowOriginPtr(Addr, IRB, ShadowTy, Alignment, /*isStore*/ true);
 
       StoreInst *NewSI = IRB.CreateAlignedStore(Shadow, ShadowPtr, Alignment);
       LLVM_DEBUG(dbgs() << "  STORE: " << *NewSI << "\n");
       (void)NewSI;
 
       if (SI->isAtomic())
         SI->setOrdering(addReleaseOrdering(SI->getOrdering()));
 
       if (MS.TrackOrigins && !SI->isAtomic())
         storeOrigin(IRB, Addr, Shadow, getOrigin(Val), OriginPtr,
-                    OriginAlignment, InstrumentWithCalls);
+                    OriginAlignment);
     }
   }
 
   // Returns true if Debug Location curresponds to multiple warnings.
   bool shouldDisambiguateWarningLocation(const DebugLoc &DebugLoc) {
     if (MS.TrackOrigins < 2)
       return false;
 
void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     else
       IRB.CreateCall(MS.WarningFn)->setCannotMerge();
     // FIXME: Insert UnreachableInst if !MS.Recover?
     // This may invalidate some of the following checks and needs to be done
     // at the very end.
   }
 
   void materializeOneCheck(IRBuilder<> &IRB, Value *ConvertedShadow,
-                           Value *Origin, bool AsCall) {
+                           Value *Origin) {
     const DataLayout &DL = F.getParent()->getDataLayout();
     unsigned TypeSizeInBits = DL.getTypeSizeInBits(ConvertedShadow->getType());
     unsigned SizeIndex = TypeSizeToSizeIndex(TypeSizeInBits);
-    if (AsCall && SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
+    if (instrumentWithCalls(ConvertedShadow) &&
+        SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
       FunctionCallee Fn = MS.MaybeWarningFn[SizeIndex];
       Value *ConvertedShadow2 =
           IRB.CreateZExt(ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
       CallBase *CB = IRB.CreateCall(
           Fn, {ConvertedShadow2,
                MS.TrackOrigins && Origin ? Origin : (Value *)IRB.getInt32(0)});
       CB->addParamAttr(0, Attribute::ZExt);
       CB->addParamAttr(1, Attribute::ZExt);
     } else {
       Value *Cmp = convertToBool(ConvertedShadow, IRB, "_mscmp");
       Instruction *CheckTerm = SplitBlockAndInsertIfThen(
           Cmp, &*IRB.GetInsertPoint(),
           /* Unreachable */ !MS.Recover, MS.ColdCallWeights);
 
       IRB.SetInsertPoint(CheckTerm);
       insertWarningFn(IRB, Origin);
       LLVM_DEBUG(dbgs() << "  CHECK: " << *Cmp << "\n");
     }
   }
 
   void materializeInstructionChecks(
-      bool InstrumentWithCalls,
       ArrayRef<ShadowOriginAndInsertPoint> InstructionChecks) {
     const DataLayout &DL = F.getParent()->getDataLayout();
     // Disable combining in some cases. TrackOrigins checks each shadow to pick
-    // correct origin. InstrumentWithCalls expects to reduce shadow using API.
-    bool Combine = !InstrumentWithCalls && !MS.TrackOrigins;
+    // correct origin.
+    bool Combine = !MS.TrackOrigins;
     Instruction *Instruction = InstructionChecks.front().OrigIns;
     Value *Shadow = nullptr;
     for (const auto &ShadowData : InstructionChecks) {
       assert(ShadowData.OrigIns == Instruction);
       IRBuilder<> IRB(Instruction);
 
       LLVM_DEBUG(dbgs() << "  SHAD0 : " << *ShadowData.Shadow << "\n");
       Value *ConvertedShadow = convertShadowToScalar(ShadowData.Shadow, IRB);
for (const auto &ShadowData : InstructionChecks) {
           if (!MS.Recover)
             return; // Always fail and stop here, not need to check the rest.
           // Skip entire instruction,
           continue;
         }
         // Fallback to runtime check, which still can be optimized out later.
       }
 
+      // Work around to keep existing tests.
+      // FIXME: update tests and remove.
+      if (ClInstrumentationWithCallThreshold >= 0 &&
+          SplittableBlocksCount >= ClInstrumentationWithCallThreshold) {
+        materializeOneCheck(IRB, ConvertedShadow, ShadowData.Origin);
+        continue;
+      }
+
       if (!Combine) {
-        materializeOneCheck(IRB, ConvertedShadow, ShadowData.Origin,
-                            InstrumentWithCalls);
+        materializeOneCheck(IRB, ConvertedShadow, ShadowData.Origin);
         continue;
       }
 
       Value *BoolShadow = convertToBool(ConvertedShadow, IRB, "_mscmp");
       Shadow = Shadow ? IRB.CreateOr(Shadow, BoolShadow, "_msor") : BoolShadow;
     }
 
     if (Shadow) {
       assert(Combine);
       IRBuilder<> IRB(Instruction);
-      materializeOneCheck(IRB, Shadow, nullptr, false);
+      materializeOneCheck(IRB, Shadow, nullptr);
     }
   }
 
-  void materializeChecks(bool InstrumentWithCalls) {
+  void materializeChecks() {
     llvm::stable_sort(InstrumentationList,
                       [](const ShadowOriginAndInsertPoint &L,
                          const ShadowOriginAndInsertPoint &R) {
                         return L.OrigIns < R.OrigIns;
                       });
 
     for (auto I = InstrumentationList.begin();
          I != InstrumentationList.end();) {
       auto J =
           std::find_if(I + 1, InstrumentationList.end(),
                        [L = I->OrigIns](const ShadowOriginAndInsertPoint &R) {
                          return L != R.OrigIns;
                        });
       // Process all checks of instruction at once.
-      materializeInstructionChecks(InstrumentWithCalls,
-                                   ArrayRef<ShadowOriginAndInsertPoint>(I, J));
+      materializeInstructionChecks(ArrayRef<ShadowOriginAndInsertPoint>(I, J));
       I = J;
     }
 
     LLVM_DEBUG(dbgs() << "DONE:\n" << F);
   }
 
   // Returns the last instruction in the new prologue
   void insertKmsanPrologue(IRBuilder<> &IRB) {
if (InstrumentLifetimeStart) {
         AllocaSet.remove(Item.second);
       }
     }
     // Poison the allocas for which we didn't instrument the corresponding
     // lifetime intrinsics.
     for (AllocaInst *AI : AllocaSet)
       instrumentAlloca(*AI);
 
-    bool InstrumentWithCalls = ClInstrumentationWithCallThreshold >= 0 &&
-                               InstrumentationList.size() + StoreList.size() >
-                                   (unsigned)ClInstrumentationWithCallThreshold;
-
     // Insert shadow value checks.
-    materializeChecks(InstrumentWithCalls);
+    materializeChecks();
 
     // Delayed instrumentation of StoreInst.
     // This may not add new address checks.
-    materializeStores(InstrumentWithCalls);
+    materializeStores();
 
     return true;
   }
 
   /// Compute the shadow type that corresponds to a given Value.
   Type *getShadowTy(Value *V) { return getShadowTy(V->getType()); }
 
   /// Compute the shadow type that corresponds to a given Type.

llvm/test/Instrumentation/MemorySanitizer/with-call-type-size.ll

define <4 x i32> @test65(<4 x i32> %vec, i65 %idx, i32 %x) sanitize_memory {
   %vec1 = insertelement <4 x i32> %vec, i32 %x, i65 %idx
   ret <4 x i32> %vec1
 }
 ; CHECK-LABEL: @test65(
 ; CHECK-NOT:     call void @__msan_maybe_warning_
 ; CHECK:         icmp ne i65 %{{.*}}, 0
 ; CHECK-NOT:     call void @__msan_maybe_warning_
 ; CHECK:         ret <4 x i32>
+
+define <4 x i32> @testUndef(<4 x i32> %vec, i32 %x) sanitize_memory {
+  %vec1 = insertelement <4 x i32> %vec, i32 undef, i32 undef
+  ret <4 x i32> %vec1
+}
+; CHECK-LABEL: @testUndef(
+; CHECK-NOT:     call void @__msan_maybe_warning_
+; CHECK:         ret <4 x i32>