[SanitizerCoverage] Add an option to gate the invocation of the tracing callbacks

Implement -sanitizer-coverage-gated-trace-callbacks to gate the invocation of the tracing callbacks based on the value of a global variable, which is stored in a specific section.
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.

Note: will add additional support for trace-cmp in a follow up PR.

Patch by Filippo Bigarella

rdar://101626834

Author: thetruestblue

clang/test/CodeGen/sanitize-coverage-gated-callbacks.c

@@ -0,0 +1,42 @@
+// 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: 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
+
+// Verify that we do not emit the __sancov_gate section for "plain" trace-pc-guard
+// GATED: section "__DATA,__sancov_gate"
+// 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
+
+int x[10];
+
+// CHECK: define{{.*}} void @foo
+void foo(int n, int m) {
+  // COM: Verify that we're emitting the call to __sanitizer_cov_trace_pc_guard upon
+  // COM: checking the value of __sancov_should_track.
+  // GATED: [[VAL:%.*]] = load i64, {{.*}}@__sancov_should_track
+  // GATED-NOT: [[VAL:%.*]] = load i64, i64* @__sancov_should_track
+  // GATED-NEXT: [[CMP:%.*]] = icmp ne i64 [[VAL]], 0
+  // 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
+  // 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
+  // GATED: [[WEIGHTS]] = !{!"branch_weights", i32 1, i32 100000}
+
+  // COM: With the non-gated instrumentation, we should not emit the
+  // COM: __sancov_should_track global.
+  // PLAIN-NOT: __sancov_should_track
+  // But we should still be emitting the calls to the callback.
+  // PLAIN: call void @__sanitizer_cov_trace_pc_guard
+  if (n) {
+    x[n] = 42;
+    if (m) {
+      x[m] = 41;
+    }
+  }
+}

llvm/include/llvm/Transforms/Utils/Instrumentation.h

@@ -161,6 +161,7 @@ struct SanitizerCoverageOptions {
   bool TraceLoads = false;
   bool TraceStores = false;
   bool CollectControlFlow = false;
+  bool GatedCallbacks = false;
 
   SanitizerCoverageOptions() = default;
 };

llvm/lib/Transforms/Instrumentation/SanitizerCoverage.cpp

@@ -16,6 +16,7 @@
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/EHPersonalities.h"
@@ -28,6 +29,8 @@
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
+#include "llvm/IR/ValueSymbolTable.h"
+#include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/SpecialCaseList.h"
 #include "llvm/Support/VirtualFileSystem.h"
@@ -82,8 +85,10 @@ const char SanCovCountersSectionName[] = "sancov_cntrs";
 const char SanCovBoolFlagSectionName[] = "sancov_bools";
 const char SanCovPCsSectionName[] = "sancov_pcs";
 const char SanCovCFsSectionName[] = "sancov_cfs";
+const char SanCovCallbackGateSectionName[] = "sancov_gate";
 
 const char SanCovLowestStackName[] = "__sancov_lowest_stack";
+const char SanCovCallbackGateName[] = "__sancov_should_track";
 
 static cl::opt<int> ClCoverageLevel(
     "sanitizer-coverage-level",
@@ -152,6 +157,12 @@ static cl::opt<bool>
     ClCollectCF("sanitizer-coverage-control-flow",
                 cl::desc("collect control flow for each function"), cl::Hidden);
 
+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::Hidden, cl::init(false));
+
 namespace {
 
 SanitizerCoverageOptions getOptions(int LegacyCoverageLevel) {
@@ -194,6 +205,7 @@ SanitizerCoverageOptions OverrideFromCL(SanitizerCoverageOptions Options) {
   Options.StackDepth |= ClStackDepth;
   Options.TraceLoads |= ClLoadTracing;
   Options.TraceStores |= ClStoreTracing;
+  Options.GatedCallbacks |= ClGatedCallbacks;
   if (!Options.TracePCGuard && !Options.TracePC &&
       !Options.Inline8bitCounters && !Options.StackDepth &&
       !Options.InlineBoolFlag && !Options.TraceLoads && !Options.TraceStores)
@@ -239,8 +251,9 @@ class ModuleSanitizerCoverage {
                                                     const char *Section);
   GlobalVariable *CreatePCArray(Function &F, ArrayRef<BasicBlock *> AllBlocks);
   void CreateFunctionLocalArrays(Function &F, ArrayRef<BasicBlock *> AllBlocks);
+  Value *CreateFunctionLocalGateCmp(IRBuilder<> &IRB);
   void InjectCoverageAtBlock(Function &F, BasicBlock &BB, size_t Idx,
-                             bool IsLeafFunc = true);
+                             Value *&FunctionGateCmp, bool IsLeafFunc = true);
   Function *CreateInitCallsForSections(Module &M, const char *CtorName,
                                        const char *InitFunctionName, Type *Ty,
                                        const char *Section);
@@ -265,6 +278,7 @@ class ModuleSanitizerCoverage {
   FunctionCallee SanCovTraceGepFunction;
   FunctionCallee SanCovTraceSwitchFunction;
   GlobalVariable *SanCovLowestStack;
+  GlobalVariable *SanCovCallbackGate;
   Type *PtrTy, *IntptrTy, *Int64Ty, *Int32Ty, *Int16Ty, *Int8Ty, *Int1Ty;
   Module *CurModule;
   std::string CurModuleUniqueId;
@@ -478,6 +492,23 @@ bool ModuleSanitizerCoverage::instrumentModule() {
   if (Options.StackDepth && !SanCovLowestStack->isDeclaration())
     SanCovLowestStack->setInitializer(Constant::getAllOnesValue(IntptrTy));
 
+  if (Options.GatedCallbacks) {
+    if (!Options.TracePCGuard) {
+      C->emitError(StringRef("'") + ClGatedCallbacks.ArgStr +
+                   "' is only supported with trace-pc-guard");
+      return true;
+    }
+
+    SanCovCallbackGate = cast<GlobalVariable>(
+        M.getOrInsertGlobal(SanCovCallbackGateName, Int64Ty));
+    SanCovCallbackGate->setSection(
+        getSectionName(SanCovCallbackGateSectionName));
+    SanCovCallbackGate->setInitializer(Constant::getNullValue(Int64Ty));
+    SanCovCallbackGate->setLinkage(GlobalVariable::LinkOnceAnyLinkage);
+    SanCovCallbackGate->setVisibility(GlobalVariable::HiddenVisibility);
+    appendToCompilerUsed(M, SanCovCallbackGate);
+  }
+
   SanCovTracePC = M.getOrInsertFunction(SanCovTracePCName, VoidTy);
   SanCovTracePCGuard =
       M.getOrInsertFunction(SanCovTracePCGuardName, VoidTy, PtrTy);
@@ -777,13 +808,22 @@ void ModuleSanitizerCoverage::CreateFunctionLocalArrays(
     FunctionPCsArray = CreatePCArray(F, AllBlocks);
 }
 
+Value *ModuleSanitizerCoverage::CreateFunctionLocalGateCmp(IRBuilder<> &IRB) {
+  auto Load = IRB.CreateLoad(Int64Ty, SanCovCallbackGate);
+  Load->setNoSanitizeMetadata();
+  auto Cmp = IRB.CreateIsNotNull(Load);
+  Cmp->setName("sancov gate cmp");
+  return Cmp;
+}
+
 bool ModuleSanitizerCoverage::InjectCoverage(Function &F,
                                              ArrayRef<BasicBlock *> AllBlocks,
                                              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, IsLeafFunc);
+    InjectCoverageAtBlock(F, *AllBlocks[i], i, FunctionGateCmp, IsLeafFunc);
   return true;
 }
 
@@ -946,6 +986,7 @@ void ModuleSanitizerCoverage::InjectTraceForCmp(
 
 void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
                                                     size_t Idx,
+                                                    Value *&FunctionGateCmp,
                                                     bool IsLeafFunc) {
   BasicBlock::iterator IP = BB.getFirstInsertionPt();
   bool IsEntryBB = &BB == &F.getEntryBlock();
@@ -971,7 +1012,23 @@ void ModuleSanitizerCoverage::InjectCoverageAtBlock(Function &F, BasicBlock &BB,
         IRB.CreateAdd(IRB.CreatePointerCast(FunctionGuardArray, IntptrTy),
                       ConstantInt::get(IntptrTy, Idx * 4)),
         PtrTy);
-    IRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
+    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();
+    } else {
+      IRB.CreateCall(SanCovTracePCGuard, GuardPtr)->setCannotMerge();
+    }
   }
   if (Options.Inline8bitCounters) {
     auto CounterPtr = IRB.CreateGEP(