[InstrPGO] Instrument sampling profile based cold function v2

clang/include/clang/Driver/Options.td

@@ -1785,6 +1785,9 @@ 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_sample_cold_function : Flag<["-"], "fprofile-sample-cold-function">, 
+    Group<f_Group>, Visibility<[ClangOption, CLOption]>,
+    HelpText<"Generate instrumented code to cold functions guided by sampling-based profile">;
 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/ToolChains/Clang.cpp

@@ -649,13 +649,24 @@ static void addPGOAndCoverageFlags(const ToolChain &TC, Compilation &C,
     }
   }
 
+  auto *SampleColdArg =
+      Args.getLastArg(options::OPT_fprofile_sample_cold_function);
+
   Arg *PGOGenArg = nullptr;
-  if (PGOGenerateArg) {
+  if (SampleColdArg) {
+    SmallString<128> Path(
+        PGOGenerateArg->getOption().matches(options::OPT_fprofile_generate_EQ)
+            ? PGOGenerateArg->getValue()
+            : "");
+    Path.append("default_%m.profraw");
+    CmdArgs.push_back("-mllvm");
+    CmdArgs.push_back(Args.MakeArgString(
+        Twine("--instrument-sample-cold-function-path=") + Path));
+  } else if (PGOGenerateArg) {
     assert(!CSPGOGenerateArg);
     PGOGenArg = PGOGenerateArg;
     CmdArgs.push_back("-fprofile-instrument=llvm");
-  }
-  if (CSPGOGenerateArg) {
+  } else if (CSPGOGenerateArg) {
     assert(!PGOGenerateArg);
     PGOGenArg = CSPGOGenerateArg;
     CmdArgs.push_back("-fprofile-instrument=csllvm");
@@ -7024,7 +7035,9 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
           options::OPT_fcs_profile_generate,
           options::OPT_fcs_profile_generate_EQ, options::OPT_fprofile_use,
           options::OPT_fprofile_use_EQ);
-      if (PGOArg)
+      auto *SampleColdArg =
+          Args.getLastArg(options::OPT_fprofile_sample_cold_function);
+      if (PGOArg && !SampleColdArg)
         D.Diag(diag::err_drv_argument_not_allowed_with)
             << "SampleUse with PGO options";
 

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -296,7 +296,13 @@ 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> InstrumentSampleColdFuncPath(
+    "instrument-sample-cold-function-path", cl::init(""),
+    cl::desc("File path for instrumenting sampling PGO guided cold functions"),
+    cl::Hidden);
+
 extern cl::opt<std::string> UseCtxProfile;
+extern cl::opt<bool> InstrumentColdFunction;
 
 namespace llvm {
 extern cl::opt<bool> EnableMemProfContextDisambiguation;
@@ -1119,6 +1125,17 @@ PassBuilder::buildModuleSimplificationPipeline(OptimizationLevel Level,
       // removed.
       MPM.addPass(
           PGOIndirectCallPromotion(true /* IsInLTO */, true /* SamplePGO */));
+
+    if (InstrumentSampleColdFuncPath.getNumOccurrences() &&
+        Phase != ThinOrFullLTOPhase::ThinLTOPostLink) {
+      assert(!InstrumentSampleColdFuncPath.empty() &&
+             "File path is requeired for instrumentation generation");
+      InstrumentColdFunction = true;
+      addPreInlinerPasses(MPM, Level, Phase);
+      addPGOInstrPasses(MPM, Level, /* RunProfileGen */ true,
+                        /* IsCS */ false, /* AtomicCounterUpdate */ false,
+                        InstrumentSampleColdFuncPath, "", PGOOpt->FS);
+    }
   }
 
   // Try to perform OpenMP specific optimizations on the module. This is a

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -319,6 +319,16 @@ static cl::opt<unsigned> PGOFunctionCriticalEdgeThreshold(
     cl::desc("Do not instrument functions with the number of critical edges "
              " greater than this threshold."));
 
+cl::opt<bool> InstrumentColdFunction(
+    "instrument-cold-function", cl::init(false), cl::Hidden,
+    cl::desc("Instrument cold functions (currently only used under sampling "
+             " PGO pipeline))"));
+
+static cl::opt<uint64_t> InstrumentColdFuncMaxEntryCount(
+    "instrument-cold-function-max-entry-count", cl::init(0), cl::Hidden,
+    cl::desc("When using --instrument-cold-function, skip instrumenting the "
+             "function whose entry count is above the given value"));
+
 extern cl::opt<unsigned> MaxNumVTableAnnotations;
 
 namespace llvm {
@@ -1891,6 +1901,10 @@ static bool skipPGOGen(const Function &F) {
     return true;
   if (F.getInstructionCount() < PGOFunctionSizeThreshold)
     return true;
+  if (InstrumentColdFunction &&
+      (!F.getEntryCount() ||
+       F.getEntryCount()->getCount() > InstrumentColdFuncMaxEntryCount))
+    return true;
   return false;
 }
 

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

@@ -0,0 +1,24 @@
+; RUN: opt < %s  --passes=pgo-instr-gen -instrument-cold-function -S  | FileCheck --check-prefixes=COLD %s
+; RUN: opt < %s  --passes=pgo-instr-gen -instrument-cold-function -instrument-cold-function-max-entry-count=1 -S  | FileCheck --check-prefixes=ENTRY-COUNT %s
+
+; COLD: call void @llvm.instrprof.increment(ptr @__profn_foo, i64  [[#]], i32 1, i32 0)
+; COLD-NOT: __profn_main
+
+; ENTRY-COUNT: call void @llvm.instrprof.increment(ptr @__profn_foo, i64  [[#]], i32 1, i32 0)
+; ENTRY-COUNT: call void @llvm.instrprof.increment(ptr @__profn_main, 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 @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}