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

[fhahn]

Summary

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.

[tejohnson]

Where/how is the returned module used in that case? How do you avoid having it accessed later in the parallel code gen case?

[fhahn]

In D95222#2516413, @tejohnson wrote:

Where/how is the returned module used in that case? How do you avoid having it accessed later in the parallel code gen case?

We should be able to use the returned module in the ParallelismLevel == 1 case, where we just pass a reference of the module to codegen. I think this behavior is needed for D94487 (llvm/lib/LTO/LTOCodeGenerator.cpp around line 608). This is a bit fragile/messy unfortunately, but I think we need it to preserve the existing libLTO behavior. It is also similar to the behavior of the previously used splitCodeGen https://github.com/llvm/llvm-project/blob/master/llvm/lib/CodeGen/ParallelCG.cpp#L49. But perhaps there's a better way to factor this out.

[tejohnson]

In D95222#2516447, @fhahn wrote:

In D95222#2516413, @tejohnson wrote:

Where/how is the returned module used in that case? How do you avoid having it accessed later in the parallel code gen case?

We should be able to use the returned module in the ParallelismLevel == 1 case, where we just pass a reference of the module to codegen. I think this behavior is needed for D94487 (llvm/lib/LTO/LTOCodeGenerator.cpp around line 608). This is a bit fragile/messy unfortunately, but I think we need it to preserve the existing libLTO behavior. It is also similar to the behavior of the previously used splitCodeGen https://github.com/llvm/llvm-project/blob/master/llvm/lib/CodeGen/ParallelCG.cpp#L49. But perhaps there's a better way to factor this out.

I see where it is resetting MergedModule at line 608 in that patch, but not where it is used afterwards. It seems a bit fragile to me that the MergedModule would be valid when parallel code gen is off and not when it is. Does anything break if MergedModule is not set again after calling backend()?

[fhahn]

In D95222#2516743, @tejohnson wrote:

In D95222#2516447, @fhahn wrote:

In D95222#2516413, @tejohnson wrote:

Where/how is the returned module used in that case? How do you avoid having it accessed later in the parallel code gen case?

We should be able to use the returned module in the ParallelismLevel == 1 case, where we just pass a reference of the module to codegen. I think this behavior is needed for D94487 (llvm/lib/LTO/LTOCodeGenerator.cpp around line 608). This is a bit fragile/messy unfortunately, but I think we need it to preserve the existing libLTO behavior. It is also similar to the behavior of the previously used splitCodeGen https://github.com/llvm/llvm-project/blob/master/llvm/lib/CodeGen/ParallelCG.cpp#L49. But perhaps there's a better way to factor this out.

I see where it is resetting MergedModule at line 608 in that patch, but not where it is used afterwards. It seems a bit fragile to me that the MergedModule would be valid when parallel code gen is off and not when it is.

Add to that that MergedModule is invalid while codegen is running....

Does anything break if MergedModule is not set again after calling backend()?

From the comments, it seems that is something that the API guarantees? @steven_wu do you know if writing back MergedModule is strictly necessary?

[steven_wu]

MergedModule is needed to support legacy API lto_codegen_write_merged_modules. You don't technical get the module back in memory, you just need to write to a file.

I forgot if there are APIs needed to work for merged module in memory but I will just implement that as first write to a file then read it back.

[steven_wu]

Offline comment to Florian is that:

The other possibility is that I don't see why SplitModule needs to own the module. Can backend just take Module& instead of grabbing the unique_ptr? That is probably the cleanest change.

[fhahn]

Updated splitCodeGen & co to pass the module by reference instead of taking ownership. splitCodeGen clones the original module before carving it up, so it does not need ownership and it simplifies the code. We can change it back to taking ownership, once it is actually required.

[steven_wu]

LGTM. @pcc does this look fine to you?

[tejohnson]

In D95222#2525197, @fhahn wrote:

Updated splitCodeGen & co to pass the module by reference instead of taking ownership. splitCodeGen clones the original module before carving it up, so it does not need ownership and it simplifies the code. We can change it back to taking ownership, once it is actually required.

