[HWASAN] Follow up for #83503 implement selective instrumentation

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -52,6 +52,7 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/RandomNumberGenerator.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
@@ -61,6 +62,7 @@
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include <optional>
+#include <random>
 
 using namespace llvm;
 
@@ -181,16 +183,23 @@ static cl::opt<bool> ClWithTls(
     cl::Hidden, cl::init(true));
 
 static cl::opt<bool>
-    CSkipHotCode("hwasan-skip-hot-code",
-                 cl::desc("Do not instument hot functions based on FDO."),
-                 cl::Hidden, cl::init(false));
+    CSelectiveInstrumentation("hwasan-selective-instrumentation",
+                              cl::desc("Use selective instrumentation"),
+                              cl::Hidden, cl::init(false));
 
-static cl::opt<int> HotPercentileCutoff("hwasan-percentile-cutoff-hot",
-                                        cl::init(0));
+static cl::opt<int> HotPercentileCutoff(
+    "hwasan-percentile-cutoff-hot", cl::init(0),
+    cl::desc("Alternative hot percentile cuttoff."
+             "By default `-profile-summary-cutoff-hot` is used."));
 
-STATISTIC(NumTotalFuncs, "Number of total funcs HWASAN");
-STATISTIC(NumInstrumentedFuncs, "Number of HWASAN instrumented funcs");
-STATISTIC(NumNoProfileSummaryFuncs, "Number of HWASAN funcs without PS");
+static cl::opt<float>
+    RandomSkipRate("hwasan-random-skip-rate", cl::init(0),
+                   cl::desc("Probability value in the range [0.0, 1.0] "
+                            "to skip instrumentation of a function."));
+
+STATISTIC(NumTotalFuncs, "Number of total funcs");
+STATISTIC(NumInstrumentedFuncs, "Number of instrumented funcs");
+STATISTIC(NumNoProfileSummaryFuncs, "Number of funcs without PS");
 
 // Mode for selecting how to insert frame record info into the stack ring
 // buffer.
@@ -291,6 +300,8 @@ class HWAddressSanitizer {
     this->CompileKernel = ClEnableKhwasan.getNumOccurrences() > 0
                               ? ClEnableKhwasan
                               : CompileKernel;
+    this->Rng =
+        RandomSkipRate.getNumOccurrences() ? M.createRNG("hwasan") : nullptr;
 
     initializeModule();
   }
