[llvm] [HWASAN] Follow up for #83503 implement selective instrumentation (PR #83942)

[Kirill Stoimenov via llvm-commits]

https://github.com/kstoimenov created https://github.com/llvm/llvm-project/pull/83942

1. Change tests to use IR instead of -stats to avoid depending on Debug mode
2. Add SkipInstrumentationRandomRate 
3. remove HWASAN from stat strings


>From 7b69b728bae8a38ce614af9f704c99f44ee72f73 Mon Sep 17 00:00:00 2001
From: Kirill Stoimenov <kstoimenov at google.com>
Date: Tue, 5 Mar 2024 01:44:32 +0000
Subject: [PATCH] [HWASAN] Change tests to use IR instead of -stats to avoid
 depnding on Debug mode, add SkipInstrumentationRandomRate and remove HWASAN
 from stat strings.

---
 .../Instrumentation/HWAddressSanitizer.cpp    | 15 +++++++--
 .../HWAddressSanitizer/pgo-opt-out-no-ps.ll   | 32 +++++++++++++------
 .../HWAddressSanitizer/pgo-opt-out.ll         | 23 +++++++------
 3 files changed, 48 insertions(+), 22 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index 4404382a85b7e3..92f8687030d609 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -83,6 +83,8 @@ static const uint64_t kDynamicShadowSentinel =
 
 static const unsigned kShadowBaseAlignment = 32;
 
+static constexpr unsigned MaxRandomRate = 1000;
+
 static cl::opt<std::string>
     ClMemoryAccessCallbackPrefix("hwasan-memory-access-callback-prefix",
                                  cl::desc("Prefix for memory access callbacks"),
@@ -188,9 +190,13 @@ static cl::opt<bool>
 static cl::opt<int> HotPercentileCutoff("hwasan-percentile-cutoff-hot",
                                         cl::init(0));
 
-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> SkipInstRandomRate(
+    "hwasan-skip-inst-random-rate", cl::init(0.0),
+    cl::desc("Probability 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.
@@ -1527,6 +1533,9 @@ void HWAddressSanitizer::sanitizeFunction(Function &F,
 
   NumTotalFuncs++;
   if (CSkipHotCode) {
+    if ((F.getGUID() % MaxRandomRate) < SkipInstRandomRate) {
+      return;
+    }
     auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
     ProfileSummaryInfo *PSI =
         MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
diff --git a/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out-no-ps.ll b/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out-no-ps.ll
index 2aa218fa1522d6..c3c2cbe2d5bb92 100644
--- a/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out-no-ps.ll
+++ b/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: opt < %s -passes='require<profile-summary>,hwasan' -S  \
 ; RUN:   -hwasan-skip-hot-code=0 | FileCheck %s --check-prefix=FULL
-; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -stats 2>&1 \
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S  \
 ; RUN:   -hwasan-skip-hot-code=1 | FileCheck %s --check-prefix=SELSAN
 
-; REQUIRES: asserts
+; FULL: @not_sanitized
+; FULL-NEXT: %x = alloca i8, i64 4
+; FULL: @sanitized_no_ps
+; FULL-SAME: @__hwasan_personality_thunk
 
-; 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-SAME: @__hwasan_personality_thunk
 
-; 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
+ }
diff --git a/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll b/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll
index 65a5f8c9689678..9cdb927f06760e 100644
--- a/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll
+++ b/llvm/test/Instrumentation/HWAddressSanitizer/pgo-opt-out.ll
@@ -1,16 +1,21 @@
-; 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-skip-hot-code=1 \
+; RUN:    | FileCheck %s --check-prefix=DEFAULT
+; RUN: opt < %s -passes='require<profile-summary>,hwasan' -S -hwasan-skip-hot-code=1 \
+; RUN:    -hwasan-percentile-cutoff-hot=700000 | FileCheck %s --check-prefix=PERCENT
 
-; REQUIRES: asserts
+; DEFAULT: @sanitized
+; DEFAULT-NEXT: %x = alloca i8, i64 4
 
-; DEFAULT: 1 hwasan - Number of total funcs HWASAN
+; PERCENT: @sanitized
+; PERCENT-SAME: @__hwasan_personality_thunk
 
-; PERCENT: 1 hwasan - Number of HWASAN instrumented funcs
-; PERCENT: 1 hwasan - Number of total funcs HWASAN
+declare void @use(ptr)
 
-define void @sanitized() sanitize_hwaddress !prof !36 { ret void }
+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}