[ctxprof] Support for "move" semantics for the contextual root (#134192)

This PR finishes what PR #133992 started.

Author: mtrofin

llvm/include/llvm/Transforms/Utils/FunctionImportUtils.h

@@ -97,29 +97,14 @@ class FunctionImportGlobalProcessing {
   /// linkage for a required promotion of a local to global scope.
   GlobalValue::LinkageTypes getLinkage(const GlobalValue *SGV, bool DoPromote);
 
+  /// The symbols with these names are moved to a different module and should be
+  /// promoted to external linkage where they are defined.
+  DenseSet<GlobalValue::GUID> SymbolsToMove;
+
 public:
   FunctionImportGlobalProcessing(Module &M, const ModuleSummaryIndex &Index,
                                  SetVector<GlobalValue *> *GlobalsToImport,
-                                 bool ClearDSOLocalOnDeclarations)
-      : M(M), ImportIndex(Index), GlobalsToImport(GlobalsToImport),
-        ClearDSOLocalOnDeclarations(ClearDSOLocalOnDeclarations) {
-    // If we have a ModuleSummaryIndex but no function to import,
-    // then this is the primary module being compiled in a ThinLTO
-    // backend compilation, and we need to see if it has functions that
-    // may be exported to another backend compilation.
-    if (!GlobalsToImport)
-      HasExportedFunctions = ImportIndex.hasExportedFunctions(M);
-
-#ifndef NDEBUG
-    SmallVector<GlobalValue *, 4> Vec;
-    // First collect those in the llvm.used set.
-    collectUsedGlobalVariables(M, Vec, /*CompilerUsed=*/false);
-    // Next collect those in the llvm.compiler.used set.
-    collectUsedGlobalVariables(M, Vec, /*CompilerUsed=*/true);
-    Used = {llvm::from_range, Vec};
-#endif
-  }
-
+                                 bool ClearDSOLocalOnDeclarations);
   void run();
 };
 

llvm/lib/Transforms/IPO/FunctionImport.cpp

@@ -182,6 +182,15 @@ static cl::opt<bool> CtxprofMoveRootsToOwnModule(
              "their own module."),
     cl::Hidden, cl::init(false));
 
+cl::list<GlobalValue::GUID> MoveSymbolGUID(
+    "thinlto-move-symbols",
+    cl::desc(
+        "Move the symbols with the given name. This will delete these symbols "
+        "wherever they are originally defined, and make sure their "
+        "linkage is External where they are imported. It is meant to be "
+        "used with the name of contextual profiling roots."),
+    cl::Hidden);
+
 namespace llvm {
 extern cl::opt<bool> EnableMemProfContextDisambiguation;
 }
@@ -1859,6 +1868,15 @@ Expected<bool> FunctionImporter::importFunctions(
   LLVM_DEBUG(dbgs() << "Starting import for Module "
                     << DestModule.getModuleIdentifier() << "\n");
   unsigned ImportedCount = 0, ImportedGVCount = 0;
+  // Before carrying out any imports, see if this module defines functions in
+  // MoveSymbolGUID. If it does, delete them here (but leave the declaration).
+  // The function will be imported elsewhere, as extenal linkage, and the
+  // destination doesn't yet have its definition.
+  DenseSet<GlobalValue::GUID> MoveSymbolGUIDSet;
+  MoveSymbolGUIDSet.insert_range(MoveSymbolGUID);
+  for (auto &F : DestModule)
+    if (!F.isDeclaration() && MoveSymbolGUIDSet.contains(F.getGUID()))
+      F.deleteBody();
 
   IRMover Mover(DestModule);
 

llvm/lib/Transforms/Utils/FunctionImportUtils.cpp

@@ -24,6 +24,31 @@ static cl::opt<bool> UseSourceFilenameForPromotedLocals(
              "This requires that the source filename has a unique name / "
              "path to avoid name collisions."));
 
+extern cl::list<GlobalValue::GUID> MoveSymbolGUID;
+
+FunctionImportGlobalProcessing::FunctionImportGlobalProcessing(
+    Module &M, const ModuleSummaryIndex &Index,
+    SetVector<GlobalValue *> *GlobalsToImport, bool ClearDSOLocalOnDeclarations)
+    : M(M), ImportIndex(Index), GlobalsToImport(GlobalsToImport),
+      ClearDSOLocalOnDeclarations(ClearDSOLocalOnDeclarations) {
+  // If we have a ModuleSummaryIndex but no function to import,
+  // then this is the primary module being compiled in a ThinLTO
+  // backend compilation, and we need to see if it has functions that
+  // may be exported to another backend compilation.
+  if (!GlobalsToImport)
+    HasExportedFunctions = ImportIndex.hasExportedFunctions(M);
+
+#ifndef NDEBUG
+  SmallVector<GlobalValue *, 4> Vec;
+  // First collect those in the llvm.used set.
+  collectUsedGlobalVariables(M, Vec, /*CompilerUsed=*/false);
+  // Next collect those in the llvm.compiler.used set.
+  collectUsedGlobalVariables(M, Vec, /*CompilerUsed=*/true);
+  Used = {llvm::from_range, Vec};
+#endif
+  SymbolsToMove.insert_range(MoveSymbolGUID);
+}
+
 /// Checks if we should import SGV as a definition, otherwise import as a
 /// declaration.
 bool FunctionImportGlobalProcessing::doImportAsDefinition(
@@ -147,7 +172,9 @@ FunctionImportGlobalProcessing::getLinkage(const GlobalValue *SGV,
     // and/or optimization, but are turned into declarations later
     // during the EliminateAvailableExternally pass.
     if (doImportAsDefinition(SGV) && !isa<GlobalAlias>(SGV))
-      return GlobalValue::AvailableExternallyLinkage;
+      return SymbolsToMove.contains(SGV->getGUID())
+                 ? GlobalValue::ExternalLinkage
+                 : GlobalValue::AvailableExternallyLinkage;
     // An imported external declaration stays external.
     return SGV->getLinkage();
 

llvm/test/ThinLTO/X86/ctxprof-separate-module.ll

@@ -22,15 +22,31 @@
 ; 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
+
+; also add the move semantics for the root:
+; RUN: llvm-lto2 run %t/m1.bc %t/m2.bc %t/m3.bc %t/6019442868614718803.bc -thinlto-move-ctxprof-trees \
+; RUN:  -thinlto-move-symbols=6019442868614718803 \
+; RUN:  -o %t/result-with-move.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:  -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
+; RUN: llvm-dis %t/result-with-move.o.1.3.import.bc -o - | FileCheck %s --check-prefix=WITHMOVE-SRC
+; RUN: llvm-dis %t/result-with-move.o.4.3.import.bc -o - | FileCheck %s --check-prefix=WITHMOVE-DEST
+; RUN: llvm-dis %t/result.o.1.3.import.bc -o - | FileCheck %s --check-prefix=WITHOUTMOVE-SRC
 ;
-;
-; CHECK: m1_f1()
-; CHECK: m2_f1()
+; CHECK: define available_externally void @m1_f1()
+; CHECK: define available_externally void @m2_f1()
 ; ABSENT: declare void @m1_f1()
 ; ABSENT-MSG: Skipping over 6019442868614718803 because its import is handled in a different module.
 ;
+; WITHMOVE-SRC: declare dso_local void @m1_f1
+; WITHMOVE-DEST: define dso_local void @m1_f1
+; WITHOUTMOVE-SRC: define dso_local void @m1_f1
 ;--- ctxprof.yaml
 Contexts: 
   -