[SampleFDO] Support enabling sample loader pass in O0 mode

llvm/include/llvm/Transforms/IPO/SampleProfile.h

@@ -41,7 +41,9 @@ 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,
+      bool UseFlattenedProfile = false);
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 
@@ -50,6 +52,8 @@ class SampleProfileLoaderPass : public PassInfoMixin<SampleProfileLoaderPass> {
   std::string ProfileRemappingFileName;
   const ThinOrFullLTOPhase LTOPhase;
   IntrusiveRefCntPtr<vfs::FileSystem> FS;
+  bool DisableSampleProfileInlining;
+  bool UseFlattenedProfile;
 };
 
 } // end namespace llvm

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -2162,6 +2162,19 @@ PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
   if (PGOOpt && PGOOpt->DebugInfoForProfiling)
     MPM.addPass(createModuleToFunctionPassAdaptor(AddDiscriminatorsPass()));
 
+  if (PGOOpt && PGOOpt->Action == PGOOptions::SampleUse) {
+    // Explicitly disable sample loader inlining and use flattened profile in O0
+    // pipeline.
+    MPM.addPass(SampleProfileLoaderPass(PGOOpt->ProfileFile,
+                                        PGOOpt->ProfileRemappingFile,
+                                        ThinOrFullLTOPhase::None, nullptr,
+                                        /*DisableSampleProfileInlining=*/true,
+                                        /*UseFlattenedProfile=*/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, Phase);
 
   // Build a minimal pipeline based on the semantics required by LLVM,

llvm/lib/Transforms/IPO/SampleProfile.cpp

@@ -469,7 +469,8 @@ class SampleProfileLoader final : public SampleProfileLoaderBaseImpl<Function> {
       std::function<AssumptionCache &(Function &)> GetAssumptionCache,
       std::function<TargetTransformInfo &(Function &)> GetTargetTransformInfo,
       std::function<const TargetLibraryInfo &(Function &)> GetTLI,
-      LazyCallGraph &CG)
+      LazyCallGraph &CG, bool DisableSampleProfileInlining,
+      bool UseFlattenedProfile)
       : SampleProfileLoaderBaseImpl(std::string(Name), std::string(RemapName),
                                     std::move(FS)),
         GetAC(std::move(GetAssumptionCache)),
@@ -478,7 +479,9 @@ class SampleProfileLoader final : public SampleProfileLoaderBaseImpl<Function> {
         AnnotatedPassName(AnnotateSampleProfileInlinePhase
                               ? llvm::AnnotateInlinePassName(InlineContext{
                                     LTOPhase, InlinePass::SampleProfileInliner})
-                              : CSINLINE_DEBUG) {}
+                              : CSINLINE_DEBUG),
+        DisableSampleProfileInlining(DisableSampleProfileInlining),
+        UseFlattenedProfile(UseFlattenedProfile) {}
 
   bool doInitialization(Module &M, FunctionAnalysisManager *FAM = nullptr);
   bool runOnModule(Module &M, ModuleAnalysisManager *AM,
@@ -592,6 +595,10 @@ class SampleProfileLoader final : public SampleProfileLoaderBaseImpl<Function> {
   // attribute.
   bool ProfAccForSymsInList;
 
+  bool DisableSampleProfileInlining;
+
+  bool UseFlattenedProfile;
+
   // External inline advisor used to replay inline decision from remarks.
   std::unique_ptr<InlineAdvisor> ExternalInlineAdvisor;
 
@@ -919,7 +926,7 @@ bool SampleProfileLoader::tryPromoteAndInlineCandidate(
     Function &F, InlineCandidate &Candidate, uint64_t SumOrigin, uint64_t &Sum,
     SmallVector<CallBase *, 8> *InlinedCallSite) {
   // Bail out early if sample-loader inliner is disabled.
-  if (DisableSampleLoaderInlining)
+  if (DisableSampleProfileInlining)
     return false;
 
   // Bail out early if MaxNumPromotions is zero.
@@ -1230,7 +1237,7 @@ bool SampleProfileLoader::tryInlineCandidate(
     InlineCandidate &Candidate, SmallVector<CallBase *, 8> *InlinedCallSites) {
   // Do not attempt to inline a candidate if
   // --disable-sample-loader-inlining is true.
-  if (DisableSampleLoaderInlining)
+  if (DisableSampleProfileInlining)
     return false;
 
   CallBase &CB = *Candidate.CallInstr;
@@ -1974,6 +1981,13 @@ bool SampleProfileLoader::doInitialization(Module &M,
 
   PSL = Reader->getProfileSymbolList();
 
+  if (DisableSampleLoaderInlining.getNumOccurrences())
+    DisableSampleProfileInlining = DisableSampleLoaderInlining;
+
+  if (UseFlattenedProfile)
+    ProfileConverter::flattenProfile(Reader->getProfiles(),
+                                     Reader->profileIsCS());
+
   // While profile-sample-accurate is on, ignore symbol list.
   ProfAccForSymsInList =
       ProfileAccurateForSymsInList && PSL && !ProfileSampleAccurate;
@@ -2304,9 +2318,12 @@ 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,
+    bool UseFlattenedProfile)
     : ProfileFileName(File), ProfileRemappingFileName(RemappingFile),
-      LTOPhase(LTOPhase), FS(std::move(FS)) {}
+      LTOPhase(LTOPhase), FS(std::move(FS)),
+      DisableSampleProfileInlining(DisableSampleProfileInlining),
+      UseFlattenedProfile(UseFlattenedProfile) {}
 
 PreservedAnalyses SampleProfileLoaderPass::run(Module &M,
                                                ModuleAnalysisManager &AM) {
@@ -2331,7 +2348,8 @@ PreservedAnalyses SampleProfileLoaderPass::run(Module &M,
       ProfileFileName.empty() ? SampleProfileFile : ProfileFileName,
       ProfileRemappingFileName.empty() ? SampleProfileRemappingFile
                                        : ProfileRemappingFileName,
-      LTOPhase, FS, GetAssumptionCache, GetTTI, GetTLI, CG);
+      LTOPhase, FS, GetAssumptionCache, GetTTI, GetTLI, CG,
+      DisableSampleProfileInlining, UseFlattenedProfile);
   if (!SampleLoader.doInitialization(M, &FAM))
     return PreservedAnalyses::all();
 

llvm/test/Other/new-pm-pgo-O0.ll

@@ -9,15 +9,19 @@
 ; 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
 ; USE_POST_LINK-NOT: Running pass: PGOInstrumentationUse
 ; 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
 }