[llvm] [SampleFDO] Support enabling sample loader pass in O0 mode (PR #113985)

[Lei Wang via llvm-commits]

https://github.com/wlei-llvm updated https://github.com/llvm/llvm-project/pull/113985

>From 9b41a3a3321c4ce0f4a7ca9f6548c601d6ecd949 Mon Sep 17 00:00:00 2001
From: wlei <wlei at fb.com>
Date: Mon, 28 Oct 2024 16:34:17 -0700
Subject: [PATCH] [SampleFDO] Support running sample loader in O0 mode

---
 llvm/include/llvm/Transforms/IPO/SampleProfile.h |  4 +++-
 llvm/lib/Passes/PassBuilderPipelines.cpp         | 11 +++++++++++
 llvm/lib/Transforms/IPO/SampleProfile.cpp        |  8 ++++++--
 llvm/test/Other/new-pm-pgo-O0.ll                 |  6 +++++-
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/Transforms/IPO/SampleProfile.h b/llvm/include/llvm/Transforms/IPO/SampleProfile.h
index e94f6ba55cd0dd..5adda3a0a16f3a 100644
--- a/llvm/include/llvm/Transforms/IPO/SampleProfile.h
+++ b/llvm/include/llvm/Transforms/IPO/SampleProfile.h
@@ -41,7 +41,8 @@ class SampleProfileLoaderPass : public PassInfoMixin<SampleProfileLoaderPass> {
   SampleProfileLoaderPass(
       std::string File = "", std::string RemappingFile = "",
       ThinOrFullLTOPhase LTOPhase = ThinOrFullLTOPhase::None,
-      IntrusiveRefCntPtr<vfs::FileSystem> FS = nullptr);
+      IntrusiveRefCntPtr<vfs::FileSystem> FS = nullptr,
+      bool DisableSampleProfileInlining = false);
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 
@@ -50,6 +51,7 @@ class SampleProfileLoaderPass : public PassInfoMixin<SampleProfileLoaderPass> {
   std::string ProfileRemappingFileName;
   const ThinOrFullLTOPhase LTOPhase;
   IntrusiveRefCntPtr<vfs::FileSystem> FS;
+  bool DisableSampleProfileInlining;
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 466fbcd7bb7703..8b8a6142d7205f 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -2153,6 +2153,17 @@ PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
   if (PGOOpt && PGOOpt->DebugInfoForProfiling)
     MPM.addPass(createModuleToFunctionPassAdaptor(AddDiscriminatorsPass()));
 
+  if (PGOOpt && PGOOpt->Action == PGOOptions::SampleUse) {
+    // Explicitly disable sample loader inlining in O0 pipeline.
+    MPM.addPass(SampleProfileLoaderPass(PGOOpt->ProfileFile,
+                                        PGOOpt->ProfileRemappingFile,
+                                        ThinOrFullLTOPhase::None, nullptr,
+                                        /*DisableSampleProfileInlining*/ true));
+    // Cache ProfileSummaryAnalysis once to avoid the potential need to insert
+    // RequireAnalysisPass for PSI before subsequent non-module passes.
+    MPM.addPass(RequireAnalysisPass<ProfileSummaryAnalysis, Module>());
+  }
+
   invokePipelineEarlySimplificationEPCallbacks(MPM, Level);
 
   // Build a minimal pipeline based on the semantics required by LLVM,
diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 5d1ec7248ae679..9cd025efb0a5d8 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -2305,9 +2305,10 @@ bool SampleProfileLoader::runOnFunction(Function &F, ModuleAnalysisManager *AM)
 }
 SampleProfileLoaderPass::SampleProfileLoaderPass(
     std::string File, std::string RemappingFile, ThinOrFullLTOPhase LTOPhase,
-    IntrusiveRefCntPtr<vfs::FileSystem> FS)
+    IntrusiveRefCntPtr<vfs::FileSystem> FS, bool DisableSampleProfileInlining)
     : ProfileFileName(File), ProfileRemappingFileName(RemappingFile),
-      LTOPhase(LTOPhase), FS(std::move(FS)) {}
+      LTOPhase(LTOPhase), FS(std::move(FS)),
+      DisableSampleProfileInlining(DisableSampleProfileInlining) {}
 
 PreservedAnalyses SampleProfileLoaderPass::run(Module &M,
                                                ModuleAnalysisManager &AM) {
@@ -2326,6 +2327,9 @@ PreservedAnalyses SampleProfileLoaderPass::run(Module &M,
 
   if (!FS)
     FS = vfs::getRealFileSystem();
+  if (!DisableSampleLoaderInlining.getNumOccurrences() &&
+      DisableSampleProfileInlining)
+    DisableSampleLoaderInlining = true;
   LazyCallGraph &CG = AM.getResult<LazyCallGraphAnalysis>(M);
 
   SampleProfileLoader SampleLoader(
diff --git a/llvm/test/Other/new-pm-pgo-O0.ll b/llvm/test/Other/new-pm-pgo-O0.ll
index d7a6a03b8e44e3..d4f662fb25ace7 100644
--- a/llvm/test/Other/new-pm-pgo-O0.ll
+++ b/llvm/test/Other/new-pm-pgo-O0.ll
@@ -9,8 +9,9 @@
 ; RUN:     |FileCheck %s --check-prefixes=USE_POST_LINK,USE
 ; RUN: opt -debug-pass-manager -passes='lto<O0>' -pgo-kind=pgo-instr-use-pipeline -profile-file='%t.profdata' %s 2>&1 \
 ; RUN:     |FileCheck %s --check-prefixes=USE_POST_LINK,USE
+; RUN: opt -debug-pass-manager -passes='default<O0>' -pgo-kind=pgo-sample-use-pipeline -profile-file='%S/Inputs/new-pm-pgo.prof' %s 2>&1 \
+; RUN:     |FileCheck %s --check-prefixes=SAMPLE_USE
 
-;
 ; GEN: Running pass: PGOInstrumentationGen
 ; USE_DEFAULT: Running pass: PGOInstrumentationUse
 ; USE_PRE_LINK: Running pass: PGOInstrumentationUse
@@ -18,6 +19,9 @@
 ; USE-NOT: Running pass: PGOIndirectCallPromotion
 ; USE-NOT: Running pass: PGOMemOPSizeOpt
 
+; SAMPLE_USE: Running pass: AddDiscriminatorsPass
+; SAMPLE_USE: Running pass: SampleProfileLoaderPass
+
 define void @foo() {
   ret void
 }