[LTO] Update splitCodeGen to take a reference to the module. (NFC)

splitCodeGen does not need to take ownership of the module, as it
currently clones the original module for each split operation.

There is an ~4 year old fixme to change that, but until this is
addressed, the function can just take a reference to the module.

This makes the transition of LTOCodeGenerator to use LTOBackend a bit
easier, because under some circumstances, LTOCodeGenerator needs to
write the original module back after codegen.

Reviewed By: tejohnson

Differential Revision: https://reviews.llvm.org/D95222

include/llvm/CodeGen/ParallelCG.h

@@ -32,14 +32,11 @@ class raw_pwrite_stream;
 ///
 /// Writes bitcode for individual partitions into output streams in BCOSs, if
 /// BCOSs is not empty.
-///
-/// \returns M if OSs.size() == 1, otherwise returns std::unique_ptr<Module>().
-std::unique_ptr<Module>
-splitCodeGen(std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs,
-             ArrayRef<llvm::raw_pwrite_stream *> BCOSs,
-             const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
-             CodeGenFileType FileType = CGFT_ObjectFile,
-             bool PreserveLocals = false);
+void splitCodeGen(
+    Module &M, ArrayRef<raw_pwrite_stream *> OSs,
+    ArrayRef<llvm::raw_pwrite_stream *> BCOSs,
+    const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
+    CodeGenFileType FileType = CGFT_ObjectFile, bool PreserveLocals = false);
 
 } // namespace llvm
 

include/llvm/LTO/LTOBackend.h