@@ -372,6 +383,7 @@ class HWAddressSanitizer {
   Module &M;
   const StackSafetyGlobalInfo *SSI;
   Triple TargetTriple;
+  std::unique_ptr<RandomNumberGenerator> Rng;
 
   /// This struct defines the shadow mapping using the rule:
   ///   shadow = (mem >> Scale) + Offset.
@@ -1526,19 +1538,26 @@ void HWAddressSanitizer::sanitizeFunction(Function &F,
     return;
 
   NumTotalFuncs++;
-  if (CSkipHotCode) {
-    auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
-    ProfileSummaryInfo *PSI =
-        MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
-    if (PSI && PSI->hasProfileSummary()) {
-      auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
-      if ((HotPercentileCutoff.getNumOccurrences() && HotPercentileCutoff >= 0)
-              ? PSI->isFunctionHotInCallGraphNthPercentile(HotPercentileCutoff,
-                                                           &F, BFI)
-              : PSI->isFunctionHotInCallGraph(&F, BFI))
+  if (CSelectiveInstrumentation) {
+    if (RandomSkipRate.getNumOccurrences()) {
+      std::bernoulli_distribution D(RandomSkipRate);
+      if (D(*Rng))
         return;
     } else {
-      ++NumNoProfileSummaryFuncs;
+      auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+      ProfileSummaryInfo *PSI =
+          MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
+      if (PSI && PSI->hasProfileSummary()) {
+        auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
+        if ((HotPercentileCutoff.getNumOccurrences() &&
+             HotPercentileCutoff >= 0)
+                ? PSI->isFunctionHotInCallGraphNthPercentile(
+                      HotPercentileCutoff, &F, BFI)
+                : PSI->isFunctionHotInCallGraph(&F, BFI))
+          return;
+      } else {
+        ++NumNoProfileSummaryFuncs;
+      }
     }
   }
   NumInstrumentedFuncs++;

llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out-no-ps.ll

@@ -1,16 +1,28 @@
-; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -stats 2>&1 \
-; RUN:   -hwasan-skip-hot-code=0 | FileCheck %s --check-prefix=FULL
-; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -stats 2>&1 \
-; RUN:   -hwasan-skip-hot-code=1 | FileCheck %s --check-prefix=SELSAN
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S  \
+; RUN:   -hwasan-selective-instrumentation=0 | FileCheck %s --check-prefix=FULL
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S  \
+; RUN:   -hwasan-selective-instrumentation=1 | FileCheck %s --check-prefix=SELSAN
 
-; REQUIRES: asserts
+; FULL: @not_sanitized
+; FULL-NEXT: %x = alloca i8, i64 4
+; FULL: @sanitized_no_ps
+; FULL-NEXT: @__hwasan_tls
 
-; FULL: 1 hwasan - Number of HWASAN instrumented funcs
-; FULL: 1 hwasan - Number of total funcs HWASAN
+; SELSAN: @not_sanitized
+; SELSAN-NEXT: %x = alloca i8, i64 4
+; SELSAN: @sanitized_no_ps
+; SELSAN-NEXT: @__hwasan_tls
 
-; SELSAN: 1 hwasan - Number of HWASAN instrumented funcs
-; SELSAN: 1 hwasan - Number of HWASAN funcs without PS
-; SELSAN: 1 hwasan - Number of total funcs HWASAN
+declare void @use(ptr)
 
-define void @not_sanitized() { ret void }
-define void @sanitized_no_ps() sanitize_hwaddress { ret void }
+define void @not_sanitized() {
+  %x = alloca i8, i64 4
+  call void @use(ptr %x)
+  ret void
+ }
+
+define void @sanitized_no_ps() sanitize_hwaddress {
+  %x = alloca i8, i64 4
+  call void @use(ptr %x)
+  ret void
+ }

llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll

@@ -1,16 +1,31 @@
-; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -stats 2>&1 \
-; RUN:   -hwasan-skip-hot-code=1 | FileCheck %s --check-prefix=DEFAULT
-; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -stats 2>&1 \
-; RUN:   -hwasan-skip-hot-code=1 -hwasan-percentile-cutoff-hot=700000 | FileCheck %s --check-prefix=PERCENT
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -hwasan-selective-instrumentation=1 \
+; RUN:    | FileCheck %s --check-prefix=DEFAULT
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -hwasan-selective-instrumentation=1 \
+; RUN:    -hwasan-percentile-cutoff-hot=700000 | FileCheck %s --check-prefix=HOT_RATE
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -hwasan-selective-instrumentation=1 \
+; RUN:    -hwasan-random-skip-rate=0.0 | FileCheck %s --check-prefix=RANDOM_RATE_0
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -hwasan-selective-instrumentation=1 \
+; RUN:    -hwasan-random-skip-rate=1.0 | FileCheck %s --check-prefix=RANDOM_RATE_1
 
-; REQUIRES: asserts
+; DEFAULT: @sanitized
+; DEFAULT-NEXT: %x = alloca i8, i64 4
 
-; DEFAULT: 1 hwasan - Number of total funcs HWASAN
+; HOT_RATE: @sanitized
+; HOT_RATE-NEXT: @__hwasan_tls
 
-; PERCENT: 1 hwasan - Number of HWASAN instrumented funcs
-; PERCENT: 1 hwasan - Number of total funcs HWASAN
+; RANDOM_RATE_0: @sanitized
+; RANDOM_RATE_0-NEXT: @__hwasan_tls
 
-define void @sanitized() sanitize_hwaddress !prof !36 { ret void }
+; RANDOM_RATE_1: @sanitized
+; RANDOM_RATE_1-NEXT: %x = alloca i8, i64 4
+
+declare void @use(ptr)
+
+define void @sanitized(i32 noundef %0) sanitize_hwaddress !prof !36 {
+  %x = alloca i8, i64 4
+  call void @use(ptr %x)
+  ret void
+}
 
 !llvm.module.flags = !{!6}
 !6 = !{i32 1, !"ProfileSummary", !7}