Fix crash in bounds checking

[jgalenson]

Summary

In r337830 I added SCEV checks to enable us to insert fewer bounds checks. Unfortunately, this sometimes crashes when multiple bounds checks are added due to SCEV caching issues. This patch splits the bounds checking pass into two phases, one that computes all the conditions (using SCEV checks) and the other that adds the new instructions.

[eugenis]

lib/Transforms/Instrumentation/BoundsChecking.cpp

162

Move GetTrapBB definition near the end of the function, where it is first used.

184

It's very uncommon to store builders in containers. Replace with an instruction pointer.

TrapInfo can be combined with WorkList, and both loops can be merged in one.

[Ka-Ka]

I found another reproducer for the problem. This one is slightly different where a different assert blow, but I guess its only a different symptom for the same fault. As this one is slightly smaller you might want to use this one as a testcase instead (or use both)? Run with "opt -bounds-checking -S".

@d = dso_local local_unnamed_addr global [1 x i16] zeroinitializer, align 1
@e = dso_local local_unnamed_addr global [1 x i16] zeroinitializer, align 1

define dso_local void @test2() local_unnamed_addr {
entry:

br label %while.cond1.preheader

while.cond1.preheader:

%0 = phi i16 [ undef, %entry ], [ %inc, %while.end ]
%1 = load i16, i16* undef, align 1
br label %while.end

while.end:

%inc = add nsw i16 %0, 1
%arrayidx = getelementptr inbounds [1 x i16], [1 x i16]* @e, i16 0, i16 %0
%2 = load i16, i16* %arrayidx, align 1
br i1 false, label %while.end6, label %while.cond1.preheader

while.end6:

ret void

}

[jgalenson]

In D49946#1179935, @Ka-Ka wrote:

I found another reproducer for the problem. This one is slightly different where a different assert blow, but I guess its only a different symptom for the same fault. As this one is slightly smaller you might want to use this one as a testcase instead (or use both)?

Thanks! I'll use both it and the first one you gave me. Luckily, it doesn't crash with this patch.

[jgalenson]

lib/Transforms/Instrumentation/BoundsChecking.cpp

184

How did you want me to combine TrapInfo with WorkList? To insert the traps, I need both the original instruction (to create the builder) and the newly-created Or. I could change Worklist to hold a tuple with an optional Or, but that seems more complicated than using two lists like this.

[eugenis]

lib/Transforms/Instrumentation/BoundsChecking.cpp

159

The loop over worklist can be replaced with a loop over instructions(F).

It would work because we no longer create new basic blocks in the loop, and new instructions are added before the current iterator position, so they will not be observed in the following iterations.

[jgalenson]

lib/Transforms/Instrumentation/BoundsChecking.cpp

159

Ah, you meant combine this loop with the previous one, not the new one I added below? Sorry, I misunderstood. That makes sense.

[jgalenson]

Friendly ping. This also fixes https://bugs.llvm.org/show_bug.cgi?id=38438.

[eugenis]

LGTM, thanks!

Sorry for the delay, I did not notice that you've updated the patch with your last comment.

[jgalenson]

No problem, and thanks for the review!

lib/Transforms/Instrumentation/BoundsChecking.cpp

static cl::opt<bool> SingleTrapBB("bounds-checking-single-trap",
                                   cl::desc("Use one trap block per function"));
 
 STATISTIC(ChecksAdded, "Bounds checks added");
 STATISTIC(ChecksSkipped, "Bounds checks skipped");
 STATISTIC(ChecksUnable, "Bounds checks unable to add");
 
 using BuilderTy = IRBuilder<TargetFolder>;
 
-/// Adds run-time bounds checks to memory accessing instructions.
+/// Gets the conditions under which memory accessing instructions will overflow.
 ///
 /// \p Ptr is the pointer that will be read/written, and \p InstVal is either
 /// the result from the load or the value being stored. It is used to determine
 /// the size of memory block that is touched.
 ///
