[ubsan][hwasan] Let mixing filters (llvm#100680)

Now the check will be enabled only if each filter is satisfied.
boomanaiden154 · Jul 26, 2024 · 0954205 · 0954205
1 parent c80c09f
commit 0954205
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 19 deletions.
diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -1528,11 +1528,7 @@ static void emitRemark(const Function &F, OptimizationRemarkEmitter &ORE,
 
 bool HWAddressSanitizer::selectiveInstrumentationShouldSkip(
     Function &F, FunctionAnalysisManager &FAM) const {
-  bool Skip = [&]() {
-    if (ClRandomSkipRate.getNumOccurrences()) {
-      std::bernoulli_distribution D(ClRandomSkipRate);
-      return !D(*Rng);
-    }
+  auto SkipHot = [&]() {
     if (!ClHotPercentileCutoff.getNumOccurrences())
       return false;
     auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
@@ -1544,7 +1540,16 @@ bool HWAddressSanitizer::selectiveInstrumentationShouldSkip(
     }
     return PSI->isFunctionHotInCallGraphNthPercentile(
         ClHotPercentileCutoff, &F, FAM.getResult<BlockFrequencyAnalysis>(F));
-  }();
+  };
+
+  auto SkipRandom = [&]() {
+    if (!ClRandomSkipRate.getNumOccurrences())
+      return false;
+    std::bernoulli_distribution D(ClRandomSkipRate);
+    return !D(*Rng);
+  };
+
+  bool Skip = SkipRandom() || SkipHot();
   emitRemark(F, FAM.getResult<OptimizationRemarkEmitterAnalysis>(F), Skip);
   return Skip;
 }

diff --git a/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp b/llvm/lib/Transforms/Instrumentation/LowerAllowCheckPass.cpp
@@ -76,13 +76,25 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
   SmallVector<std::pair<IntrinsicInst *, bool>, 16> ReplaceWithValue;
   std::unique_ptr<RandomNumberGenerator> Rng;
 
-  auto ShouldRemove = [&](bool IsHot) {
-    if (!RandomRate.getNumOccurrences())
-      return IsHot;
+  auto GetRng = [&]() -> RandomNumberGenerator & {
     if (!Rng)
       Rng = F.getParent()->createRNG(F.getName());
-    std::bernoulli_distribution D(RandomRate);
-    return !D(*Rng);
+    return *Rng;
+  };
+
+  auto ShouldRemoveHot = [&](const BasicBlock &BB) {
+    return HotPercentileCutoff.getNumOccurrences() && PSI &&
+           PSI->isHotCountNthPercentile(
+               HotPercentileCutoff, 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);
   };
 
   for (BasicBlock &BB : F) {
@@ -96,13 +108,7 @@ static bool removeUbsanTraps(Function &F, const BlockFrequencyInfo &BFI,
       case Intrinsic::allow_runtime_check: {
         ++NumChecksTotal;
 
-        bool IsHot = false;
-        if (PSI) {
-          uint64_t Count = BFI.getBlockProfileCount(&BB).value_or(0);
-          IsHot = PSI->isHotCountNthPercentile(HotPercentileCutoff, Count);
-        }
-
-        bool ToRemove = ShouldRemove(IsHot);
+        bool ToRemove = ShouldRemove(BB);
         ReplaceWithValue.push_back({
             II,
             ToRemove,

diff --git a/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll b/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll
@@ -2,6 +2,8 @@
 ; RUN: opt < %s -passes='require<profile-summary>,hwasan' -pass-remarks=hwasan -pass-remarks-missed=hwasan -S -hwasan-percentile-cutoff-hot=990000 2>&1 | FileCheck %s --check-prefix=NONE
 ; RUN: opt < %s -passes='require<profile-summary>,hwasan' -pass-remarks=hwasan -pass-remarks-missed=hwasan -S -hwasan-random-rate=1.0 2>&1 | FileCheck %s --check-prefix=ALL
 ; RUN: opt < %s -passes='require<profile-summary>,hwasan' -pass-remarks=hwasan -pass-remarks-missed=hwasan -S -hwasan-random-rate=0.0 2>&1 | FileCheck %s --check-prefix=NONE
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -pass-remarks=hwasan -pass-remarks-missed=hwasan -S -hwasan-random-rate=1.0 -hwasan-percentile-cutoff-hot=990000 2>&1 | FileCheck %s --check-prefix=NONE
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -pass-remarks=hwasan -pass-remarks-missed=hwasan -S -hwasan-random-rate=0.0 -hwasan-percentile-cutoff-hot=700000 2>&1 | FileCheck %s --check-prefix=NONE
 
 ; ALL: remark: <unknown>:0:0: Sanitized: F=sanitize
 ; ALL: @sanitized

diff --git a/llvm/test/Transforms/lower-builtin-allow-check.ll b/llvm/test/Transforms/lower-builtin-allow-check.ll
@@ -309,7 +309,7 @@ define dso_local noundef i32 @veryHot(ptr noundef readonly %0) !prof !39 {
 ; ALL70-LABEL: define dso_local noundef i32 @veryHot(
 ; ALL70-SAME: ptr noundef readonly [[TMP0:%.*]]) !prof [[PROF17:![0-9]+]] {
 ; ALL70-NEXT:    [[CHK:%.*]] = icmp eq ptr [[TMP0]], null
-; ALL70-NEXT:    [[HOT:%.*]] = xor i1 true, true
+; ALL70-NEXT:    [[HOT:%.*]] = xor i1 false, true
 ; ALL70-NEXT:    [[TMP2:%.*]] = or i1 [[CHK]], [[HOT]]
 ; ALL70-NEXT:    br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]]
 ; ALL70:       3: