[InstrPGO] Avoid using global variable to fix potential data race (#114364)

In #109837, it sets a global
variable(`PGOInstrumentColdFunctionOnly`) in PassBuilderPipelines.cpp
which introduced a data race detected by TSan. To fix this, I decouple
the flag setting, the flags are now set
separately(`instrument-cold-function-only-path` is required to be used
with `--pgo-instrument-cold-function-only`).

Author: wlei-llvm

clang/include/clang/Driver/Options.td

@@ -1786,6 +1786,12 @@ defm debug_info_for_profiling : BoolFOption<"debug-info-for-profiling",
   PosFlag<SetTrue, [], [ClangOption, CC1Option],
           "Emit extra debug info to make sample profile more accurate">,
   NegFlag<SetFalse>>;
+def fprofile_generate_cold_function_coverage : Flag<["-"], "fprofile-generate-cold-function-coverage">, 
+    Group<f_Group>, Visibility<[ClangOption, CLOption]>,
+    HelpText<"Generate instrumented code to collect coverage info for cold functions into default.profraw file (overridden by '=' form of option or LLVM_PROFILE_FILE env var)">;
+def fprofile_generate_cold_function_coverage_EQ : Joined<["-"], "fprofile-generate-cold-function-coverage=">, 
+    Group<f_Group>, Visibility<[ClangOption, CLOption]>, MetaVarName<"<directory>">,
+    HelpText<"Generate instrumented code to collect coverage info for cold functions into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)">; 
 def fprofile_instr_generate : Flag<["-"], "fprofile-instr-generate">,
     Group<f_Group>, Visibility<[ClangOption, CLOption]>,
     HelpText<"Generate instrumented code to collect execution counts into default.profraw file (overridden by '=' form of option or LLVM_PROFILE_FILE env var)">;

clang/lib/Driver/ToolChain.cpp

@@ -899,7 +899,9 @@ bool ToolChain::needsProfileRT(const ArgList &Args) {
          Args.hasArg(options::OPT_fprofile_instr_generate) ||
          Args.hasArg(options::OPT_fprofile_instr_generate_EQ) ||
          Args.hasArg(options::OPT_fcreate_profile) ||
-         Args.hasArg(options::OPT_forder_file_instrumentation);
+         Args.hasArg(options::OPT_forder_file_instrumentation) ||
+         Args.hasArg(options::OPT_fprofile_generate_cold_function_coverage) ||
+         Args.hasArg(options::OPT_fprofile_generate_cold_function_coverage_EQ);
 }
 
 bool ToolChain::needsGCovInstrumentation(const llvm::opt::ArgList &Args) {

clang/lib/Driver/ToolChains/Clang.cpp

@@ -632,6 +632,28 @@ static void addPGOAndCoverageFlags(const ToolChain &TC, Compilation &C,
     }
   }
 
+  if (auto *ColdFuncCoverageArg = Args.getLastArg(
+          options::OPT_fprofile_generate_cold_function_coverage,
+          options::OPT_fprofile_generate_cold_function_coverage_EQ)) {
+    SmallString<128> Path(
+        ColdFuncCoverageArg->getOption().matches(
+            options::OPT_fprofile_generate_cold_function_coverage_EQ)
+            ? ColdFuncCoverageArg->getValue()
+            : "");
+    llvm::sys::path::append(Path, "default_%m.profraw");
+    // FIXME: Idealy the file path should be passed through
+    // `-fprofile-instrument-path=`(InstrProfileOutput), however, this field is
+    // shared with other profile use path(see PGOOptions), we need to refactor
+    // PGOOptions to make it work.
+    CmdArgs.push_back("-mllvm");
+    CmdArgs.push_back(Args.MakeArgString(
+        Twine("--instrument-cold-function-only-path=") + Path));
+    CmdArgs.push_back("-mllvm");
+    CmdArgs.push_back("--pgo-instrument-cold-function-only");
+    CmdArgs.push_back("-mllvm");
+    CmdArgs.push_back("--pgo-function-entry-coverage");
+  }
+
   Arg *PGOGenArg = nullptr;
   if (PGOGenerateArg) {
     assert(!CSPGOGenerateArg);

clang/test/CodeGen/pgo-cold-function-coverage.c

@@ -0,0 +1,19 @@
+// Test -fprofile-generate-cold-function-coverage 
+
+// RUN: rm -rf %t && split-file %s %t
+// RUN: %clang --target=x86_64 -O2 -fprofile-generate-cold-function-coverage=/xxx/yyy/ -fprofile-sample-accurate -fprofile-sample-use=%t/pgo-cold-func.prof  -S -emit-llvm -o - %t/pgo-cold-func.c | FileCheck %s
+
+// CHECK: @__llvm_profile_filename = {{.*}} c"/xxx/yyy/default_%m.profraw\00"
+
+// CHECK: store i8 0, ptr @__profc_bar, align 1
+// CHECK-NOT: @__profc_foo 
+
+//--- pgo-cold-func.prof
+foo:1:1
+ 1: 1
+
+//--- pgo-cold-func.c
+int bar(int x) { return x;}
+int foo(int x) { 
+    return x;
+}

clang/test/Driver/fprofile-generate-cold-function-coverage.c

@@ -0,0 +1,9 @@
+// RUN: %clang -### -c -fprofile-generate-cold-function-coverage %s 2>&1 | FileCheck %s
+// CHECK: "--instrument-cold-function-only-path=default_%m.profraw" 
+// CHECK: "--pgo-instrument-cold-function-only"
+// CHECK: "--pgo-function-entry-coverage"
+// CHECK-NOT:  "-fprofile-instrument"
+// CHECK-NOT:  "-fprofile-instrument-path=
+
+// RUN: %clang -### -c -fprofile-generate-cold-function-coverage=dir %s 2>&1 | FileCheck %s --check-prefix=CHECK-EQ
+// CHECK-EQ: "--instrument-cold-function-only-path=dir{{/|\\\\}}default_%m.profraw" 

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -296,7 +296,14 @@ static cl::opt<bool> UseLoopVersioningLICM(
     "enable-loop-versioning-licm", cl::init(false), cl::Hidden,
     cl::desc("Enable the experimental Loop Versioning LICM pass"));
 
+static cl::opt<std::string> InstrumentColdFuncOnlyPath(
+    "instrument-cold-function-only-path", cl::init(""),
+    cl::desc("File path for cold function only instrumentation(requires use "
+             "with --pgo-instrument-cold-function-only)"),
+    cl::Hidden);
+
 extern cl::opt<std::string> UseCtxProfile;
+extern cl::opt<bool> PGOInstrumentColdFunctionOnly;
 
 namespace llvm {
 extern cl::opt<bool> EnableMemProfContextDisambiguation;
@@ -1183,8 +1190,16 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
   const bool IsCtxProfUse =
       !UseCtxProfile.empty() && Phase == ThinOrFullLTOPhase::ThinLTOPreLink;
 
+  assert(
+      (InstrumentColdFuncOnlyPath.empty() || PGOInstrumentColdFunctionOnly) &&
+      "--instrument-cold-function-only-path is provided but "
+      "--pgo-instrument-cold-function-only is not enabled");
+  const bool IsColdFuncOnlyInstrGen = PGOInstrumentColdFunctionOnly &&
+                                      IsPGOPreLink &&
+                                      !InstrumentColdFuncOnlyPath.empty();
+
   if (IsPGOInstrGen || IsPGOInstrUse || IsMemprofUse || IsCtxProfGen ||
-      IsCtxProfUse)
+      IsCtxProfUse || IsColdFuncOnlyInstrGen)
     addPreInlinerPasses(MPM, Level, Phase);
 
   // Add all the requested passes for instrumentation PGO, if requested.
@@ -1206,6 +1221,11 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
       return MPM;
     addPostPGOLoopRotation(MPM, Level);
     MPM.addPass(PGOCtxProfLoweringPass());
+  } else if (IsColdFuncOnlyInstrGen) {
+    addPGOInstrPasses(
+        MPM, Level, /* RunProfileGen */ true, /* IsCS */ false,
+        /* AtomicCounterUpdate */ false, InstrumentColdFuncOnlyPath,
+        /* ProfileRemappingFile */ "", IntrusiveRefCntPtr<vfs::FileSystem>());
   }
 
   if (IsPGOInstrGen || IsPGOInstrUse || IsCtxProfGen)

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -319,6 +319,20 @@ static cl::opt<unsigned> PGOFunctionCriticalEdgeThreshold(
     cl::desc("Do not instrument functions with the number of critical edges "
              " greater than this threshold."));
 
+static cl::opt<uint64_t> PGOColdInstrumentEntryThreshold(
+    "pgo-cold-instrument-entry-threshold", cl::init(0), cl::Hidden,
+    cl::desc("For cold function instrumentation, skip instrumenting functions "
+             "whose entry count is above the given value."));
+
+static cl::opt<bool> PGOTreatUnknownAsCold(
+    "pgo-treat-unknown-as-cold", cl::init(false), cl::Hidden,
+    cl::desc("For cold function instrumentation, treat count unknown(e.g. "
+             "unprofiled) functions as cold."));
+
+cl::opt<bool> PGOInstrumentColdFunctionOnly(
+    "pgo-instrument-cold-function-only", cl::init(false), cl::Hidden,
+    cl::desc("Enable cold function only instrumentation."));
+
 extern cl::opt<unsigned> MaxNumVTableAnnotations;
 
 namespace llvm {
@@ -1897,6 +1911,11 @@ static bool skipPGOGen(const Function &F) {
     return true;
   if (F.getInstructionCount() < PGOFunctionSizeThreshold)
     return true;
+  if (PGOInstrumentColdFunctionOnly) {
+    if (auto EntryCount = F.getEntryCount())
+      return EntryCount->getCount() > PGOColdInstrumentEntryThreshold;
+    return !PGOTreatUnknownAsCold;
+  }
   return false;
 }
 

llvm/test/Transforms/PGOProfile/instr-gen-cold-function.ll

@@ -0,0 +1,35 @@
+; RUN: opt < %s  --passes=pgo-instr-gen -pgo-instrument-cold-function-only -pgo-function-entry-coverage -S  | FileCheck --check-prefixes=COLD %s
+; RUN: opt < %s  --passes=pgo-instr-gen -pgo-instrument-cold-function-only -pgo-function-entry-coverage -pgo-cold-instrument-entry-threshold=1 -S  | FileCheck --check-prefixes=ENTRY-COUNT %s
+; RUN: opt < %s  --passes=pgo-instr-gen -pgo-instrument-cold-function-only -pgo-function-entry-coverage -pgo-treat-unknown-as-cold -S  | FileCheck --check-prefixes=UNKNOWN-FUNC %s
+
+; COLD: call void @llvm.instrprof.cover(ptr @__profn_foo, i64  [[#]], i32 1, i32 0)
+; COLD-NOT: __profn_main
+; COLD-NOT: __profn_bar
+
+; ENTRY-COUNT: call void @llvm.instrprof.cover(ptr @__profn_foo, i64  [[#]], i32 1, i32 0)
+; ENTRY-COUNT: call void @llvm.instrprof.cover(ptr @__profn_main, i64 [[#]], i32 1, i32 0)
+
+; UNKNOWN-FUNC: call void @llvm.instrprof.cover(ptr @__profn_bar, i64  [[#]], i32 1, i32 0)
+; UNKNOWN-FUNC: call void @llvm.instrprof.cover(ptr @__profn_foo, i64  [[#]], i32 1, i32 0)
+
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @bar() {
+entry:
+  ret void
+}
+
+define void @foo() !prof !0 {
+entry:
+  ret void
+}
+
+define i32 @main() !prof !1 {
+entry:
+  ret i32 0
+}
+
+!0 = !{!"function_entry_count", i64 0}
+!1 = !{!"function_entry_count", i64 1}