[SanitizerCoverage] Add gated tracing callbacks support to trace-cmp

thetruestblue · thetruestblue · commit dd972e68de67 · 2024-10-21T14:31:48.000-07:00
The option -sanitizer-coverage-gated-trace-callbacks gates the invocation of the
trace-pc-guard callbacks based on the value of a global variable, which is stored
in a specific section.
In this commit, we extend this feature to trace-cmp and gate the cmp callbacks to
the same variable used for trace-pc-guard.

Update SanitizerCoverage doc with this flag.

rdar://135404160

Patch by: Andrea Fioraldi
diff --git a/clang/docs/SanitizerCoverage.rst b/clang/docs/SanitizerCoverage.rst
@@ -385,6 +385,20 @@ Users need to implement a single function to capture the CF table at startup:
     // the collected control flow.
   }
 
+Gated Trace Callbacks
+=====================
+
+Gate the invocation of the tracing callbacks with
+``-sanitizer-coverage-gated-trace-callbacks``.
+
+When this option is enabled, the instrumentation will not call into the
+runtime-provided callbacks for tracing, thus only incurring in a trivial
+branch without going through a function call.
+
+It is up to the runtime to toggle the value of the global variable in order to
+enable tracing.
+
+This option is only supported for trace-pc-guard and trace-cmp.
 
 Disabling instrumentation with ``__attribute__((no_sanitize("coverage")))``
 ===========================================================================
diff --git a/clang/test/CodeGen/sanitize-coverage-gated-callbacks.c b/clang/test/CodeGen/sanitize-coverage-gated-callbacks.c
@@ -1,5 +1,7 @@
 // RUN: %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=trace-pc-guard -mllvm -sanitizer-coverage-gated-trace-callbacks=1 -o - | FileCheck %s --check-prefixes=CHECK,GATED
 // RUN: %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=trace-pc-guard -mllvm -sanitizer-coverage-gated-trace-callbacks=0 -o - | FileCheck %s --check-prefixes=CHECK,PLAIN
+// RUN: %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=trace-pc-guard,trace-cmp -mllvm -sanitizer-coverage-gated-trace-callbacks=1 -o - | FileCheck %s --check-prefixes=CHECK,GATED,GATEDCMP
+// RUN: %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=trace-pc-guard,trace-cmp -mllvm -sanitizer-coverage-gated-trace-callbacks=0 -o - | FileCheck %s --check-prefixes=CHECK,PLAIN,PLAINCMP
 // RUN: not %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=trace-pc -mllvm -sanitizer-coverage-gated-trace-callbacks=1 -o /dev/null 2>&1 | FileCheck %s --check-prefixes=INCOMPATIBLE
 // RUN: not %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=inline-8bit-counters -mllvm -sanitizer-coverage-gated-trace-callbacks=1 -o /dev/null 2>&1 | FileCheck %s --check-prefixes=INCOMPATIBLE
 // RUN: not %clang %s -target arm64-apple-darwin -emit-llvm -S -fsanitize-coverage=inline-bool-flag -mllvm -sanitizer-coverage-gated-trace-callbacks=1 -o /dev/null 2>&1 | FileCheck %s --check-prefixes=INCOMPATIBLE
@@ -9,7 +11,7 @@
 // PLAIN-NOT: section "__DATA,__sancov_gate"
 
 // Produce an error for all incompatible sanitizer coverage modes.
-// INCOMPATIBLE: error: 'sanitizer-coverage-gated-trace-callbacks' is only supported with trace-pc-guard
+// INCOMPATIBLE: error: 'sanitizer-coverage-gated-trace-callbacks' is only supported with trace-pc-guard and trace-cmp
 
 int x[10];
 
