[ctxprof] Option to move a whole tree to its own module

[mtrofin]

Modules may contain a mix of functions that participate or don't participate in callgraphs covered by a contextual profile. We currently have been importing all the functions under a context root in the module defining that root, but if the other functions there are covered by flat profiles, the result is difficult to reason about.

This patch allows moving everything under a context root (and that root) in its own module. For now, we expect a module with a filename matching the GUID of the function be present in the set of modules known by the linker. This mechanism can be improved in a later patch.

Subsequent patches will handle implementing "move" instead of "import" semantics for the root function (because we want to make sure only one version of the root exists - so the optimizations we perform are actually the ones being observed at runtime).

[mtrofin]

• [ctxprof] Support for "move" semantics for the contextual root #134192
• [ctxprof] Don't import roots elsewhere #134012
• [ctxprof][nfc] Make computeImportForFunction a member of ModuleImportsManager #134011
• [ctxprof] Option to move a whole tree to its own module #133992 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-lto

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

Changes

Modules may contain a mix of functions that participate or don't participate in callgraphs covered by a contextual profile. We currently have been importing all the functions under a context root in the module defining that root, but if the other functions there are covered by flat profiles, the result is difficult to reason about.

This patch allows moving everything under a context root (and that root) in its own module. For now, we expect a module with a filename matching the GUID of the function be present in the set of modules known by the linker. This mechanism can be improved in a later patch.

Subsequent patches will handle implementing "move" instead of "import" semantics for the root function (because we want to make sure only one version of the root exists - so the optimizations we perform are actually the ones being observed at runtime).

---

Full diff: https://github.com/llvm/llvm-project/pull/133992.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/FunctionImport.cpp (+26-5)
• (added) llvm/test/ThinLTO/X86/ctxprof-separate-module.ll (+56)

diff --git a/llvm/lib/Transforms/IPO/FunctionImport.cpp b/llvm/lib/Transforms/IPO/FunctionImport.cpp
index f0daf1a558316..faa052bb4d5b6 100644
--- a/llvm/lib/Transforms/IPO/FunctionImport.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionImport.cpp
@@ -38,6 +38,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO/Internalize.h"
@@ -175,6 +176,12 @@ static cl::opt<std::string> WorkloadDefinitions(
 
 extern cl::opt<std::string> UseCtxProfile;
 
+static cl::opt<bool> CtxprofMoveRootsToOwnModule(
+    "thinlto-move-ctxprof-trees",
+    cl::desc("Move contextual profiling roots and the graphs under them in "
+             "their own module."),
+    cl::Hidden, cl::init(false));
+
 namespace llvm {
 extern cl::opt<bool> EnableMemProfContextDisambiguation;
 }
@@ -535,7 +542,13 @@ class WorkloadImportsManager : public ModuleImportsManager {
   computeImportForModule(const GVSummaryMapTy &DefinedGVSummaries,
                          StringRef ModName,
                          FunctionImporter::ImportMapTy &ImportList) override {
-    auto SetIter = Workloads.find(ModName);
+    StringRef Filename = ModName;
+    if (CtxprofMoveRootsToOwnModule) {
+      Filename = sys::path::filename(ModName);
+      Filename = Filename.substr(0, Filename.find_last_of('.'));
+    }
+    auto SetIter = Workloads.find(Filename);
+
     if (SetIter == Workloads.end()) {
       LLVM_DEBUG(dbgs() << "[Workload] " << ModName
                         << " does not contain the root of any context.\n");
@@ -748,10 +761,18 @@ class WorkloadImportsManager : public ModuleImportsManager {
                           << RootVI.getSummaryList().size() << ". Skipping.\n");
         continue;
       }
-      StringRef RootDefiningModule =
-          RootVI.getSummaryList().front()->modulePath();
-      LLVM_DEBUG(dbgs() << "[Workload] Root defining module for " << RootGuid
-                        << " is : " << RootDefiningModule << "\n");
+      std::string RootDefiningModule =
+          RootVI.getSummaryList().front()->modulePath().str();
+      if (CtxprofMoveRootsToOwnModule) {
+        RootDefiningModule = std::to_string(RootGuid);
+        LLVM_DEBUG(
+            dbgs() << "[Workload] Moving " << RootGuid
+                   << " to a module with the filename without extension : "
+                   << RootDefiningModule << "\n");
+      } else {
+        LLVM_DEBUG(dbgs() << "[Workload] Root defining module for " << RootGuid
+                          << " is : " << RootDefiningModule << "\n");
+      }
       auto &Set = Workloads[RootDefiningModule];
       Root.getContainedGuids(ContainedGUIDs);
       for (auto Guid : ContainedGUIDs)
diff --git a/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll b/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll
new file mode 100644
index 0000000000000..c7891d336cc89
--- /dev/null
+++ b/llvm/test/ThinLTO/X86/ctxprof-separate-module.ll
@@ -0,0 +1,56 @@
+; 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
+; used when computing the GUIDs of internal linkage symbols.
+;
+; Set up
+; RUN: rm -rf %t
+; RUN: mkdir -p %t
+; RUN: split-file %s %t
+;
+; 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/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:  -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
+;
+;
+; CHECK: m1_f1()
+; CHECK: m2_f1()
+;
+;--- ctxprof.yaml
+Contexts: 
+  -
+    Guid: 6019442868614718803
+    TotalRootEntryCount: 5
+    Counters: [1]
+    Callsites:  
+      - -
+          Guid: 15593096274670919754
+          Counters: [1]
+
+;--- m1.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"
+
+define dso_local void @m1_f1() {
+  ret void
+}
+
+;--- m2.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"
+
+define dso_local void @m2_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"

[snehasish]

Did you mean to add reviewers instead of assignees? This PR does not show up on my graphite dashboard unless I'm added to reviewers.

[mtrofin]

ya. Managed to fall in the confusion of the UI

[snehasish]

lgtm

[mtrofin]

Merge activity

• Apr 2, 9:11 PM EDT: A user started a stack merge that includes this pull request via Graphite.
• Apr 2, 9:13 PM EDT: Graphite rebased this pull request as part of a merge.
• Apr 2, 9:15 PM EDT: A user merged this pull request with Graphite.