[ubsan] Parse and use <cutoffs[0,1,2]=70000;cutoffs[5,6,8]=90000> in LowerAllowCheckPass

llvm/include/llvm/Transforms/Instrumentation/LowerAllowCheckPass.h

@@ -25,14 +25,16 @@ namespace llvm {
 class LowerAllowCheckPass : public PassInfoMixin<LowerAllowCheckPass> {
 public:
   struct Options {
-    std::vector<unsigned int> placeholder; // TODO: cutoffs
+    std::vector<unsigned int> cutoffs;
   };
 
   explicit LowerAllowCheckPass(LowerAllowCheckPass::Options Opts)
       : Opts(std::move(Opts)) {};
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 
   static bool IsRequested();
+  void printPipeline(raw_ostream &OS,
+                     function_ref<StringRef(StringRef)> MapClassName2PassName);
 
 private:
   LowerAllowCheckPass::Options Opts;

llvm/lib/Passes/PassBuilder.cpp

@@ -828,11 +828,65 @@ parseLowerAllowCheckPassOptions(StringRef Params) {
     StringRef ParamName;
     std::tie(ParamName, Params) = Params.split(';');
 
-    return make_error<StringError>(
-        formatv("invalid LowerAllowCheck pass parameter '{0}' ", ParamName)
-            .str(),
-        inconvertibleErrorCode());
+    // Format is <cutoffs[1,2,3]=70000;cutoffs[5,6,8]=90000>
+    //
+    // Parsing allows duplicate indices (last one takes precedence).
+    // It would technically be in spec to specify
+    //   cutoffs[0]=70000,cutoffs[1]=90000,cutoffs[0]=80000,...
+    if (ParamName.starts_with("cutoffs[")) {
+      StringRef IndicesStr;
+      StringRef CutoffStr;
+
+      std::tie(IndicesStr, CutoffStr) = ParamName.split("]=");
+      //       cutoffs[1,2,3
+      //                   70000
+
+      int cutoff;
+      if (CutoffStr.getAsInteger(0, cutoff))
+        return make_error<StringError>(
+            formatv("invalid LowerAllowCheck pass cutoffs parameter '{0}' "
+                    "({1})",
+                    CutoffStr, Params)
+                .str(),
+            inconvertibleErrorCode());
+
+      if (!IndicesStr.consume_front("cutoffs[") || IndicesStr == "")
+        return make_error<StringError>(
+            formatv("invalid LowerAllowCheck pass index parameter '{0}' "
+                    "({1})",
+                    IndicesStr, CutoffStr)
+                .str(),
+            inconvertibleErrorCode());
+
+      while (IndicesStr != "") {
+        StringRef firstIndexStr;
+        std::tie(firstIndexStr, IndicesStr) = IndicesStr.split('|');
+
+        unsigned int index;
+        if (firstIndexStr.getAsInteger(0, index))
+          return make_error<StringError>(
+              formatv("invalid LowerAllowCheck pass index parameter '{0}' "
+                      "({1}) {2}",
+                      firstIndexStr, IndicesStr)
+                  .str(),
+              inconvertibleErrorCode());
+
+        // In the common case (sequentially increasing indices), we will issue
+        // O(n) resize requests. We assume the underlying data structure has
+        // O(1) runtime for each added element.
+        if (index >= Result.cutoffs.size())
+          Result.cutoffs.resize(index + 1, 0);
+
+        Result.cutoffs[index] = cutoff;
+      }
+    } else {
+      return make_error<StringError>(
+          formatv("invalid LowerAllowCheck pass parameter '{0}' ", ParamName)
+              .str(),
+          inconvertibleErrorCode());
+    }
   }
+
   return Result;
 }
 

llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp

@@ -20,6 +20,7 @@
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/RandomNumberGenerator.h"
 #include <memory>
 #include <random>
@@ -71,7 +72,8 @@ static void emitRemark(IntrinsicInst *II, OptimizationRemarkEmitter &ORE,
 
 static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
                              const ProfileSummaryInfo *PSI,
-                             OptimizationRemarkEmitter &ORE) {
+                             OptimizationRemarkEmitter &ORE,
+                             const std::vector<unsigned int> &cutoffs) {
   SmallVector<std::pair<IntrinsicInst *, bool>, 16> ReplaceWithValue;
   std::unique_ptr<RandomNumberGenerator> Rng;
 
@@ -81,19 +83,32 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
     return *Rng;
   };
 
-  auto ShouldRemoveHot = [&](const BasicBlock &BB) {
-    return HotPercentileCutoff.getNumOccurrences() && PSI &&
-           PSI->isHotCountNthPercentile(
-               HotPercentileCutoff, BFI.getBlockProfileCount(&BB).value_or(0));
+  auto GetCutoff = [&](const IntrinsicInst *II) -> unsigned {
+    if (HotPercentileCutoff.getNumOccurrences())
+      return HotPercentileCutoff;
+    else if (II->getIntrinsicID() == Intrinsic::allow_ubsan_check) {
+      auto *Kind = cast<ConstantInt>(II->getArgOperand(0));
+      if (Kind->getZExtValue() < cutoffs.size())
+        return cutoffs[Kind->getZExtValue()];
+    }
+
+    return 0;
+  };
+
+  auto ShouldRemoveHot = [&](const BasicBlock &BB, unsigned int cutoff) {
+    return (cutoff == 1000000) ||
+           (PSI && PSI->isHotCountNthPercentile(
+                       cutoff, BFI.getBlockProfileCount(&BB).value_or(0)));
   };
 
   auto ShouldRemoveRandom = [&]() {
     return RandomRate.getNumOccurrences() &&
            !std::bernoulli_distribution(RandomRate)(GetRng());
   };
 
-  auto ShouldRemove = [&](const BasicBlock &BB) {
-    return ShouldRemoveRandom() || ShouldRemoveHot(BB);
+  auto ShouldRemove = [&](const IntrinsicInst *II) {
+    unsigned int cutoff = GetCutoff(II);
+    return ShouldRemoveRandom() || ShouldRemoveHot(*(II->getParent()), cutoff);
   };
 
   for (BasicBlock &BB : F) {
@@ -107,7 +122,8 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
       case Intrinsic::allow_runtime_check: {
         ++NumChecksTotal;
 
-        bool ToRemove = ShouldRemove(BB);
+        bool ToRemove = ShouldRemove(II);
+
         ReplaceWithValue.push_back({
             II,
             ToRemove,
@@ -142,11 +158,37 @@ PreservedAnalyses LowerAllowCheckPass::run(Function &F,
   OptimizationRemarkEmitter &ORE =
       AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
 
-  return removeUbsanTraps(F, BFI, PSI, ORE) ? PreservedAnalyses::none()
-                                            : PreservedAnalyses::all();
+  return removeUbsanTraps(F, BFI, PSI, ORE, Opts.cutoffs)
+             ? PreservedAnalyses::none()
+             : PreservedAnalyses::all();
 }
 
 bool LowerAllowCheckPass::IsRequested() {
   return RandomRate.getNumOccurrences() ||
          HotPercentileCutoff.getNumOccurrences();
 }
+
+void LowerAllowCheckPass::printPipeline(
+    raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
+  static_cast<PassInfoMixin<LowerAllowCheckPass> *>(this)->printPipeline(
+      OS, MapClassName2PassName);
+  OS << "<";
+
+  // Format is <cutoffs[0,1,2]=70000;cutoffs[5,6,8]=90000>
+  // but it's equally valid to specify
+  //   cutoffs[0]=70000;cutoffs[1]=70000;cutoffs[2]=70000;cutoffs[5]=90000;...
+  // and that's what we do here. It is verbose but valid and easy to verify
+  // correctness.
+  // TODO: print shorter output by combining adjacent runs, etc.
+  int i = 0;
+  for (unsigned int cutoff : Opts.cutoffs) {
+    if (cutoff > 0) {
+      if (i > 0)
+        OS << ";";
+      OS << "cutoffs[" << i << "]=" << cutoff;
+    }
+
+    i++;
+  }
+  OS << '>';
+}

llvm/test/Transforms/lower-builtin-allow-check.ll

@@ -2,6 +2,14 @@
 ; RUN: opt < %s -passes='function(lower-allow-check)' -S | FileCheck %s --check-prefixes=NOPROFILE
 ; RUN: opt < %s -passes='function(lower-allow-check)' -lower-allow-check-random-rate=0 -S | FileCheck %s --check-prefixes=NONE
 ; RUN: opt < %s -passes='function(lower-allow-check)' -lower-allow-check-random-rate=1 -S | FileCheck %s --check-prefixes=ALL
+;
+; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check<cutoffs[22]=990000>)' -S | FileCheck %s --check-prefixes=HOT99
+; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check<cutoffs[22]=700000>)' -S | FileCheck %s --check-prefixes=HOT70
+; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check<cutoffs[22]=990000>)' -lower-allow-check-random-rate=0 -S | FileCheck %s --check-prefixes=NONE99
+; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check<cutoffs[22]=700000>)' -lower-allow-check-random-rate=1 -S | FileCheck %s --check-prefixes=ALL70
+;
+; -lower-allow-check-percentile-cutoff is deprecated and will be removed in the future;
+; use the cutoffs parameter to the lower-allow-check pass, as shown above.
 ; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check)' -lower-allow-check-percentile-cutoff-hot=990000 -S | FileCheck %s --check-prefixes=HOT99
 ; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check)' -lower-allow-check-percentile-cutoff-hot=700000 -S | FileCheck %s --check-prefixes=HOT70
 ; RUN: opt < %s -passes='require<profile-summary>,function(lower-allow-check)' -lower-allow-check-random-rate=0 -lower-allow-check-percentile-cutoff-hot=990000 -S | FileCheck %s --check-prefixes=NONE99