[BoundsChecking] Add support for runtime handlers

[vitalybuka]

This is a step forward to have reporting consistent with other UBSAN checks.

Runtime and clang parts are here #120515.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Vitaly Buka (vitalybuka)

Changes

This is a step forward to have reporting consistent with other UBSAN checks.

---

Full diff: https://github.com/llvm/llvm-project/pull/120513.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp (+81-21)
• (modified) llvm/test/Instrumentation/BoundsChecking/runtimes.ll (+6-6)

diff --git a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
index c86d967716a5a0..d84f3ae73b062e 100644
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -8,6 +8,7 @@
 
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
 #include "llvm/Analysis/ScalarEvolution.h"
@@ -104,6 +105,29 @@ static Value *getBoundsCheckCond(Value *Ptr, Value *InstVal,
   return Or;
 }
 
+static CallInst *InsertTrap(BuilderTy &IRB) {
+  if (!DebugTrapBB)
+    return IRB.CreateIntrinsic(Intrinsic::trap, {}, {});
+  return IRB.CreateIntrinsic(
+      Intrinsic::ubsantrap, {},
+      ConstantInt::get(IRB.getInt8Ty(),
+                       IRB.GetInsertBlock()->getParent()->size()));
+}
+
+static CallInst *InsertCall(BuilderTy &IRB, bool MayReturn, StringRef Name) {
+  Function *Fn = IRB.GetInsertBlock()->getParent();
+  LLVMContext &Ctx = Fn->getContext();
+  llvm::AttrBuilder B(Ctx);
+  B.addAttribute(llvm::Attribute::NoUnwind);
+  if (!MayReturn)
+    B.addAttribute(llvm::Attribute::NoReturn);
+  FunctionCallee Callee = Fn->getParent()->getOrInsertFunction(
+      Name,
+      llvm::AttributeList::get(Ctx, llvm::AttributeList::FunctionIndex, B),
+      Type::getVoidTy(Ctx));
+  return IRB.CreateCall(Callee);
+}
+
 /// Adds run-time bounds checks to memory accessing instructions.
 ///
 /// \p Or is the condition that should guard the trap.
@@ -126,20 +150,53 @@ static void insertBoundsCheck(Value *Or, BuilderTy &IRB, GetTrapBBT GetTrapBB) {
   BasicBlock *Cont = OldBB->splitBasicBlock(SplitI);
   OldBB->getTerminator()->eraseFromParent();
 
+  BasicBlock *TrapBB = GetTrapBB(IRB, Cont);
+
   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);
+    BranchInst::Create(TrapBB, OldBB);
     return;
   }
 
   // Create the conditional branch.
-  BranchInst::Create(GetTrapBB(IRB), Cont, Or, OldBB);
+  BranchInst::Create(TrapBB, Cont, Or, OldBB);
 }
 
+struct ReportingOpts {
+  bool MayReturn = false;
+  bool UseTrap = false;
+  bool MinRuntime = false;
+  StringRef Name;
+
+  ReportingOpts(BoundsCheckingPass::ReportingMode Mode) {
+    switch (Mode) {
+    case BoundsCheckingPass::ReportingMode::Trap:
+      UseTrap = true;
+      break;
+    case BoundsCheckingPass::ReportingMode::MinRuntime:
+      Name = "__ubsan_handle_local_out_of_bounds_minimal";
+      MinRuntime = true;
+      MayReturn = true;
+      break;
+    case BoundsCheckingPass::ReportingMode::MinRuntimeAbort:
+      Name = "__ubsan_handle_local_out_of_bounds_minimal_abort";
+      MinRuntime = true;
+      break;
+    case BoundsCheckingPass::ReportingMode::FullRuntime:
+      Name = "__ubsan_handle_local_out_of_bounds";
+      MayReturn = true;
+      break;
+    case BoundsCheckingPass::ReportingMode::FullRuntimeAbort:
+      Name = "__ubsan_handle_local_out_of_bounds_abort";
+      break;
+    }
+  }
+};
+
 static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
