[SanitizerCoverage] Add gated tracing callbacks support to trace-cmp (#113227)

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

Author: thetruestblue

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")))``
 ===========================================================================

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 or trace-cmp
 
 int x[10];
 
@@ -23,6 +25,11 @@ void foo(int n, int m) {
   // GATED-NEXT: br i1 [[CMP]], label %[[L_TRUE:.*]], label %[[L_FALSE:.*]], !prof [[WEIGHTS:!.+]]
   // GATED: [[L_TRUE]]:
   // GATED-NEXT:   call void @__sanitizer_cov_trace_pc_guard
+  // COM: Check the trace-cmp instrumentation of the if (n) branch
+  // GATEDCMP: [[OPERAND:%.*]] = load i32, {{.*}}
+  // GATEDCMP-NEXT: br i1 [[CMP]], label %[[L_TRUE_1:.*]], label %[[L_FALSE_1:.*]]
+  // GATEDCMP: [[L_TRUE_1]]:
+  // GATEDCMP-NEXT:   call void @__sanitizer_cov_trace_const_cmp4(i32 0, i32 [[OPERAND]])
   // GATED:   br i1 [[CMP]], label %[[L_TRUE_2:.*]], label %[[L_FALSE_2:.*]]
   // GATED: [[L_TRUE_2]]:
   // GATED-NEXT:   call void @__sanitizer_cov_trace_pc_guard
@@ -33,10 +40,12 @@ void foo(int n, int m) {
   // PLAIN-NOT: __sancov_should_track
   // But we should still be emitting the calls to the callback.
   // PLAIN: call void @__sanitizer_cov_trace_pc_guard
+  // PLAINCMP: [[OPERAND:%.*]] = load i32, {{.*}}
+  // PLAINCMP-NEXT: call void @__sanitizer_cov_trace_const_cmp4(i32 0, i32 [[OPERAND]])
   if (n) {
     x[n] = 42;
     if (m) {
       x[m] = 41;
     }
   }
-}
+}

llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp

@@ -158,8 +158,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 {
@@ -234,17 +234,19 @@ 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);
+                      Value *&FunctionGateCmp, bool IsLeafFunc);
   GlobalVariable *CreateFunctionLocalArrayInSection(size_t NumElements,
                                                     Function &F, Type *Ty,
                                                     const char *Section);
@@ -494,9 +496,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;
     }
 
@@ -725,10 +727,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);
@@ -837,10 +840,10 @@ Instruction *ModuleSanitizerCoverage::CreateGateBranch(Function &F,
 
 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;
@@ -874,7 +877,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);
@@ -905,7 +909,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});
+      }
     }
   }
 }
@@ -969,7 +979,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);
@@ -997,8 +1008,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)});
+      }
     }
   }
 }