[UBSAN] Rename remove-traps to lower-allow-check

clang/lib/CodeGen/BackendUtil.cpp

@@ -73,18 +73,17 @@
 #include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/InstrProfiling.h"
 #include "llvm/Transforms/Instrumentation/KCFI.h"
+#include "llvm/Transforms/Instrumentation/LowerBuiltinHotPass.h"
 #include "llvm/Transforms/Instrumentation/MemProfiler.h"
 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
-#include "llvm/Transforms/Instrumentation/RemoveTrapsPass.h"
 #include "llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h"
 #include "llvm/Transforms/Instrumentation/SanitizerCoverage.h"
 #include "llvm/Transforms/Instrumentation/ThreadSanitizer.h"
 #include "llvm/Transforms/ObjCARC.h"
 #include "llvm/Transforms/Scalar/EarlyCSE.h"
 #include "llvm/Transforms/Scalar/GVN.h"
 #include "llvm/Transforms/Scalar/JumpThreading.h"
-#include "llvm/Transforms/Scalar/SimplifyCFG.h"
 #include "llvm/Transforms/Utils/Debugify.h"
 #include "llvm/Transforms/Utils/EntryExitInstrumenter.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
@@ -100,9 +99,10 @@ using namespace llvm;
 namespace llvm {
 extern cl::opt<bool> PrintPipelinePasses;
 
-cl::opt<bool> ClRemoveTraps("clang-remove-traps", cl::Optional,
-                            cl::desc("Insert remove-traps pass."),
-                            cl::init(false));
+static cl::opt<bool>
+    ClLowerBuiltinHot("clang-lower-builtin-hot", cl::Optional,
+                      cl::desc("Insert lower-builtin-hot pass."),
+                      cl::init(false));
 
 // Experiment to move sanitizers earlier.
 static cl::opt<bool> ClSanitizeOnOptimizerEarlyEP(
@@ -751,18 +751,13 @@ static void addSanitizers(const Triple &TargetTriple,
     PB.registerOptimizerLastEPCallback(SanitizersCallback);
   }
 
-  if (ClRemoveTraps) {
+  if (ClLowerBuiltinHot) {
     // We can optimize after inliner, and PGO profile matching. The hook below
     // is called at the end `buildFunctionSimplificationPipeline`, which called
     // from `buildInlinerPipeline`, which called after profile matching.
     PB.registerScalarOptimizerLateEPCallback(
         [](FunctionPassManager &FPM, OptimizationLevel Level) {
-          // RemoveTrapsPass expects trap blocks preceded by conditional
-          // branches, which usually is not the case without SimplifyCFG.
-          // TODO: Remove `SimplifyCFGPass` after switching to dedicated
-          // intrinsic.
-          FPM.addPass(SimplifyCFGPass());
-          FPM.addPass(RemoveTrapsPass());
+          FPM.addPass(LowerBuiltinHotPass());
         });
   }
 }

clang/lib/CodeGen/CGExpr.cpp

@@ -56,7 +56,14 @@ using namespace CodeGen;
 // Experiment to make sanitizers easier to debug
 static llvm::cl::opt<bool> ClSanitizeDebugDeoptimization(
     "ubsan-unique-traps", llvm::cl::Optional,
-    llvm::cl::desc("Deoptimize traps for UBSAN so there is 1 trap per check"),
+    llvm::cl::desc("Deoptimize traps for UBSAN so there is 1 trap per check."),
+    llvm::cl::init(false));
+
+// TODO: Introduce frontend options to enabled per sanitizers, similar to
+// `fsanitize-trap`.
+static llvm::cl::opt<bool> ClSanitizeExpHot(
+    "ubsan-exp-hot", llvm::cl::Optional,
+    llvm::cl::desc("Pass UBSAN checks if `llvm.experimental.hot()` is true."),
     llvm::cl::init(false));
 
 //===--------------------------------------------------------------------===//
@@ -3805,6 +3812,12 @@ void CodeGenFunction::EmitTrapCheck(llvm::Value *Checked,
                                     SanitizerHandler CheckHandlerID) {
   llvm::BasicBlock *Cont = createBasicBlock("cont");
 
+  if (ClSanitizeExpHot) {
+    Checked =
+        Builder.CreateOr(Checked, Builder.CreateCall(CGM.getIntrinsic(
+                                      llvm::Intrinsic::experimental_hot)));
+  }
+
   // If we're optimizing, collapse all calls to trap down to just one per
   // check-type per function to save on code size.
   if ((int)TrapBBs.size() <= CheckHandlerID)

clang/test/CodeGen/lower-builtin-hot.c

@@ -0,0 +1,32 @@
+// RUN: %clang_cc1 -O1 %s -o - -emit-llvm -fsanitize=signed-integer-overflow -fsanitize-trap=signed-integer-overflow -mllvm -ubsan-exp-hot | FileCheck %s
+// RUN: %clang_cc1 -O1 %s -o - -emit-llvm -fsanitize=signed-integer-overflow -fsanitize-trap=signed-integer-overflow -mllvm -ubsan-exp-hot -mllvm -clang-lower-builtin-hot -mllvm -lower-builtin-hot-random-rate=1 %s -o - | FileCheck %s --check-prefixes=REMOVE
+
+#include <stdbool.h>
+
+int test(int x) {
+  return x + 123;
+}
+
+// CHECK-LABEL: define {{.*}}i32 @test(
+// CHECK: call { i32, i1 } @llvm.sadd.with.overflow.i32(
+// CHECK: trap:
+// CHECK-NEXT: call void @llvm.ubsantrap(i8 0)
+// CHECK-NEXT: unreachable
+
+// REMOVE-LABEL: define {{.*}}i32 @test(
+// REMOVE: add i32 %x, 123
+// REMOVE-NEXT: ret i32
+
+
+bool experimental_hot() __asm("llvm.experimental.hot");
+
+bool test_asm() {
+  return experimental_hot();
+}
+
+// CHECK-LABEL: define {{.*}}i1 @test_asm(
+// CHECK: [[R:%.*]] = tail call zeroext i1 @llvm.experimental.hot()
+// CHECK: ret i1 [[R]]
+
+// REMOVE-LABEL: define {{.*}}i1 @test_asm(
+// REMOVE: ret i1 true

llvm/docs/LangRef.rst

@@ -27639,6 +27639,54 @@ constant `true`. However it is always correct to replace
 it with any other `i1` value. Any pass can
 freely do it if it can benefit from non-default lowering.
 
+'``llvm.experimental.hot``' Intrinsic
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+::
+
+      declare i1 @llvm.experimental.hot()
+
+Overview:
+"""""""""
+
+This intrinsic returns ``true`` iff it's known that containing basic block is
+hot in profile.
+
+When used with profile based optimization allows to change program behaviour
+depending on the code hotness.
+
+Arguments:
+""""""""""
+
+None.
+
+Semantics:
+""""""""""
+
+The intrinsic ``@llvm.experimental.hot()`` returns either ``true`` or ``false``,
+depending on profile used. Expresion is evaluated as ``true`` iff profile and
+summary are available and profile counter for the block reach hotness threshold.
+For each evaluation of a call to this intrinsic, the program must be valid and
+correct both if it returns ``true`` and if it returns ``false``.
+
+When used in a branch condition, it allows us to choose between
+two alternative correct solutions for the same problem, like
+in example below:
+
+.. code-block:: text
+
+    %cond = call i1 @llvm.experimental.hot()
+    br i1 %cond, label %fast_path, label %slow_path
+
+  fast_path:
+    ; Omit diagnostics.
+
+  slow_path:
+    ; Additional diagnostics.
+
 
 '``llvm.load.relative``' Intrinsic
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

llvm/include/llvm/IR/Intrinsics.td

@@ -1722,6 +1722,11 @@ def int_debugtrap : Intrinsic<[]>,
 def int_ubsantrap : Intrinsic<[], [llvm_i8_ty],
                               [IntrNoReturn, IntrCold, ImmArg<ArgIndex<0>>]>;
 
+// Return true if profile counter for containing block is hot.
+def int_experimental_hot : Intrinsic<[llvm_i1_ty], [],
+                                     [IntrInaccessibleMemOnly, IntrWriteMem,
+                                      IntrWillReturn, NoUndef<RetIndex>]>;
+
 // Support for dynamic deoptimization (or de-specialization)
 def int_experimental_deoptimize : Intrinsic<[llvm_any_ty], [llvm_vararg_ty],
                                             [Throws]>;

llvm/include/llvm/Transforms/Instrumentation/RemoveTrapsPass.h → llvm/include/llvm/Transforms/Instrumentation/LowerBuiltinHotPass.h

@@ -1,4 +1,4 @@
-//===- RemoveTrapsPass.h ----------------------------------------*- C++ -*-===//
+//===- LowerBuiltinHotPass.h ------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -11,8 +11,8 @@
 ///
 //===----------------------------------------------------------------------===//
 
-#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_UBSANOPTIMIZATIONPASS_H
-#define LLVM_TRANSFORMS_INSTRUMENTATION_UBSANOPTIMIZATIONPASS_H
+#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_LOWERBUILTINHOTPASS_H
+#define LLVM_TRANSFORMS_INSTRUMENTATION_LOWERBUILTINHOTPASS_H
 
 #include "llvm/IR/Function.h"
 #include "llvm/IR/PassManager.h"
@@ -22,7 +22,7 @@ namespace llvm {
 
 // This pass is responsible for removing optional traps, like llvm.ubsantrap
 // from the hot code.
-class RemoveTrapsPass : public PassInfoMixin<RemoveTrapsPass> {
+class LowerBuiltinHotPass : public PassInfoMixin<LowerBuiltinHotPass> {
 public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -7276,6 +7276,12 @@ void SelectionDAGBuilder::visitIntrinsicCall(const CallInst &I,
     setValue(&I, getValue(I.getArgOperand(0)));
     return;
 
+  case Intrinsic::experimental_hot:
+    // Default lowering to false. It's intended to be lowered as soon as profile
+    // is avalible to unblock other optimizations.
+    setValue(&I, DAG.getConstant(0, sdl, MVT::i1));
+    return;
+
   case Intrinsic::ubsantrap:
   case Intrinsic::debugtrap:
   case Intrinsic::trap: {

llvm/lib/Passes/PassBuilder.cpp

@@ -172,12 +172,12 @@
 #include "llvm/Transforms/Instrumentation/InstrOrderFile.h"
 #include "llvm/Transforms/Instrumentation/InstrProfiling.h"
 #include "llvm/Transforms/Instrumentation/KCFI.h"
+#include "llvm/Transforms/Instrumentation/LowerBuiltinHotPass.h"
 #include "llvm/Transforms/Instrumentation/MemProfiler.h"
 #include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
 #include "llvm/Transforms/Instrumentation/PGOForceFunctionAttrs.h"
 #include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
 #include "llvm/Transforms/Instrumentation/PoisonChecking.h"
-#include "llvm/Transforms/Instrumentation/RemoveTrapsPass.h"
 #include "llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h"
 #include "llvm/Transforms/Instrumentation/SanitizerCoverage.h"
 #include "llvm/Transforms/Instrumentation/ThreadSanitizer.h"

llvm/lib/Passes/PassRegistry.def

@@ -422,7 +422,7 @@ FUNCTION_PASS("print<uniformity>", UniformityInfoPrinterPass(dbgs()))
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass())
 FUNCTION_PASS("reg2mem", RegToMemPass())
-FUNCTION_PASS("remove-traps", RemoveTrapsPass())
+FUNCTION_PASS("lower-builtin-hot", LowerBuiltinHotPass())
 FUNCTION_PASS("safe-stack", SafeStackPass(TM))
 FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass())
 FUNCTION_PASS("scalarizer", ScalarizerPass())

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -2827,6 +2827,14 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
         }))
       return nullptr;
     break;
+  case Intrinsic::experimental_hot: {
+    // The intrinsic declaration includes sideeffects to avoid it moved. This
+    // prevents removing even if the intrinsic is unused. We should remove
+    // unused ones to enabled other optimizations.
+    if (CI.use_empty())
+      return eraseInstFromFunction(CI);
+    break;
+  }
   case Intrinsic::assume: {
     Value *IIOperand = II->getArgOperand(0);
     SmallVector<OperandBundleDef, 4> OpBundles;

llvm/lib/Transforms/Instrumentation/CMakeLists.txt

@@ -13,11 +13,11 @@ add_llvm_component_library(LLVMInstrumentation
   InstrOrderFile.cpp
   InstrProfiling.cpp
   KCFI.cpp
+  LowerBuiltinHotPass.cpp
   PGOForceFunctionAttrs.cpp
   PGOInstrumentation.cpp
   PGOMemOPSizeOpt.cpp
   PoisonChecking.cpp
-  RemoveTrapsPass.cpp
   SanitizerCoverage.cpp
   SanitizerBinaryMetadata.cpp
   ValueProfileCollector.cpp

llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp → llvm/lib/Transforms/Instrumentation/LowerBuiltinHotPass.cpp

@@ -1,16 +1,17 @@
-//===- RemoveTrapsPass.cpp --------------------------------------*- C++ -*-===//
+//===- LowerBuiltinHotPass.cpp ----------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/Transforms/Instrumentation/RemoveTrapsPass.h"
+#include "llvm/Transforms/Instrumentation/LowerBuiltinHotPass.h"
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/IR/Constant.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
@@ -20,15 +21,15 @@
 
 using namespace llvm;
 
-#define DEBUG_TYPE "remove-traps"
+#define DEBUG_TYPE "lower-builtin-hot"
 
 static cl::opt<int> HotPercentileCutoff(
-    "remove-traps-percentile-cutoff-hot", cl::init(0),
+    "lower-builtin-hot-percentile-cutoff", cl::init(0),
     cl::desc("Alternative hot percentile cuttoff. By default "
              "`-profile-summary-cutoff-hot` is used."));
 
 static cl::opt<float>
-    RandomRate("remove-traps-random-rate", cl::init(0.0),
+    RandomRate("lower-builtin-hot-random-rate", cl::init(0.0),
                cl::desc("Probability value in the range [0.0, 1.0] of "
                         "unconditional pseudo-random checks removal."));
 
@@ -37,7 +38,7 @@ STATISTIC(NumChecksRemoved, "Number of removed checks");
 
 static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
                              const ProfileSummaryInfo *PSI) {
-  SmallVector<IntrinsicInst *, 16> Remove;
+  SmallVector<std::pair<IntrinsicInst *, Value *>, 16> ReplaceWithValue;
   std::unique_ptr<RandomNumberGenerator> Rng;
 
   auto ShouldRemove = [&](bool IsHot) {
@@ -56,26 +57,27 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
         continue;
       auto ID = II->getIntrinsicID();
       switch (ID) {
-      case Intrinsic::ubsantrap: {
+      case Intrinsic::experimental_hot: {
         ++NumChecksTotal;
 
         bool IsHot = false;
         if (PSI) {
-          uint64_t Count = 0;
-          for (const auto *PR : predecessors(&BB))
-            Count += BFI.getBlockProfileCount(PR).value_or(0);
-
+          uint64_t Count = BFI.getBlockProfileCount(&BB).value_or(0);
           IsHot =
               HotPercentileCutoff.getNumOccurrences()
                   ? (HotPercentileCutoff > 0 &&
                      PSI->isHotCountNthPercentile(HotPercentileCutoff, Count))
                   : PSI->isHotCount(Count);
         }
 
-        if (ShouldRemove(IsHot)) {
-          Remove.push_back(II);
+        bool ToRemove = ShouldRemove(IsHot);
+        ReplaceWithValue.push_back({
+            II,
+            ToRemove ? Constant::getAllOnesValue(II->getType())
+                     : (Constant::getNullValue(II->getType())),
+        });
+        if (ToRemove)
           ++NumChecksRemoved;
-        }
         break;
       }
       default:
@@ -84,14 +86,16 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
     }
   }
 
-  for (IntrinsicInst *I : Remove)
+  for (auto [I, V] : ReplaceWithValue) {
+    I->replaceAllUsesWith(V);
     I->eraseFromParent();
+  }
 
-  return !Remove.empty();
+  return !ReplaceWithValue.empty();
 }
 
-PreservedAnalyses RemoveTrapsPass::run(Function &F,
-                                       FunctionAnalysisManager &AM) {
+PreservedAnalyses LowerBuiltinHotPass::run(Function &F,
+                                           FunctionAnalysisManager &AM) {
   if (F.isDeclaration())
     return PreservedAnalyses::all();
   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);