diff --git a/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp b/llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp
@@ -159,8 +159,8 @@ static cl::opt<bool>
 
 static cl::opt<bool> ClGatedCallbacks(
     "sanitizer-coverage-gated-trace-callbacks",
-    cl::desc("Gate the invocation of the tracing callbacks on a global "
-             "variable. Currently only supported for trace-pc-guard."),
+    cl::desc("Gate the invocation of the tracing callbacks on a global variable"
+             ". Currently only supported for trace-pc-guard and trace-cmp."),
     cl::Hidden, cl::init(false));
 
 namespace {
@@ -235,22 +235,26 @@ class ModuleSanitizerCoverage {
   void instrumentFunction(Function &F);
   void InjectCoverageForIndirectCalls(Function &F,
                                       ArrayRef<Instruction *> IndirCalls);
-  void InjectTraceForCmp(Function &F, ArrayRef<Instruction *> CmpTraceTargets);
+  void InjectTraceForCmp(Function &F, ArrayRef<Instruction *> CmpTraceTargets,
+                         Value *&FunctionGateCmp);
   void InjectTraceForDiv(Function &F,
                          ArrayRef<BinaryOperator *> DivTraceTargets);
   void InjectTraceForGep(Function &F,
                          ArrayRef<GetElementPtrInst *> GepTraceTargets);
   void InjectTraceForLoadsAndStores(Function &F, ArrayRef<LoadInst *> Loads,
                                     ArrayRef<StoreInst *> Stores);
   void InjectTraceForSwitch(Function &F,
-                            ArrayRef<Instruction *> SwitchTraceTargets);
+                            ArrayRef<Instruction *> SwitchTraceTargets,
+                            Value *&FunctionGateCmp);
   bool InjectCoverage(Function &F, ArrayRef<BasicBlock *> AllBlocks,
-                      bool IsLeafFunc = true);
+                      Value *&FunctionGateCmp, bool IsLeafFunc = true);
   GlobalVariable *CreateFunctionLocalArrayInSection(size_t NumElements,
                                                     Function &F, Type *Ty,
                                                     const char *Section);
   GlobalVariable *CreatePCArray(Function &F, ArrayRef<BasicBlock *> AllBlocks);
   void CreateFunctionLocalArrays(Function &F, ArrayRef<BasicBlock *> AllBlocks);
+  Instruction *CreateGateBranch(Function &F, Value *&FunctionGateCmp,
+                                Instruction *I);
   Value *CreateFunctionLocalGateCmp(IRBuilder<> &IRB);
   void InjectCoverageAtBlock(Function &F, BasicBlock &BB, size_t Idx,
                              Value *&FunctionGateCmp, bool IsLeafFunc = true);
