[ubsan][pgo] Pass to remove ubsan checks based on profile data (#83471)

UBSAN checks can be too expensive to be used
in release binaries. However not all code affect
performace in the same way. Removing small
number of checks in hot code we can performance
loss, preserving most of the checks.

Author: vitalybuka

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

@@ -0,0 +1,32 @@
+//===- RemoveTrapsPass.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.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+/// \file
+/// This file provides the interface for the pass responsible for removing
+/// expensive ubsan checks.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_UBSANOPTIMIZATIONPASS_H
+#define LLVM_TRANSFORMS_INSTRUMENTATION_UBSANOPTIMIZATIONPASS_H
+
+#include "llvm/IR/Function.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
+
+namespace llvm {
+
+// This pass is responsible for removing optional traps, like llvm.ubsantrap
+// from the hot code.
+class RemoveTrapsPass : public PassInfoMixin<RemoveTrapsPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
+} // namespace llvm
+
+#endif

llvm/lib/Passes/PassBuilder.cpp

@@ -177,6 +177,7 @@
 #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,6 +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("safe-stack", SafeStackPass(TM))
 FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass())
 FUNCTION_PASS("scalarizer", ScalarizerPass())

llvm/lib/Transforms/Instrumentation/CMakeLists.txt

@@ -17,6 +17,7 @@ add_llvm_component_library(LLVMInstrumentation
   PGOInstrumentation.cpp
   PGOMemOPSizeOpt.cpp
   PoisonChecking.cpp
+  RemoveTrapsPass.cpp
   SanitizerCoverage.cpp
   SanitizerBinaryMetadata.cpp
   ValueProfileCollector.cpp

llvm/lib/Transforms/Instrumentation/RemoveTrapsPass.cpp

@@ -0,0 +1,104 @@
+//===- RemoveTrapsPass.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/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/Support/RandomNumberGenerator.h"
+#include <memory>
+#include <random>
+
+using namespace llvm;
+
+#define DEBUG_TYPE "remove-traps"
+
+static cl::opt<int> HotPercentileCutoff(
+    "remove-traps-percentile-cutoff-hot", 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),
+               cl::desc("Probability value in the range [0.0, 1.0] of "
+                        "unconditional pseudo-random checks removal."));
+
+STATISTIC(NumChecksTotal, "Number of checks");
+STATISTIC(NumChecksRemoved, "Number of removed checks");
+
+static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
+                             const ProfileSummaryInfo *PSI) {
+  SmallVector<IntrinsicInst *, 16> Remove;
+  std::unique_ptr<RandomNumberGenerator> Rng;
+
+  auto ShouldRemove = [&](bool IsHot) {
+    if (!RandomRate.getNumOccurrences())
+      return IsHot;
+    if (!Rng)
+      Rng = F.getParent()->createRNG(F.getName());
+    std::bernoulli_distribution D(RandomRate);
+    return D(*Rng);
+  };
+
+  for (BasicBlock &BB : F) {
+    for (Instruction &I : BB) {
+      IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
+      if (!II)
+        continue;
+      auto ID = II->getIntrinsicID();
+      switch (ID) {
+      case Intrinsic::ubsantrap: {
+        ++NumChecksTotal;
+
+        bool IsHot = false;
+        if (PSI) {
+          uint64_t Count = 0;
+          for (const auto *PR : predecessors(&BB))
+            Count += BFI.getBlockProfileCount(PR).value_or(0);
+
+          IsHot =
+              HotPercentileCutoff.getNumOccurrences()
+                  ? (HotPercentileCutoff > 0 &&
+                     PSI->isHotCountNthPercentile(HotPercentileCutoff, Count))
+                  : PSI->isHotCount(Count);
+        }
+
+        if (ShouldRemove(IsHot)) {
+          Remove.push_back(II);
+          ++NumChecksRemoved;
+        }
+        break;
+      }
+      default:
+        break;
+      }
+    }
+  }
+
+  for (IntrinsicInst *I : Remove)
+    I->eraseFromParent();
+
+  return !Remove.empty();
+}
+
+PreservedAnalyses RemoveTrapsPass::run(Function &F,
+                                       FunctionAnalysisManager &AM) {
+  if (F.isDeclaration())
+    return PreservedAnalyses::all();
+  auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+  ProfileSummaryInfo *PSI =
+      MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
+  BlockFrequencyInfo &BFI = AM.getResult<BlockFrequencyAnalysis>(F);
+
+  return removeUbsanTraps(F, BFI, PSI) ? PreservedAnalyses::none()
+                                       : PreservedAnalyses::all();
+}