[llvm] f3587d4 - [sancov] Apply branch weights when checking counters (#89458)

[via llvm-commits]

Author: Vitaly Buka
Date: 2024-04-19T17:01:18-07:00
New Revision: f3587d41064ce22330231baee1553b210777f3e3

URL: https://github.com/llvm/llvm-project/commit/f3587d41064ce22330231baee1553b210777f3e3
DIFF: https://github.com/llvm/llvm-project/commit/f3587d41064ce22330231baee1553b210777f3e3.diff

LOG: [sancov] Apply branch weights when checking counters (#89458)

It reduces instrumentation overhead by ~50%.

Added: 
    

Modified: 
    llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
    llvm/test/Instrumentation/SanitizerCoverage/inline-bool-flag.ll
    llvm/test/Instrumentation/SanitizerCoverage/stack-depth.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
index 17c1c442384220..239e7c602a5471 100644
--- a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
+++ b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
@@ -25,6 +25,7 @@
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/CommandLine.h"
@@ -979,8 +980,9 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
         FunctionBoolArray->getValueType(), FunctionBoolArray,
         {ConstantInt::get(IntptrTy, 0), ConstantInt::get(IntptrTy, Idx)});
     auto Load = IRB.CreateLoad(Int1Ty, FlagPtr);
-    auto ThenTerm =
-        SplitBlockAndInsertIfThen(IRB.CreateIsNull(Load), &*IP, false);
+    auto ThenTerm = SplitBlockAndInsertIfThen(
+        IRB.CreateIsNull(Load), &*IP, false,
+        MDBuilder(IRB.getContext()).createBranchWeights(1, (1 << 20) - 1));
     IRBuilder<> ThenIRB(ThenTerm);
     auto Store = ThenIRB.CreateStore(ConstantInt::getTrue(Int1Ty), FlagPtr);
     Load->setNoSanitizeMetadata();
@@ -997,7 +999,9 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
     auto FrameAddrInt = IRB.CreatePtrToInt(FrameAddrPtr, IntptrTy);
     auto LowestStack = IRB.CreateLoad(IntptrTy, SanCovLowestStack);
     auto IsStackLower = IRB.CreateICmpULT(FrameAddrInt, LowestStack);
-    auto ThenTerm = SplitBlockAndInsertIfThen(IsStackLower, &*IP, false);
+    auto ThenTerm = SplitBlockAndInsertIfThen(
+        IsStackLower, &*IP, false,
+        MDBuilder(IRB.getContext()).createBranchWeights(1, (1 << 20) - 1));
     IRBuilder<> ThenIRB(ThenTerm);
     auto Store = ThenIRB.CreateStore(FrameAddrInt, SanCovLowestStack);
     LowestStack->setNoSanitizeMetadata();

diff  --git a/llvm/test/Instrumentation/SanitizerCoverage/inline-bool-flag.ll b/llvm/test/Instrumentation/SanitizerCoverage/inline-bool-flag.ll
index 1380e6ac090cc2..a6d7ae67d8453c 100644
--- a/llvm/test/Instrumentation/SanitizerCoverage/inline-bool-flag.ll
+++ b/llvm/test/Instrumentation/SanitizerCoverage/inline-bool-flag.ll
@@ -19,7 +19,7 @@ define void @foo() {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i1, ptr @__sancov_gen_, align 1, !nosanitize [[META0:![0-9]+]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i1 [[TMP0]], false
-; CHECK-NEXT:    br i1 [[TMP1]], label [[TMP2:%.*]], label [[TMP3:%.*]]
+; CHECK-NEXT:    br i1 [[TMP1]], label [[TMP2:%.*]], label [[TMP3:%.*]], !prof [[PROF1:![0-9]+]]
 ; CHECK:       2:
 ; CHECK-NEXT:    store i1 true, ptr @__sancov_gen_, align 1, !nosanitize [[META0]]
 ; CHECK-NEXT:    br label [[TMP3]]
@@ -34,4 +34,5 @@ entry:
 ; CHECK: attributes #[[ATTR0:[0-9]+]] = { nounwind }
 ;.
 ; CHECK: [[META0]] = !{}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 1, i32 1048575}
 ;.

diff  --git a/llvm/test/Instrumentation/SanitizerCoverage/stack-depth.ll b/llvm/test/Instrumentation/SanitizerCoverage/stack-depth.ll
index 908bca6cdf74b3..00547afd1b9ef2 100644
--- a/llvm/test/Instrumentation/SanitizerCoverage/stack-depth.ll
+++ b/llvm/test/Instrumentation/SanitizerCoverage/stack-depth.ll
@@ -43,7 +43,7 @@ define i32 @bar() {
 ; L1-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[TMP0]] to i64
 ; L1-NEXT:    [[TMP2:%.*]] = load i64, ptr @__sancov_lowest_stack, align 8, !nosanitize [[META0:![0-9]+]]
 ; L1-NEXT:    [[TMP3:%.*]] = icmp ult i64 [[TMP1]], [[TMP2]]
-; L1-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP5:%.*]]
+; L1-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP5:%.*]], !prof [[PROF1:![0-9]+]]
 ; L1:       4:
 ; L1-NEXT:    store i64 [[TMP1]], ptr @__sancov_lowest_stack, align 8, !nosanitize [[META0]]
 ; L1-NEXT:    br label [[TMP5]]
@@ -58,7 +58,7 @@ define i32 @bar() {
 ; L3-NEXT:    [[TMP1:%.*]] = ptrtoint ptr [[TMP0]] to i64
 ; L3-NEXT:    [[TMP2:%.*]] = load i64, ptr @__sancov_lowest_stack, align 8, !nosanitize [[META0:![0-9]+]]
 ; L3-NEXT:    [[TMP3:%.*]] = icmp ult i64 [[TMP1]], [[TMP2]]
-; L3-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP5:%.*]]
+; L3-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP5:%.*]], !prof [[PROF1:![0-9]+]]
 ; L3:       4:
 ; L3-NEXT:    store i64 [[TMP1]], ptr @__sancov_lowest_stack, align 8, !nosanitize [[META0]]
 ; L3-NEXT:    br label [[TMP5]]
@@ -90,6 +90,8 @@ define weak_odr hidden ptr @_ZTW21__sancov_lowest_stack() {
 ; L3: attributes #[[ATTR2]] = { nomerge }
 ;.
 ; L1: [[META0]] = !{}
+; L1: [[PROF1]] = !{!"branch_weights", i32 1, i32 1048575}
 ;.
 ; L3: [[META0]] = !{}
+; L3: [[PROF1]] = !{!"branch_weights", i32 1, i32 1048575}
 ;.