-/// \p GetTrapBB is a callable that returns the trap BB to use on failure.
-///
-/// Returns true if any change was made to the IR, false otherwise.
-template <typename GetTrapBBT>
-static bool instrumentMemAccess(Value *Ptr, Value *InstVal,
+/// Returns the condition under which the access will overflow.
+static Value *getBoundsCheckCond(Value *Ptr, Value *InstVal,
                                  const DataLayout &DL, TargetLibraryInfo &TLI,
                                  ObjectSizeOffsetEvaluator &ObjSizeEval,
-                                BuilderTy &IRB, GetTrapBBT GetTrapBB,
-                                ScalarEvolution &SE) {
+                                 BuilderTy &IRB, ScalarEvolution &SE) {
   uint64_t NeededSize = DL.getTypeStoreSize(InstVal->getType());
   LLVM_DEBUG(dbgs() << "Instrument " << *Ptr << " for " << Twine(NeededSize)
                     << " bytes\n");
 
   SizeOffsetEvalType SizeOffset = ObjSizeEval.compute(Ptr);
 
   if (!ObjSizeEval.bothKnown(SizeOffset)) {
     ++ChecksUnable;
-    return false;
+    return nullptr;
   }
 
   Value *Size   = SizeOffset.first;
   Value *Offset = SizeOffset.second;
   ConstantInt *SizeCI = dyn_cast<ConstantInt>(Size);
 
   Type *IntTy = DL.getIntPtrType(Ptr->getType());
   Value *NeededSizeVal = ConstantInt::get(IntTy, NeededSize);
Value *Cmp3 = SizeRange.sub(OffsetRange)
                     : IRB.CreateICmpULT(ObjSize, NeededSizeVal);
   Value *Or = IRB.CreateOr(Cmp2, Cmp3);
   if ((!SizeCI || SizeCI->getValue().slt(0)) &&
       !SizeRange.getSignedMin().isNonNegative()) {
     Value *Cmp1 = IRB.CreateICmpSLT(Offset, ConstantInt::get(IntTy, 0));
     Or = IRB.CreateOr(Cmp1, Or);
   }
 
+  return Or;
+}
+
+/// Adds run-time bounds checks to memory accessing instructions.
+///
+/// \p Or is the condition that should guard the trap.
+///
+/// \p GetTrapBB is a callable that returns the trap BB to use on failure.
+template <typename GetTrapBBT>
+static void insertBoundsCheck(Value *Or, BuilderTy IRB, GetTrapBBT GetTrapBB) {
   // check if the comparison is always false
   ConstantInt *C = dyn_cast_or_null<ConstantInt>(Or);
   if (C) {
     ++ChecksSkipped;
     // If non-zero, nothing to do.
     if (!C->getZExtValue())
-      return true;
+      return;
   }
   ++ChecksAdded;
 
   BasicBlock::iterator SplitI = IRB.GetInsertPoint();
   BasicBlock *OldBB = SplitI->getParent();
   BasicBlock *Cont = OldBB->splitBasicBlock(SplitI);
   OldBB->getTerminator()->eraseFromParent();
 
   if (C) {
     // If we have a constant zero, unconditionally branch.
     // FIXME: We should really handle this differently to bypass the splitting
     // the block.
     BranchInst::Create(GetTrapBB(IRB), OldBB);
-    return true;
+    return;
   }
 
   // Create the conditional branch.
   BranchInst::Create(GetTrapBB(IRB), Cont, Or, OldBB);
-  return true;
 }
 
 static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
                               ScalarEvolution &SE) {
   const DataLayout &DL = F.getParent()->getDataLayout();
   ObjectSizeOffsetEvaluator ObjSizeEval(DL, &TLI, F.getContext(),
                                            /*RoundToAlign=*/true);
 
   // check HANDLE_MEMORY_INST in include/llvm/Instruction.def for memory
   // touching instructions
   std::vector<Instruction *> WorkList;
   for (Instruction &I : instructions(F)) {
     if (isa<LoadInst>(I) || isa<StoreInst>(I) || isa<AtomicCmpXchgInst>(I) ||
         isa<AtomicRMWInst>(I))
         WorkList.push_back(&I);
   }
 
   // Create a trapping basic block on demand using a callback. Depending on
   // flags, this will either create a single block for the entire function or
   // will create a fresh block every time it is called.
   BasicBlock *TrapBB = nullptr;
   auto GetTrapBB = [&TrapBB](BuilderTy &IRB) {
     if (TrapBB && SingleTrapBB)
       return TrapBB;
 
     Function *Fn = IRB.GetInsertBlock()->getParent();
     // FIXME: This debug location doesn't make a lot of sense in the
     // `SingleTrapBB` case.
     auto DebugLoc = IRB.getCurrentDebugLocation();
     IRBuilder<>::InsertPointGuard Guard(IRB);
     TrapBB = BasicBlock::Create(Fn->getContext(), "trap", Fn);
     IRB.SetInsertPoint(TrapBB);
 
     auto *F = Intrinsic::getDeclaration(Fn->getParent(), Intrinsic::trap);
     CallInst *TrapCall = IRB.CreateCall(F, {});
     TrapCall->setDoesNotReturn();
     TrapCall->setDoesNotThrow();
     TrapCall->setDebugLoc(DebugLoc);
     IRB.CreateUnreachable();
 
     return TrapBB;
   };
 
-  bool MadeChange = false;
+  SmallVector<std::pair<Value *, BuilderTy>, 4> TrapInfo;
   for (Instruction *Inst : WorkList) {
+    Value *Or = nullptr;
     BuilderTy IRB(Inst->getParent(), BasicBlock::iterator(Inst), TargetFolder(DL));
     if (LoadInst *LI = dyn_cast<LoadInst>(Inst)) {
-      MadeChange |= instrumentMemAccess(LI->getPointerOperand(), LI, DL, TLI,
-                                        ObjSizeEval, IRB, GetTrapBB, SE);
+      Or = getBoundsCheckCond(LI->getPointerOperand(), LI, DL, TLI,
+                              ObjSizeEval, IRB, SE);
     } else if (StoreInst *SI = dyn_cast<StoreInst>(Inst)) {
-      MadeChange |=
-          instrumentMemAccess(SI->getPointerOperand(), SI->getValueOperand(),
-                              DL, TLI, ObjSizeEval, IRB, GetTrapBB, SE);
+      Or = getBoundsCheckCond(SI->getPointerOperand(), SI->getValueOperand(),
+                              DL, TLI, ObjSizeEval, IRB, SE);
     } else if (AtomicCmpXchgInst *AI = dyn_cast<AtomicCmpXchgInst>(Inst)) {
-      MadeChange |=
-          instrumentMemAccess(AI->getPointerOperand(), AI->getCompareOperand(),
-                              DL, TLI, ObjSizeEval, IRB, GetTrapBB, SE);
+      Or = getBoundsCheckCond(AI->getPointerOperand(), AI->getCompareOperand(),
+                              DL, TLI, ObjSizeEval, IRB, SE);
     } else if (AtomicRMWInst *AI = dyn_cast<AtomicRMWInst>(Inst)) {
-      MadeChange |=
-          instrumentMemAccess(AI->getPointerOperand(), AI->getValOperand(), DL,
-                              TLI, ObjSizeEval, IRB, GetTrapBB, SE);
+      Or = getBoundsCheckCond(AI->getPointerOperand(), AI->getValOperand(), DL,
+                              TLI, ObjSizeEval, IRB, SE);
     } else {
       llvm_unreachable("unknown Instruction type");
     }
+    if (Or)
+      TrapInfo.push_back(std::make_pair(Or, IRB));
   }
-  return MadeChange;
+
+  for (const auto &Entry : TrapInfo)
+    insertBoundsCheck(Entry.first, Entry.second, GetTrapBB);
+
+  return !TrapInfo.empty();
 }
 
 PreservedAnalyses BoundsCheckingPass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
 
   if (!addBoundsChecking(F, TLI, SE))
     return PreservedAnalyses::all();