Will that cause issues if ownership is required down the line? If so, perhaps just remove the fixme about avoiding cloning the module for the last partition.

llvm/include/llvm/CodeGen/ParallelCG.h

34–35

Update comment

[fhahn]

In D95222#2526788, @tejohnson wrote:

In D95222#2525197, @fhahn wrote:

Updated splitCodeGen & co to pass the module by reference instead of taking ownership. splitCodeGen clones the original module before carving it up, so it does not need ownership and it simplifies the code. We can change it back to taking ownership, once it is actually required.

Will that cause issues if ownership is required down the line? If so, perhaps just remove the fixme about avoiding cloning the module for the last partition.

Yes, I think some additional changes will be needed to address the fixme. I think the fixme is still valuable, so rather than removing it I added a note about callers not expecting the module to change, so whoever addresses the fixme is aware that further changes are need. (The simplest fix would be to just clone the original module once).

What do you think?

[tejohnson]

llvm/include/llvm/CodeGen/ParallelCG.h

35–39

I don't think this comment is accurate. Looking at splitCodeGen, in the case of parallelism=1 it will call codegen() on the input Module, which will apply passes to emit the Module, which in turn will modify it. In the case of parallelism>1, it will call SplitModule, which will externalize some symbols before doing the cloning (because the default of PreserveLocals is false).

[fhahn]

Undo comment change.

[fhahn]

llvm/include/llvm/CodeGen/ParallelCG.h

35–39

Yes you are right, that wasn't right. I meant to say it preserves the semantics of the original module. I removed the comment.

[tejohnson]

lgtm

llvm/include/llvm/CodeGen/ParallelCG.h

 
 /// Split M into OSs.size() partitions, and generate code for each. Takes a
 /// factory function for the TargetMachine TMFactory. Writes OSs.size() output
 /// files to the output streams in OSs. The resulting output files if linked
 /// together are intended to be equivalent to the single output file that would
 /// have been code generated from M.
 ///
 /// 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,
+    CodeGenFileType FileType = CGFT_ObjectFile, bool PreserveLocals = false);
 
 } // namespace llvm
 
 #endif

llvm/include/llvm/LTO/LTOBackend.h

 bool opt(const Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
          bool IsThinLTO, ModuleSummaryIndex *ExportSummary,
          const ModuleSummaryIndex *ImportSummary,
          const std::vector<uint8_t> &CmdArgs);
 
 /// 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,
                   Module &M, const ModuleSummaryIndex &CombinedIndex,
                   const FunctionImporter::ImportMapTy &ImportList,
                   const GVSummaryMapTy &DefinedGlobals,
                   MapVector<StringRef, BitcodeModule> &ModuleMap,
                   const std::vector<uint8_t> &CmdArgs = std::vector<uint8_t>());

llvm/include/llvm/Transforms/Utils/SplitModule.h

 /// FIXME: This function does not deal with the somewhat subtle symbol
 /// visibility issues around module splitting, including (but not limited to):
 ///
 /// - Internal symbols should not collide with symbols defined outside the
 ///   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);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_SPLITMODULE_H

llvm/lib/CodeGen/ParallelCG.cpp

static void codegen(Module *M, llvm::raw_pwrite_stream &OS,
   assert(TM && "Failed to create target machine!");
 
   legacy::PassManager CodeGenPasses;
   if (TM->addPassesToEmitFile(CodeGenPasses, OS, nullptr, FileType))
     report_fatal_error("Failed to setup codegen");
   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) {
   assert(BCOSs.empty() || BCOSs.size() == OSs.size());
 
   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
   // on destruction.
   {
     ThreadPool CodegenThreadPool(hardware_concurrency(OSs.size()));
     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
           // (while still on the main thread, in order to avoid data races) and
           // spinning up new threads which deserialize the partitions into
           // separate contexts.
           // FIXME: Provide a more direct way to do this in LLVM.
           SmallString<0> BC;
SplitModule(
                 codegen(MPartInCtx.get(), *ThreadOS, TMFactory, FileType);
               },
               // Pass BC using std::move to ensure that it get moved rather than
               // copied into the thread's context.
               std::move(BC));
         },
         PreserveLocals);
   }