@@ -42,8 +42,8 @@ bool opt(const Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
 /// Runs a regular LTO backend. The regular LTO backend can also act as the
 /// regular LTO phase of ThinLTO, which may need to access the combined index.
 Error backend(const Config &C, AddStreamFn AddStream,
-              unsigned ParallelCodeGenParallelismLevel,
-              std::unique_ptr<Module> M, ModuleSummaryIndex &CombinedIndex);
+              unsigned ParallelCodeGenParallelismLevel, Module &M,
+              ModuleSummaryIndex &CombinedIndex);
 
 /// Runs a ThinLTO backend.
 Error thinBackend(const Config &C, unsigned Task, AddStreamFn AddStream,

include/llvm/Transforms/Utils/SplitModule.h

@@ -33,7 +33,7 @@ class Module;
 /// - Internal symbols defined in module-level inline asm should be visible to
 ///   each partition.
 void SplitModule(
-    std::unique_ptr<Module> M, unsigned N,
+    Module &M, unsigned N,
     function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
     bool PreserveLocals = false);
 

lib/CodeGen/ParallelCG.cpp

@@ -36,8 +36,8 @@ static void codegen(Module *M, llvm::raw_pwrite_stream &OS,
   CodeGenPasses.run(*M);
 }
 
-std::unique_ptr<Module> llvm::splitCodeGen(
-    std::unique_ptr<Module> M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
+void llvm::splitCodeGen(
+    Module &M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
     ArrayRef<llvm::raw_pwrite_stream *> BCOSs,
     const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
     CodeGenFileType FileType, bool PreserveLocals) {
@@ -45,9 +45,9 @@ std::unique_ptr<Module> llvm::splitCodeGen(
 
   if (OSs.size() == 1) {
     if (!BCOSs.empty())
-      WriteBitcodeToFile(*M, *BCOSs[0]);
-    codegen(M.get(), *OSs[0], TMFactory, FileType);
-    return M;
+      WriteBitcodeToFile(M, *BCOSs[0]);
+    codegen(&M, *OSs[0], TMFactory, FileType);
+    return;
   }
 
   // Create ThreadPool in nested scope so that threads will be joined
@@ -57,7 +57,7 @@ std::unique_ptr<Module> llvm::splitCodeGen(
     int ThreadCount = 0;
 
     SplitModule(
-        std::move(M), OSs.size(),
+        M, OSs.size(),
         [&](std::unique_ptr<Module> MPart) {
           // We want to clone the module in a new context to multi-thread the
           // codegen. We do it by serializing partition modules to bitcode
@@ -95,6 +95,4 @@ std::unique_ptr<Module> llvm::splitCodeGen(
         },
         PreserveLocals);
   }
-
-  return {};
 }

lib/LTO/LTO.cpp

@@ -1101,9 +1101,9 @@ Error LTO::runRegularLTO(AddStreamFn AddStream) {
   }
 
   if (!RegularLTO.EmptyCombinedModule || Conf.AlwaysEmitRegularLTOObj) {
-    if (Error Err = backend(
-            Conf, AddStream, RegularLTO.ParallelCodeGenParallelismLevel,
-            std::move(RegularLTO.CombinedModule), ThinLTO.CombinedIndex))
+    if (Error Err =
+            backend(Conf, AddStream, RegularLTO.ParallelCodeGenParallelismLevel,
+                    *RegularLTO.CombinedModule, ThinLTO.CombinedIndex))
       return Err;
   }
 

lib/LTO/LTOBackend.cpp

@@ -456,8 +456,7 @@ static void codegen(const Config &Conf, TargetMachine *TM,
 
 static void splitCodeGen(const Config &C, TargetMachine *TM,
                          AddStreamFn AddStream,
-                         unsigned ParallelCodeGenParallelismLevel,
-                         std::unique_ptr<Module> Mod,
+                         unsigned ParallelCodeGenParallelismLevel, Module &Mod,
                          const ModuleSummaryIndex &CombinedIndex) {
   ThreadPool CodegenThreadPool(
       heavyweight_hardware_concurrency(ParallelCodeGenParallelismLevel));
@@ -465,7 +464,7 @@ static void splitCodeGen(const Config &C, TargetMachine *TM,
   const Target *T = &TM->getTarget();
 
   SplitModule(
-      std::move(Mod), ParallelCodeGenParallelismLevel,
+      Mod, ParallelCodeGenParallelismLevel,
       [&](std::unique_ptr<Module> MPart) {
         // We want to clone the module in a new context to multi-thread the
         // codegen. We do it by serializing partition modules to bitcode
@@ -532,27 +531,26 @@ Error lto::finalizeOptimizationRemarks(
 }
 
 Error lto::backend(const Config &C, AddStreamFn AddStream,
-                   unsigned ParallelCodeGenParallelismLevel,
-                   std::unique_ptr<Module> Mod,
+                   unsigned ParallelCodeGenParallelismLevel, Module &Mod,
                    ModuleSummaryIndex &CombinedIndex) {
-  Expected<const Target *> TOrErr = initAndLookupTarget(C, *Mod);
+  Expected<const Target *> TOrErr = initAndLookupTarget(C, Mod);
   if (!TOrErr)
     return TOrErr.takeError();
 
-  std::unique_ptr<TargetMachine> TM = createTargetMachine(C, *TOrErr, *Mod);
+  std::unique_ptr<TargetMachine> TM = createTargetMachine(C, *TOrErr, Mod);
 
   if (!C.CodeGenOnly) {
-    if (!opt(C, TM.get(), 0, *Mod, /*IsThinLTO=*/false,
+    if (!opt(C, TM.get(), 0, Mod, /*IsThinLTO=*/false,
              /*ExportSummary=*/&CombinedIndex, /*ImportSummary=*/nullptr,
              /*CmdArgs*/ std::vector<uint8_t>()))
       return Error::success();
   }
 
   if (ParallelCodeGenParallelismLevel == 1) {
-    codegen(C, TM.get(), AddStream, 0, *Mod, CombinedIndex);
+    codegen(C, TM.get(), AddStream, 0, Mod, CombinedIndex);
   } else {
-    splitCodeGen(C, TM.get(), AddStream, ParallelCodeGenParallelismLevel,
-                 std::move(Mod), CombinedIndex);
+    splitCodeGen(C, TM.get(), AddStream, ParallelCodeGenParallelismLevel, Mod,
+                 CombinedIndex);
   }
   return Error::success();
 }

lib/LTO/LTOCodeGenerator.cpp

@@ -616,14 +616,9 @@ bool LTOCodeGenerator::compileOptimized(ArrayRef<raw_pwrite_stream *> Out) {
   // for splitting
   restoreLinkageForExternals();
 
-  // Do code generation. We need to preserve the module in case the client calls
-  // writeMergedModules() after compilation, but we only need to allow this at
-  // parallelism level 1. This is achieved by having splitCodeGen return the
-  // original module at parallelism level 1 which we then assign back to
-  // MergedModule.
-  MergedModule = splitCodeGen(std::move(MergedModule), Out, {},
-                              [&]() { return createTargetMachine(); }, FileType,
-                              ShouldRestoreGlobalsLinkage);
+  splitCodeGen(
+      *MergedModule, Out, {}, [&]() { return createTargetMachine(); }, FileType,
+      ShouldRestoreGlobalsLinkage);
 
   // If statistics were requested, save them to the specified file or
   // print them out after codegen.

lib/Transforms/Utils/SplitModule.cpp

@@ -95,13 +95,12 @@ static void addAllGlobalValueUsers(ClusterMapType &GVtoClusterMap,
 // globalized.
 // Try to balance pack those partitions into N files since this roughly equals
 // thread balancing for the backend codegen step.
-static void findPartitions(Module *M, ClusterIDMapType &ClusterIDMap,
+static void findPartitions(Module &M, ClusterIDMapType &ClusterIDMap,
                            unsigned N) {
   // At this point module should have the proper mix of globals and locals.
   // As we attempt to partition this module, we must not change any
   // locals to globals.
-  LLVM_DEBUG(dbgs() << "Partition module with (" << M->size()
-                    << ")functions\n");
+  LLVM_DEBUG(dbgs() << "Partition module with (" << M.size() << ")functions\n");
   ClusterMapType GVtoClusterMap;
   ComdatMembersType ComdatMembers;
 
@@ -144,9 +143,9 @@ static void findPartitions(Module *M, ClusterIDMapType &ClusterIDMap,
       addAllGlobalValueUsers(GVtoClusterMap, &GV, &GV);
   };
 
-  llvm::for_each(M->functions(), recordGVSet);
-  llvm::for_each(M->globals(), recordGVSet);
-  llvm::for_each(M->aliases(), recordGVSet);
+  llvm::for_each(M.functions(), recordGVSet);
+  llvm::for_each(M.globals(), recordGVSet);
+  llvm::for_each(M.aliases(), recordGVSet);
 
   // Assigned all GVs to merged clusters while balancing number of objects in
   // each.
@@ -247,31 +246,32 @@ static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {
 }
 
 void llvm::SplitModule(
-    std::unique_ptr<Module> M, unsigned N,
+    Module &M, unsigned N,
     function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback,
     bool PreserveLocals) {
   if (!PreserveLocals) {
-    for (Function &F : *M)
+    for (Function &F : M)
       externalize(&F);
-    for (GlobalVariable &GV : M->globals())
+    for (GlobalVariable &GV : M.globals())
       externalize(&GV);
-    for (GlobalAlias &GA : M->aliases())
+    for (GlobalAlias &GA : M.aliases())
       externalize(&GA);
-    for (GlobalIFunc &GIF : M->ifuncs())
+    for (GlobalIFunc &GIF : M.ifuncs())
       externalize(&GIF);
   }
 
   // This performs splitting without a need for externalization, which might not
   // always be possible.
   ClusterIDMapType ClusterIDMap;
-  findPartitions(M.get(), ClusterIDMap, N);
+  findPartitions(M, ClusterIDMap, N);
 
   // FIXME: We should be able to reuse M as the last partition instead of
-  // cloning it.
+  // cloning it. Note that the callers at the moment expect the module to
+  // be preserved, so will need some adjustments as well.
   for (unsigned I = 0; I < N; ++I) {
     ValueToValueMapTy VMap;
     std::unique_ptr<Module> MPart(
-        CloneModule(*M, VMap, [&](const GlobalValue *GV) {
+        CloneModule(M, VMap, [&](const GlobalValue *GV) {
           if (ClusterIDMap.count(GV))
             return (ClusterIDMap[GV] == I);
           else

tools/llvm-split/llvm-split.cpp

@@ -52,25 +52,28 @@ int main(int argc, char **argv) {
   }
 
   unsigned I = 0;
-  SplitModule(std::move(M), NumOutputs, [&](std::unique_ptr<Module> MPart) {
-    std::error_code EC;
-    std::unique_ptr<ToolOutputFile> Out(
-        new ToolOutputFile(OutputFilename + utostr(I++), EC, sys::fs::OF_None));
-    if (EC) {
-      errs() << EC.message() << '\n';
-      exit(1);
-    }
+  SplitModule(
+      *M, NumOutputs,
+      [&](std::unique_ptr<Module> MPart) {
+        std::error_code EC;
+        std::unique_ptr<ToolOutputFile> Out(new ToolOutputFile(
+            OutputFilename + utostr(I++), EC, sys::fs::OF_None));
+        if (EC) {
+          errs() << EC.message() << '\n';
+          exit(1);
+        }
 
-    if (verifyModule(*MPart, &errs())) {
-      errs() << "Broken module!\n";
-      exit(1);
-    }
+        if (verifyModule(*MPart, &errs())) {
+          errs() << "Broken module!\n";
+          exit(1);
+        }
 
-    WriteBitcodeToFile(*MPart, Out->os());
+        WriteBitcodeToFile(*MPart, Out->os());
 
-    // Declare success.
-    Out->keep();
-  }, PreserveLocals);
+        // Declare success.
+        Out->keep();
+      },
+      PreserveLocals);
 
   return 0;
 }