-                              ScalarEvolution &SE) {
+                              ScalarEvolution &SE, const ReportingOpts &Opts) {
   if (F.hasFnAttribute(Attribute::NoSanitizeBounds))
     return false;
 
@@ -180,37 +237,40 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
   // 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) {
+  BasicBlock *ReuseTrapBB = nullptr;
+  auto GetTrapBB = [&ReuseTrapBB, &Opts](BuilderTy &IRB, BasicBlock *Cont) {
     Function *Fn = IRB.GetInsertBlock()->getParent();
     auto DebugLoc = IRB.getCurrentDebugLocation();
     IRBuilder<>::InsertPointGuard Guard(IRB);
 
-    if (TrapBB && SingleTrapBB && !DebugTrapBB)
-      return TrapBB;
+    // 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.
+    if (ReuseTrapBB)
+      return ReuseTrapBB;
 
-    TrapBB = BasicBlock::Create(Fn->getContext(), "trap", Fn);
+    BasicBlock *TrapBB = BasicBlock::Create(Fn->getContext(), "trap", Fn);
     IRB.SetInsertPoint(TrapBB);
 
-    Intrinsic::ID IntrID = DebugTrapBB ? Intrinsic::ubsantrap : Intrinsic::trap;
+    CallInst *TrapCall = Opts.UseTrap
+                             ? InsertTrap(IRB)
+                             : InsertCall(IRB, Opts.MayReturn, Opts.Name);
 
-    CallInst *TrapCall;
-    if (DebugTrapBB) {
-      TrapCall = IRB.CreateIntrinsic(
-          IntrID, {}, ConstantInt::get(IRB.getInt8Ty(), Fn->size()));
+    TrapCall->setDoesNotThrow();
+    TrapCall->setDebugLoc(DebugLoc);
+    if (Opts.MayReturn) {
+      IRB.CreateBr(Cont);
     } else {
-      TrapCall = IRB.CreateIntrinsic(IntrID, {}, {});
+      TrapCall->setDoesNotReturn();
+      IRB.CreateUnreachable();
     }
 
-    TrapCall->setDoesNotReturn();
-    TrapCall->setDoesNotThrow();
-    TrapCall->setDebugLoc(DebugLoc);
-    IRB.CreateUnreachable();
+    if (!Opts.MayReturn && SingleTrapBB && !DebugTrapBB)
+      ReuseTrapBB = TrapBB;
 
     return TrapBB;
   };
 
-  // Add the checks.
   for (const auto &Entry : TrapInfo) {
     Instruction *Inst = Entry.first;
     BuilderTy IRB(Inst->getParent(), BasicBlock::iterator(Inst), TargetFolder(DL));
@@ -224,7 +284,7 @@ PreservedAnalyses BoundsCheckingPass::run(Function &F, FunctionAnalysisManager &
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
 
-  if (!addBoundsChecking(F, TLI, SE))
+  if (!addBoundsChecking(F, TLI, SE, ReportingOpts(Mode)))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
@@ -251,4 +311,4 @@ void BoundsCheckingPass::printPipeline(
     OS << "<rt-abort>";
     break;
   }
-}
\ No newline at end of file
+}
diff --git a/llvm/test/Instrumentation/BoundsChecking/runtimes.ll b/llvm/test/Instrumentation/BoundsChecking/runtimes.ll
index fd27694c155d2b..357f92aca85c08 100644
--- a/llvm/test/Instrumentation/BoundsChecking/runtimes.ll
+++ b/llvm/test/Instrumentation/BoundsChecking/runtimes.ll
@@ -38,8 +38,8 @@ define void @f1(i64 %x) nounwind {
 ; RT-NEXT:    [[TMP8:%.*]] = load i128, ptr [[TMP2]], align 4
 ; RT-NEXT:    ret void
 ; RT:       [[TRAP]]:
-; RT-NEXT:    call void @llvm.trap() #[[ATTR2:[0-9]+]]
-; RT-NEXT:    unreachable
+; RT-NEXT:    call void @__ubsan_handle_local_out_of_bounds() #[[ATTR0]]
+; RT-NEXT:    br label %[[BB7]]
 ;
 ; RTABORT-LABEL: define void @f1(
 ; RTABORT-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
@@ -54,7 +54,7 @@ define void @f1(i64 %x) nounwind {
 ; RTABORT-NEXT:    [[TMP8:%.*]] = load i128, ptr [[TMP2]], align 4
 ; RTABORT-NEXT:    ret void
 ; RTABORT:       [[TRAP]]:
-; RTABORT-NEXT:    call void @llvm.trap() #[[ATTR2:[0-9]+]]
+; RTABORT-NEXT:    call void @__ubsan_handle_local_out_of_bounds_abort() #[[ATTR1:[0-9]+]]
 ; RTABORT-NEXT:    unreachable
 ;
 ; MINRT-LABEL: define void @f1(
@@ -70,8 +70,8 @@ define void @f1(i64 %x) nounwind {
 ; MINRT-NEXT:    [[TMP8:%.*]] = load i128, ptr [[TMP2]], align 4
 ; MINRT-NEXT:    ret void
 ; MINRT:       [[TRAP]]:
-; MINRT-NEXT:    call void @llvm.trap() #[[ATTR2:[0-9]+]]
-; MINRT-NEXT:    unreachable
+; MINRT-NEXT:    call void @__ubsan_handle_local_out_of_bounds_minimal() #[[ATTR0]]
+; MINRT-NEXT:    br label %[[BB7]]
 ;
 ; MINRTABORT-LABEL: define void @f1(
 ; MINRTABORT-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
@@ -86,7 +86,7 @@ define void @f1(i64 %x) nounwind {
 ; MINRTABORT-NEXT:    [[TMP8:%.*]] = load i128, ptr [[TMP2]], align 4
 ; MINRTABORT-NEXT:    ret void
 ; MINRTABORT:       [[TRAP]]:
-; MINRTABORT-NEXT:    call void @llvm.trap() #[[ATTR2:[0-9]+]]
+; MINRTABORT-NEXT:    call void @__ubsan_handle_local_out_of_bounds_minimal_abort() #[[ATTR1:[0-9]+]]
 ; MINRTABORT-NEXT:    unreachable
 ;
   %1 = alloca i128, i64 %x

[vitalybuka]

Can't land #120515 without this one

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.