[llvm] [ctxprof] Don't import roots elsewhere (PR #134012)

[Mircea Trofin via llvm-commits]

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

>From 319acce9ffa7e6c265516362c60efa145f83bfe2 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 1 Apr 2025 16:49:47 -0700
Subject: [PATCH] [ctxprof] Don't import roots elsewhere

---
 llvm/lib/Transforms/IPO/FunctionImport.cpp    | 17 +++++++++++++
 .../ThinLTO/X86/ctxprof-separate-module.ll    | 24 ++++++++++++++++---
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/FunctionImport.cpp b/llvm/lib/Transforms/IPO/FunctionImport.cpp
index 43807a8feb36e..3d9fb7b12b5d5 100644
--- a/llvm/lib/Transforms/IPO/FunctionImport.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionImport.cpp
@@ -516,6 +516,7 @@ class ModuleImportsManager {
       const ModuleSummaryIndex &Index,
       DenseMap<StringRef, FunctionImporter::ExportSetTy> *ExportLists = nullptr)
       : IsPrevailing(IsPrevailing), Index(Index), ExportLists(ExportLists) {}
+  virtual bool canImport(ValueInfo VI) { return true; }
 
 public:
   virtual ~ModuleImportsManager() = default;
@@ -544,6 +545,10 @@ class WorkloadImportsManager : public ModuleImportsManager {
   // determine if a module's import list should be done by the base
   // ModuleImportsManager or by us.
   StringMap<DenseSet<ValueInfo>> Workloads;
+  // Track the roots to avoid importing them due to other callers. We want there
+  // to be only one variant, for which we optimize according to the contextual
+  // profile.
+  DenseSet<ValueInfo> Roots;
 
   void
   computeImportForModule(const GVSummaryMapTy &DefinedGVSummaries,
@@ -783,12 +788,15 @@ class WorkloadImportsManager : public ModuleImportsManager {
       }
       auto &Set = Workloads[RootDefiningModule];
       Root.getContainedGuids(ContainedGUIDs);
+      Roots.insert(RootVI);
       for (auto Guid : ContainedGUIDs)
         if (auto VI = Index.getValueInfo(Guid))
           Set.insert(VI);
     }
   }
 
+  bool canImport(ValueInfo VI) override { return !Roots.contains(VI); }
+
 public:
   WorkloadImportsManager(
       function_ref<bool(GlobalValue::GUID, const GlobalValueSummary *)>
@@ -886,6 +894,15 @@ void ModuleImportsManager::computeImportForFunction(
       continue;
     }
 
+    if (!canImport(VI)) {
+      LLVM_DEBUG(
+          dbgs() << "Skipping over " << VI.getGUID()
+                 << " because its import is handled in a different module.");
+      assert(VI.getSummaryList().size() == 1 &&
+             "The root was expected to be an external symbol");
+      continue;
+    }
+
     auto GetBonusMultiplier = [](CalleeInfo::HotnessType Hotness) -> float {
       if (Hotness == CalleeInfo::HotnessType::Hot)
         return ImportHotMultiplier;
diff --git a/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll b/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll
index c7891d336cc89..391fe21a1b638 100644
--- a/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll
+++ b/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll
@@ -1,3 +1,4 @@
+; REQUIRES: asserts
 ; Test workload based importing via -thinlto-pgo-ctx-prof with moving the whole
 ; graph to a new module.
 ; Use external linkage symbols so we don't depend on module paths which are
@@ -10,19 +11,25 @@
 ;
 ; RUN: opt -module-summary -passes=assign-guid,ctx-instr-gen %t/m1.ll -o %t/m1.bc
 ; RUN: opt -module-summary -passes=assign-guid,ctx-instr-gen %t/m2.ll -o %t/m2.bc
+; RUN: opt -module-summary -passes=assign-guid,ctx-instr-gen %t/m3.ll -o %t/m3.bc
 ; RUN: opt -module-summary -passes=assign-guid,ctx-instr-gen %t/6019442868614718803.ll -o %t/6019442868614718803.bc
 
 ; RUN: llvm-ctxprof-util fromYAML --input %t/ctxprof.yaml --output %t/ctxprof.bitstream
-; RUN: llvm-lto2 run %t/m1.bc %t/m2.bc %t/6019442868614718803.bc -thinlto-move-ctxprof-trees \
+; RUN: llvm-lto2 run %t/m1.bc %t/m2.bc %t/m3.bc %t/6019442868614718803.bc -thinlto-move-ctxprof-trees \
 ; RUN:  -o %t/result.o -save-temps \
 ; RUN:  -use-ctx-profile=%t/ctxprof.bitstream \
 ; RUN:  -r %t/m1.bc,m1_f1,plx \
-; RUN:  -r %t/m2.bc,m2_f1,plx
-; RUN: llvm-dis %t/result.o.3.3.import.bc -o - | FileCheck %s
+; RUN:  -r %t/m2.bc,m2_f1,plx \
+; RUN:  -r %t/m3.bc,m1_f1 \
+; RUN:  -r %t/m3.bc,m3_f1,plx -debug-only=function-import 2>&1 | FileCheck %s --check-prefix=ABSENT-MSG
+; RUN: llvm-dis %t/result.o.4.3.import.bc -o - | FileCheck %s
+; RUN: llvm-dis %t/result.o.3.3.import.bc -o - | FileCheck %s --check-prefix=ABSENT
 ;
 ;
 ; CHECK: m1_f1()
 ; CHECK: m2_f1()
+; ABSENT: declare void @m1_f1()
+; ABSENT-MSG: Skipping over 6019442868614718803 because its import is handled in a different module.
 ;
 ;--- ctxprof.yaml
 Contexts: 
@@ -51,6 +58,17 @@ define dso_local void @m2_f1() {
   ret void
 }
 
+;--- m3.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+declare void @m1_f1()
+
+define dso_local void @m3_f1() {
+  call void @m1_f1()
+  ret void
+}
+
 ;--- 6019442868614718803.ll
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-linux-gnu"