[llvm] [ctx_prof] Insert the ctx prof flattener after the module inliner (PR #107499)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/107499

>From 0c148abec93afa5f95444f4d863683b8bd453084 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Thu, 5 Sep 2024 12:52:56 -0700
Subject: [PATCH] [ctx_prof] Insert the ctx prof flattener after the module
 inliner

---
 llvm/lib/Passes/PassBuilderPipelines.cpp     | 19 ++++++++++++++-----
 llvm/lib/Transforms/IPO/ModuleInliner.cpp    |  6 ++++--
 llvm/test/Analysis/CtxProfAnalysis/inline.ll | 17 +++++++++++++++++
 llvm/test/Other/opt-hot-cold-split.ll        |  2 +-
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 7f9e1362e7ef23..4126e563619c5b 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -25,6 +25,7 @@
 #include "llvm/Analysis/ScopedNoAliasAA.h"
 #include "llvm/Analysis/TypeBasedAliasAnalysis.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
 #include "llvm/Passes/OptimizationLevel.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/CommandLine.h"
@@ -1011,6 +1012,11 @@ PassBuilder::buildModuleInlinerPipeline(OptimizationLevel Level,
   IP.EnableDeferral = false;
 
   MPM.addPass(ModuleInlinerPass(IP, UseInlineAdvisor, Phase));
+  if (!UseCtxProfile.empty() && Phase == ThinOrFullLTOPhase::ThinLTOPostLink) {
+    MPM.addPass(GlobalOptPass());
+    MPM.addPass(GlobalDCEPass());
+    MPM.addPass(PGOCtxProfFlatteningPass());
+  }
 
   MPM.addPass(createModuleToFunctionPassAdaptor(
       buildFunctionSimplificationPipeline(Level, Phase),
@@ -1734,11 +1740,14 @@ ModulePassManager PassBuilder::buildThinLTODefaultPipeline(
     MPM.addPass(GlobalDCEPass());
     return MPM;
   }
-
-  // Add the core simplification pipeline.
-  MPM.addPass(buildModuleSimplificationPipeline(
-      Level, ThinOrFullLTOPhase::ThinLTOPostLink));
-
+  if (!UseCtxProfile.empty()) {
+    MPM.addPass(
+        buildModuleInlinerPipeline(Level, ThinOrFullLTOPhase::ThinLTOPostLink));
+  } else {
+    // Add the core simplification pipeline.
+    MPM.addPass(buildModuleSimplificationPipeline(
+        Level, ThinOrFullLTOPhase::ThinLTOPostLink));
+  }
   // Now add the optimization pipeline.
   MPM.addPass(buildModuleOptimizationPipeline(
       Level, ThinOrFullLTOPhase::ThinLTOPostLink));
diff --git a/llvm/lib/Transforms/IPO/ModuleInliner.cpp b/llvm/lib/Transforms/IPO/ModuleInliner.cpp
index b7e4531c8e390d..542c319b880747 100644
--- a/llvm/lib/Transforms/IPO/ModuleInliner.cpp
+++ b/llvm/lib/Transforms/IPO/ModuleInliner.cpp
@@ -241,8 +241,10 @@ PreservedAnalyses ModuleInlinerPass::run(Module &M,
           // the post-inline cleanup and the next DevirtSCCRepeatedPass
           // iteration because the next iteration may not happen and we may
           // miss inlining it.
-          if (tryPromoteCall(*ICB))
-            NewCallee = ICB->getCalledFunction();
+          // FIXME: enable for ctxprof.
+          if (!CtxProf)
+            if (tryPromoteCall(*ICB))
+              NewCallee = ICB->getCalledFunction();
         }
         if (NewCallee)
           if (!NewCallee->isDeclaration())
diff --git a/llvm/test/Analysis/CtxProfAnalysis/inline.ll b/llvm/test/Analysis/CtxProfAnalysis/inline.ll
index 875bc4938653b9..9381418c4e3f12 100644
--- a/llvm/test/Analysis/CtxProfAnalysis/inline.ll
+++ b/llvm/test/Analysis/CtxProfAnalysis/inline.ll
@@ -31,6 +31,23 @@
 ; CHECK-NEXT:    %call2 = call i32 @a(i32 %x) #1
 ; CHECK-NEXT:    br label %exit
 
+; Make sure the postlink thinlto pipeline is aware of ctxprof
+; RUN: opt -passes='thinlto<O2>' -use-ctx-profile=%t/profile.ctxprofdata \
+; RUN:   %t/module.ll -S -o - | FileCheck %s --check-prefix=PIPELINE
+
+; PIPELINE-LABEL: define i32 @entrypoint
+; PIPELINE-SAME: !prof ![[ENTRYPOINT_COUNT:[0-9]+]]
+; PIPELINE-LABEL: loop.i:
+; PIPELINE:         br i1 %cond.i, label %loop.i, label %exit, !prof ![[LOOP_BW_INL:[0-9]+]]
+; PIPELINE-LABEL: define i32 @a
+; PIPELINE-LABEL: loop:
+; PIPELINE:         br i1 %cond, label %loop, label %exit, !prof ![[LOOP_BW_ORIG:[0-9]+]]
+
+; PIPELINE: ![[ENTRYPOINT_COUNT]] = !{!"function_entry_count", i64 10}
+; These are the weights of the inlined @a, where the counters were 2, 100 (2 for entry, 100 for loop)
+; PIPELINE: ![[LOOP_BW_INL]] = !{!"branch_weights", i32 98, i32 2}
+; These are the weights of the un-inlined @a, where the counters were 8, 500 (8 for entry, 500 for loop)
+; PIPELINE: ![[LOOP_BW_ORIG]] = !{!"branch_weights", i32 492, i32 8}
 
 ;--- module.ll
 define i32 @entrypoint(i32 %x) !guid !0 {
diff --git a/llvm/test/Other/opt-hot-cold-split.ll b/llvm/test/Other/opt-hot-cold-split.ll
index 21c713d35bb746..cd290dcc306570 100644
--- a/llvm/test/Other/opt-hot-cold-split.ll
+++ b/llvm/test/Other/opt-hot-cold-split.ll
@@ -2,7 +2,7 @@
 ; RUN: opt -mtriple=x86_64-- -hot-cold-split=true -passes='lto-pre-link<Os>' -debug-pass-manager < %s -o /dev/null 2>&1 | FileCheck %s -check-prefix=LTO-PRELINK-Os
 ; RUN: opt -mtriple=x86_64-- -hot-cold-split=true -passes='thinlto-pre-link<Os>' -debug-pass-manager < %s -o /dev/null 2>&1 | FileCheck %s -check-prefix=THINLTO-PRELINK-Os
 ; RUN: opt -mtriple=x86_64-- -hot-cold-split=true -passes='lto<Os>' -debug-pass-manager < %s -o /dev/null 2>&1 | FileCheck %s -check-prefix=LTO-POSTLINK-Os
-; RUN: opt -mtriple=x86_64-- -hot-cold-split=true -passes='thinlto<Os>' -debug-pass-manager < %s -o /dev/null 2>&1 | FileCheck %s -check-prefix=THINLTO-POSTLINK-Os
+; RUN: opt -mtriple=x86_64-- -hot-cold-split=true -LINK-Os
 
 ; REQUIRES: asserts