[ctxprof] Option to move a whole tree to its own module (#133992)

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).

Author: mtrofin

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,14 @@ 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);
+      // Drop the file extension.
+      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 +762,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)

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"