[HWASAN] Implement selective instrumentation based on profiling information

llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp

@@ -15,11 +15,14 @@
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
@@ -177,6 +180,22 @@ static cl::opt<bool> ClWithTls(
              "platforms that support this"),
     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));
+
+static cl::opt<int> HotPercentileCutoff("hwasan-percentile-cutoff-hot",
+                                        cl::init(0));
+
+STATISTIC(NumTotalFuncs, "Number of funcs seen by HWASAN");
+STATISTIC(NumHwasanCtors, "Number of HWASAN ctors");
+STATISTIC(NumNoSanitizeFuncs, "Number of no-sanitize HWASAN funcs");
+STATISTIC(NumConsideredFuncs, "Number of funcs considered for HWASAN");
+STATISTIC(NumInstrumentedFuncs, "Number of HWASAN instrumented funcs");
+STATISTIC(NumNoProfileSummaryFuncs, "Number of HWASAN funcs without PS");
+STATISTIC(NumSkippedHotFuncs, "Number of skipped hot HWASAN funcs");
+
 // Mode for selecting how to insert frame record info into the stack ring
 // buffer.
 enum RecordStackHistoryMode {
@@ -1501,11 +1520,39 @@ bool HWAddressSanitizer::instrumentStack(memtag::StackInfo &SInfo,
 
 void HWAddressSanitizer::sanitizeFunction(Function &F,
                                           FunctionAnalysisManager &FAM) {
-  if (&F == HwasanCtorFunction)
+  NumTotalFuncs++;
+  if (&F == HwasanCtorFunction) {
+    NumHwasanCtors++;
     return;
+  }
 
-  if (!F.hasFnAttribute(Attribute::SanitizeHWAddress))
+  if (!F.hasFnAttribute(Attribute::SanitizeHWAddress)) {
+    NumNoSanitizeFuncs++;
     return;
+  }
+
+  NumConsideredFuncs++;
+  if (CSkipHotCode) {
+    auto &MAMProxy = FAM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+    ProfileSummaryInfo *PSI =
+        MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
+    if (PSI != nullptr && PSI->hasProfileSummary()) {
+      auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
+      const bool is_hot =
+          (HotPercentileCutoff.getNumOccurrences() && HotPercentileCutoff >= 0)
+              ? PSI->isFunctionHotInCallGraphNthPercentile(HotPercentileCutoff,
+                                                           &F, BFI)
+              : PSI->isFunctionEntryHot(&F);
+
+      if (is_hot) {
+        ++NumSkippedHotFuncs;
+        return;
+      }
+    } else {
+      ++NumNoProfileSummaryFuncs;
+    }
+  }
+  NumInstrumentedFuncs++;
 
   LLVM_DEBUG(dbgs() << "Function: " << F.getName() << "\n");
 

llvm/test/Instrumentation/HWAddressSanitizer/selsan-no-ps.ll

@@ -0,0 +1,16 @@
+; 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
+
+; FULL: 1 hwasan - Number of funcs considered for HWASAN
+; FULL: 1 hwasan - Number of HWASAN ctors
+; FULL: 1 hwasan - Number of HWASAN instrumented funcs
+
+; SELSAN: 1 hwasan - Number of funcs considered for HWASAN
+; SELSAN: 1 hwasan - Number of HWASAN ctors
+; SELSAN: 1 hwasan - Number of HWASAN instrumented funcs
+; SELSAN: 1 hwasan - Number of HWASAN funcs without PS
+
+define void @not_sanitized() { ret void }
+define void @sanitized_no_ps() sanitize_hwaddress { ret void }

llvm/test/Instrumentation/HWAddressSanitizer/selsan.ll

@@ -0,0 +1,32 @@
+; 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
+
+; DEFAULT: 1 hwasan - Number of funcs considered for HWASAN
+; DEFAULT: 1 hwasan - Number of skipped hot HWASAN funcs
+
+; PERCENT: 1 hwasan - Number of funcs considered for HWASAN
+; PERCENT: 1 hwasan - Number of HWASAN instrumented funcs
+
+define void @sanitized() sanitize_hwaddress !prof !36 { ret void }
+
+!llvm.module.flags = !{!6}
+!6 = !{i32 1, !"ProfileSummary", !7}
+!7 = !{!8, !9, !10, !11, !12, !13, !14, !17}
+!8 = !{!"ProfileFormat", !"InstrProf"}
+!9 = !{!"TotalCount", i64 30000}
+!10 = !{!"MaxCount", i64 10000}
+!11 = !{!"MaxInternalCount", i64 10000}
+!12 = !{!"MaxFunctionCount", i64 10000}
+!13 = !{!"NumCounts", i64 3}
+!14 = !{!"NumFunctions", i64 5}
+!17 = !{!"DetailedSummary", !18}
+!18 = !{!19, !29, !30, !32, !34}
+!19 = !{i32 10000, i64 10000, i32 3}
+!29 = !{i32 950000, i64 5000, i32 3}
+!30 = !{i32 990000, i64 500, i32 4}
+!32 = !{i32 999900, i64 250, i32 4}
+!34 = !{i32 999999, i64 1, i32 6}
+
+!36 = !{!"function_entry_count", i64 1000}