[llvm] [ctx_prof] Automatically convert available external linkage to local for modules with contextual roots (PR #109203)

[Mircea Trofin via llvm-commits]

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

>From 1ba77f56fc8dca81c5c48dd0085dcc08a5075e0a Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Wed, 18 Sep 2024 14:18:23 -0700
Subject: [PATCH] [ctx_prof] Automatically convert available external linkage
 to local for modules with contextual roots

---
 llvm/lib/Transforms/IPO/ElimAvailExtern.cpp   | 19 ++++++++++++++-----
 .../transform-to-local.ll                     | 18 ++++++++++++++++--
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp b/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp
index 2b34d3b5a56ea4..d3d27de4218c87 100644
--- a/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp
+++ b/llvm/lib/Transforms/IPO/ElimAvailExtern.cpp
@@ -14,6 +14,7 @@
 #include "llvm/Transforms/IPO/ElimAvailExtern.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/CtxProfAnalysis.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Function.h"
@@ -88,7 +89,7 @@ static void convertToLocalCopy(Module &M, Function &F) {
   ++NumConversions;
 }
 
-static bool eliminateAvailableExternally(Module &M) {
+static bool eliminateAvailableExternally(Module &M, bool Convert) {
   bool Changed = false;
 
   // Drop initializers of available externally global variables.
@@ -112,7 +113,7 @@ static bool eliminateAvailableExternally(Module &M) {
     if (F.isDeclaration() || !F.hasAvailableExternallyLinkage())
       continue;
 
-    if (ConvertToLocal)
+    if (Convert || ConvertToLocal)
       convertToLocalCopy(M, F);
     else
       deleteFunction(F);
@@ -125,8 +126,16 @@ static bool eliminateAvailableExternally(Module &M) {
 }
 
 PreservedAnalyses
-EliminateAvailableExternallyPass::run(Module &M, ModuleAnalysisManager &) {
-  if (!eliminateAvailableExternally(M))
-    return PreservedAnalyses::all();
+EliminateAvailableExternallyPass::run(Module &M, ModuleAnalysisManager &MAM) {
+  auto *CtxProf = MAM.getCachedResult<CtxProfAnalysis>(M);
+  // Convert to local instead of eliding if we use contextual profiling in this
+  // module. This is because the IPO decisions performed with contextual
+  // information will likely differ from decisions made without. For a function
+  // that's imported, its optimizations will, thus, differ, and be specialized
+  // for this contextual information. Eliding it in favor of the original would
+  // undo these optimizations.
+  if (!eliminateAvailableExternally(M, /*Convert=*/(CtxProf && !!(*CtxProf))))
+    ;
+  return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }
diff --git a/llvm/test/Transforms/EliminateAvailableExternally/transform-to-local.ll b/llvm/test/Transforms/EliminateAvailableExternally/transform-to-local.ll
index 786cc260d331c6..4908fba62e3bfe 100644
--- a/llvm/test/Transforms/EliminateAvailableExternally/transform-to-local.ll
+++ b/llvm/test/Transforms/EliminateAvailableExternally/transform-to-local.ll
@@ -1,6 +1,16 @@
 ; REQUIRES: asserts
 ; RUN: opt -passes=elim-avail-extern -avail-extern-to-local -stats -S 2>&1 < %s | FileCheck %s
+;
+; RUN: echo '[{"Guid":1234, "Counters": [1]}]' | llvm-ctxprof-util fromJSON --input=- --output=%t_profile.ctxprofdata
+;
+; Because we pass a contextual profile with a root defined in this module, we expect the outcome to be the same as-if
+; we passed -avail-extern-to-local, i.e. available_externally don't get elided and instead get converted to local linkage
+; RUN: opt -passes='assign-guid,require<ctx-prof-analysis>,elim-avail-extern' -use-ctx-profile=%t_profile.ctxprofdata -stats -S 2>&1 < %s | FileCheck %s
 
+; If the profile doesn't apply to this module, available_externally won't get converted to internal linkage, and will be
+; removed instead.
+; RUN: echo '[{"Guid":5678, "Counters": [1]}]' | llvm-ctxprof-util fromJSON --input=- --output=%t_profile_bad.ctxprofdata
+; RUN: opt -passes='assign-guid,require<ctx-prof-analysis>,elim-avail-extern' -use-ctx-profile=%t_profile_bad.ctxprofdata -stats -S 2>&1 < %s | FileCheck %s --check-prefix=NOOP
 
 declare void @call_out(ptr %fct)
 
@@ -12,13 +22,15 @@ define available_externally hidden void @g() {
   ret void
 }
 
-define void @hello(ptr %g) {
+define void @hello(ptr %g) !guid !0 {
   call void @f()
   %f = load ptr, ptr @f
   call void @call_out(ptr %f)
   ret void
 }
 
+!0 = !{i64 1234}
+
 ; CHECK: define internal void @f.__uniq.{{[0-9|a-f]*}}()
 ; CHECK: declare hidden void @g()
 ; CHECK: call void @f.__uniq.{{[0-9|a-f]*}}()
@@ -26,4 +38,6 @@ define void @hello(ptr %g) {
 ; CHECK-NEXT: call void @call_out(ptr %f)
 ; CHECK: Statistics Collected
 ; CHECK: 1 elim-avail-extern - Number of functions converted
-; CHECK: 1 elim-avail-extern - Number of functions removed
\ No newline at end of file
+; CHECK: 1 elim-avail-extern - Number of functions removed
+
+; NOOP: 2 elim-avail-extern - Number of functions removed
\ No newline at end of file