test/Instrumentation/BoundsChecking/many-traps-2.ll

+; RUN: opt < %s -bounds-checking -S | FileCheck %s
+@array = internal global [1819 x i16] zeroinitializer, section ".bss,bss"
+@offsets = external dso_local global [10 x i16]
+
+; CHECK-LABEL: @test
+define dso_local void @test() {
+bb1:
+  br label %bb19
+
+bb20:
+  %_tmp819 = load i16, i16* null
+; CHECK: br {{.*}} %trap
+  %_tmp820 = sub nsw i16 9, %_tmp819
+  %_tmp821 = sext i16 %_tmp820 to i64
+  %_tmp822 = getelementptr [10 x i16], [10 x i16]* @offsets, i16 0, i64 %_tmp821
+  %_tmp823 = load i16, i16* %_tmp822
+  br label %bb33
+
+bb34:
+  %_tmp907 = zext i16 %i__7.107.0 to i64
+  %_tmp908 = getelementptr [1819 x i16], [1819 x i16]* @array, i16 0, i64 %_tmp907
+  store i16 0, i16* %_tmp908
+; CHECK: br {{.*}} %trap
+  %_tmp910 = add i16 %i__7.107.0, 1
+  br label %bb33
+
+bb33:
+  %i__7.107.0 = phi i16 [ undef, %bb20 ], [ %_tmp910, %bb34 ]
+  %_tmp913 = add i16 %_tmp823, 191
+  %_tmp914 = icmp ult i16 %i__7.107.0, %_tmp913
+  br i1 %_tmp914, label %bb34, label %bb19
+
+bb19:
+  %_tmp976 = icmp slt i16 0, 10
+  br i1 %_tmp976, label %bb20, label %bb39
+
+bb39:
+  ret void
+}