@@ -493,9 +497,9 @@ bool ModuleSanitizerCoverage::instrumentModule() {
     SanCovLowestStack->setInitializer(Constant::getAllOnesValue(IntptrTy));
 
   if (Options.GatedCallbacks) {
-    if (!Options.TracePCGuard) {
+    if (!Options.TracePCGuard && !Options.TraceCmp) {
       C->emitError(StringRef("'") + ClGatedCallbacks.ArgStr +
-                   "' is only supported with trace-pc-guard");
+                   "' is only supported with trace-pc-guard or trace-cmp");
       return true;
     }
 
@@ -724,10 +728,11 @@ void ModuleSanitizerCoverage::instrumentFunction(Function &F) {
   if (Options.CollectControlFlow)
     createFunctionControlFlow(F);
 
-  InjectCoverage(F, BlocksToInstrument, IsLeafFunc);
+  Value *FunctionGateCmp = nullptr;
+  InjectCoverage(F, BlocksToInstrument, FunctionGateCmp, IsLeafFunc);
   InjectCoverageForIndirectCalls(F, IndirCalls);
-  InjectTraceForCmp(F, CmpTraceTargets);
-  InjectTraceForSwitch(F, SwitchTraceTargets);
+  InjectTraceForCmp(F, CmpTraceTargets, FunctionGateCmp);
+  InjectTraceForSwitch(F, SwitchTraceTargets, FunctionGateCmp);
   InjectTraceForDiv(F, DivTraceTargets);
   InjectTraceForGep(F, GepTraceTargets);
   InjectTraceForLoadsAndStores(F, Loads, Stores);
@@ -816,12 +821,30 @@ Value *ModuleSanitizerCoverage::CreateFunctionLocalGateCmp(IRBuilder<> &IRB) {
   return Cmp;
 }
 
+Instruction *ModuleSanitizerCoverage::CreateGateBranch(Function &F,
+                                                       Value *&FunctionGateCmp,
+                                                       Instruction *IP) {
+  if (!FunctionGateCmp) {
+    // Create this in the entry block
+    BasicBlock &BB = F.getEntryBlock();
+    BasicBlock::iterator IP = BB.getFirstInsertionPt();
+    IP = PrepareToSplitEntryBlock(BB, IP);
+    IRBuilder<> EntryIRB(&*IP);
+    FunctionGateCmp = CreateFunctionLocalGateCmp(EntryIRB);
+  }
+  // Set the branch weights in order to minimize the price paid when the
+  // gate is turned off, allowing the default enablement of this
+  // instrumentation with as little of a performance cost as possible
+  auto Weights = MDBuilder(*C).createBranchWeights(1, 100000);
+  return SplitBlockAndInsertIfThen(FunctionGateCmp, IP, false, Weights);
+}
+
 bool ModuleSanitizerCoverage::InjectCoverage(Function &F,
                                              ArrayRef<BasicBlock *> AllBlocks,
+                                             Value *&FunctionGateCmp,
                                              bool IsLeafFunc) {
   if (AllBlocks.empty()) return false;
   CreateFunctionLocalArrays(F, AllBlocks);
-  Value *FunctionGateCmp = nullptr;
   for (size_t i = 0, N = AllBlocks.size(); i < N; i++)
     InjectCoverageAtBlock(F, *AllBlocks[i], i, FunctionGateCmp, IsLeafFunc);
   return true;
@@ -855,7 +878,8 @@ void ModuleSanitizerCoverage::InjectCoverageForIndirectCalls(
 //      {NumCases, ValueSizeInBits, Case0Value, Case1Value, Case2Value, ... })
 
 void ModuleSanitizerCoverage::InjectTraceForSwitch(
-    Function &, ArrayRef<Instruction *> SwitchTraceTargets) {
+    Function &F, ArrayRef<Instruction *> SwitchTraceTargets,
+    Value *&FunctionGateCmp) {
   for (auto *I : SwitchTraceTargets) {
     if (SwitchInst *SI = dyn_cast<SwitchInst>(I)) {
       InstrumentationIRBuilder IRB(I);
@@ -886,7 +910,13 @@ void ModuleSanitizerCoverage::InjectTraceForSwitch(
           *CurModule, ArrayOfInt64Ty, false, GlobalVariable::InternalLinkage,
           ConstantArray::get(ArrayOfInt64Ty, Initializers),
           "__sancov_gen_cov_switch_values");
-      IRB.CreateCall(SanCovTraceSwitchFunction, {Cond, GV});
+      if (Options.GatedCallbacks) {
+        auto GateBranch = CreateGateBranch(F, FunctionGateCmp, I);
+        IRBuilder<> GateIRB(GateBranch);
+        GateIRB.CreateCall(SanCovTraceSwitchFunction, {Cond, GV});
+      } else {
+        IRB.CreateCall(SanCovTraceSwitchFunction, {Cond, GV});
+      }
     }
   }
 }
@@ -950,7 +980,8 @@ void ModuleSanitizerCoverage::InjectTraceForLoadsAndStores(
 }
 
 void ModuleSanitizerCoverage::InjectTraceForCmp(
-    Function &, ArrayRef<Instruction *> CmpTraceTargets) {
+    Function &F, ArrayRef<Instruction *> CmpTraceTargets,
+    Value *&FunctionGateCmp) {
   for (auto *I : CmpTraceTargets) {
     if (ICmpInst *ICMP = dyn_cast<ICmpInst>(I)) {
       InstrumentationIRBuilder IRB(ICMP);
@@ -978,8 +1009,15 @@ void ModuleSanitizerCoverage::InjectTraceForCmp(
       }
 
       auto Ty = Type::getIntNTy(*C, TypeSize);
-      IRB.CreateCall(CallbackFunc, {IRB.CreateIntCast(A0, Ty, true),
-              IRB.CreateIntCast(A1, Ty, true)});
+      if (Options.GatedCallbacks) {
+        auto GateBranch = CreateGateBranch(F, FunctionGateCmp, I);
+        IRBuilder<> GateIRB(GateBranch);
+        GateIRB.CreateCall(CallbackFunc, {GateIRB.CreateIntCast(A0, Ty, true),
+                                          GateIRB.CreateIntCast(A1, Ty, true)});
+      } else {
+        IRB.CreateCall(CallbackFunc, {IRB.CreateIntCast(A0, Ty, true),
+                                      IRB.CreateIntCast(A1, Ty, true)});
+      }
     }
   }
 }
@@ -1013,19 +1051,10 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
                       ConstantInt::get(IntptrTy, Idx * 4)),
         PtrTy);
     if (Options.GatedCallbacks) {
-      if (!FunctionGateCmp) {
-        // Create this in the entry block
-        assert(IsEntryBB);
-        FunctionGateCmp = CreateFunctionLocalGateCmp(IRB);
-      }
-      // Set the branch weights in order to minimize the price paid when the
-      // gate is turned off, allowing the default enablement of this
-      // instrumentation with as little of a performance cost as possible
-      auto Weights = MDBuilder(*C).createBranchWeights(1, 100000);
-      auto ThenTerm =
-          SplitBlockAndInsertIfThen(FunctionGateCmp, &*IP, false, Weights);
-      IRBuilder<> ThenIRB(ThenTerm);
-      ThenIRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
+      Instruction *I = &*IP;
+      auto GateBranch = CreateGateBranch(F, FunctionGateCmp, I);
+      IRBuilder<> GateIRB(GateBranch);
+      GateIRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
     } else {
       IRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
     }
