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

Block a context root from being imported by its callers. 

Suppose that happened. Its caller - usually a message pump - inlines its copy of the root. Then it (the root) and whatever it calls will be the non-contextually optimized callee versions.

Author: mtrofin

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,11 @@ 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. "Variants" refers to copies due to importing - we want there to be
+  // just one instance of this function.
+  DenseSet<ValueInfo> Roots;
 
   void
   computeImportForModule(const GVSummaryMapTy &DefinedGVSummaries,
@@ -783,12 +789,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 +895,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;

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"