-
-  return {};
 }

llvm/lib/LTO/LTO.cpp

if (!Conf.CodeGenOnly) {
     RegularLTO.CombinedModule->addModuleFlag(Module::Error, "LTOPostLink", 1);
 
     if (Conf.PostInternalizeModuleHook &&
         !Conf.PostInternalizeModuleHook(0, *RegularLTO.CombinedModule))
       return Error::success();
   }
 
   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;
   }
 
   return finalizeOptimizationRemarks(std::move(*DiagFileOrErr));
 }
 
 static const char *libcallRoutineNames[] = {
 #define HANDLE_LIBCALL(code, name) name,

llvm/lib/LTO/LTOBackend.cpp

static void codegen(const Config &Conf, TargetMachine *TM,
   CodeGenPasses.run(Mod);
 
   if (DwoOut)
     DwoOut->keep();
 }
 
 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));
   unsigned ThreadCount = 0;
   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
         // (while still on the main thread, in order to avoid data races) and
         // spinning up new threads which deserialize the partitions into
         // separate contexts.
         // FIXME: Provide a more direct way to do this in LLVM.
         SmallString<0> BC;
Error lto::finalizeOptimizationRemarks(
   if (!DiagOutputFile)
     return Error::success();
   DiagOutputFile->keep();
   DiagOutputFile->os().flush();
   return Error::success();
 }
 
 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();
 }
 
 static void dropDeadSymbols(Module &Mod, const GVSummaryMapTy &DefinedGlobals,
                             const ModuleSummaryIndex &Index) {
   std::vector<GlobalValue*> DeadGVs;
   for (auto &GV : Mod.global_values())

llvm/lib/LTO/LTOCodeGenerator.cpp

bool LTOCodeGenerator::compileOptimized(ArrayRef<raw_pwrite_stream *> Out) {
   // the ObjCARCContractPass must be run, so do it unconditionally here.
   preCodeGenPasses.add(createObjCARCContractPass());
   preCodeGenPasses.run(*MergedModule);
 
   // Re-externalize globals that may have been internalized to increase scope
   // 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,
 
   // If statistics were requested, save them to the specified file or
   // print them out after codegen.
   if (StatsFile)
     PrintStatisticsJSON(StatsFile->os());
   else if (AreStatisticsEnabled())
     PrintStatistics();
 

llvm/lib/Transforms/Utils/SplitModule.cpp

for (auto *U : V->users()) {
     }
   }
 }
 
 // Find partitions for module in the way that no locals need to be
 // 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;
 
   auto recordGVSet = [&GVtoClusterMap, &ComdatMembers](GlobalValue &GV) {
     if (GV.isDeclaration())
       return;
 
     if (!GV.hasName())
if (const Function *F = dyn_cast<Function>(&GV)) {
         addAllGlobalValueUsers(GVtoClusterMap, F, BA);
       }
     }
 
     if (GV.hasLocalLinkage())
       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.
   auto CompareClusters = [](const std::pair<unsigned, unsigned> &a,
                             const std::pair<unsigned, unsigned> &b) {
     if (a.second || b.second)
       return a.second > b.second;
     else
static bool isInPartition(const GlobalValue *GV, unsigned I, unsigned N) {
   MD5 H;
   MD5::MD5Result R;
   H.update(Name);
   H.final(R);
   return (R[0] | (R[1] << 8)) % N == I;
 }
 
 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
             return isInPartition(GV, I, N);
         }));
     if (I != 0)
       MPart->setModuleInlineAsm("");
     ModuleCallback(std::move(MPart));
   }
 }

llvm/tools/llvm-split/llvm-split.cpp

int main(int argc, char **argv) {
   std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
 
   if (!M) {
     Err.print(argv[0], errs());
     return 1;
   }
 
   unsigned I = 0;
-  SplitModule(std::move(M), NumOutputs, [&](std::unique_ptr<Module> MPart) {
+  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));
+        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);
         }
 
         WriteBitcodeToFile(*MPart, Out->os());
 
         // Declare success.
         Out->keep();
-  }, PreserveLocals);
+      },
+      PreserveLocals);
 
   return 0;
 }