[LTO] SplitCodeGen new API

[davide]

Summary

.. so that we don't have to create TM in two places in lld. See http://reviews.llvm.org/D18999 for reference.
Hopefully should work, assuming TheTarget is thread-safe. Thanks to Michael Spencer for help with this patch.

[mehdi_amini]

Looks OK, see some inline comments.

include/llvm/CodeGen/ParallelCG.h

49 ↗

(On Diff #53804)

Is there a better name for F?
Also can you add a doxygen for this API (you can refer to the other saying this is a variant that takes a factory function for the TargetMachine)

49 ↗

(On Diff #53804)

Also: const &

lib/CodeGen/ParallelCG.cpp

29 ↗

(On Diff #53804)

Use another name than F, and take a const &.

51 ↗

(On Diff #53804)

Do you really mean = here? I'd use &.

[Bigcheese]

lib/CodeGen/ParallelCG.cpp

50–54 ↗

(On Diff #53804)

clang-format.

[davide]

Comments addressed.

[davide]

Correct upload.

[mehdi_amini]

LGTM, with a few more fixes (see inline)

include/llvm/CodeGen/ParallelCG.h

49 ↗

(On Diff #53815)

Add in the documentation the fact that TMFactory needs to be thread safe on the client side.

lib/CodeGen/ParallelCG.cpp

30 ↗

(On Diff #53815)

still missing const & here.

99 ↗

(On Diff #53815)

&TMFactory (avoid copy here).

[davide]

Thank you Mehdi, r266390.

llvm/trunk/include/llvm/CodeGen/ParallelCG.h

splitCodeGen(std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs,
              ArrayRef<llvm::raw_pwrite_stream *> BCOSs, StringRef CPU,
              StringRef Features, const TargetOptions &Options,
              Reloc::Model RM = Reloc::Default,
              CodeModel::Model CM = CodeModel::Default,
              CodeGenOpt::Level OL = CodeGenOpt::Default,
              TargetMachine::CodeGenFileType FT = TargetMachine::CGFT_ObjectFile,
              bool PreserveLocals = false);
 
+/// Split M into OSs.size() partitions, and generate code for each.
+/// It is a variant that takes a factory function for the TargetMachine
+/// TMFactory. TMFactory needs to be thread safe on the client side.
+/// See the other splitCodeGen() for a more detailed description.
+///
+/// \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,
+             TargetMachine::CodeGenFileType FT = TargetMachine::CGFT_ObjectFile,
+             bool PreserveLocals = false);
+
 } // namespace llvm
 
 #endif

llvm/trunk/lib/CodeGen/ParallelCG.cpp

 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 
 using namespace llvm;
 
-static void codegen(Module *M, llvm::raw_pwrite_stream &OS,
-                    const Target *TheTarget, StringRef CPU, StringRef Features,
-                    const TargetOptions &Options, Reloc::Model RM,
-                    CodeModel::Model CM, CodeGenOpt::Level OL,
-                    TargetMachine::CodeGenFileType FileType) {
-  std::unique_ptr<TargetMachine> TM(TheTarget->createTargetMachine(
-      M->getTargetTriple(), CPU, Features, Options, RM, CM, OL));
-
+static void
+codegen(Module *M, llvm::raw_pwrite_stream &OS,
+        const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
+  std::unique_ptr<TargetMachine> TM = TMFactory();
   legacy::PassManager CodeGenPasses;
   if (TM->addPassesToEmitFile(CodeGenPasses, OS, 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,
+std::unique_ptr<Module>
+llvm::splitCodeGen(std::unique_ptr<Module> M, ArrayRef<raw_pwrite_stream *> OSs,
                    ArrayRef<llvm::raw_pwrite_stream *> BCOSs, StringRef CPU,
-    StringRef Features, const TargetOptions &Options, Reloc::Model RM,
-    CodeModel::Model CM, CodeGenOpt::Level OL,
-    TargetMachine::CodeGenFileType FileType, bool PreserveLocals) {
+                   StringRef Features, const TargetOptions &Options,
+                   Reloc::Model RM, CodeModel::Model CM, CodeGenOpt::Level OL,
+                   TargetMachine::CodeGenFileType FileType,
+                   bool PreserveLocals) {
   StringRef TripleStr = M->getTargetTriple();
   std::string ErrMsg;
+
   const Target *TheTarget = TargetRegistry::lookupTarget(TripleStr, ErrMsg);
   if (!TheTarget)
     report_fatal_error(Twine("Target not found: ") + ErrMsg);
+  return splitCodeGen(std::move(M), OSs, BCOSs, [&]() {
+    return std::unique_ptr<TargetMachine>(TheTarget->createTargetMachine(
+        TripleStr, CPU, Features, Options, RM, CM, OL));
+  }, FileType, PreserveLocals);
+}
 
+std::unique_ptr<Module> llvm::splitCodeGen(
+    std::unique_ptr<Module> M, ArrayRef<llvm::raw_pwrite_stream *> OSs,
+    ArrayRef<llvm::raw_pwrite_stream *> BCOSs,
+    const std::function<std::unique_ptr<TargetMachine>()> &TMFactory,
+    TargetMachine::CodeGenFileType FileType, bool PreserveLocals) {
   assert(BCOSs.empty() || BCOSs.size() == OSs.size());
 
   if (OSs.size() == 1) {
     if (!BCOSs.empty())
       WriteBitcodeToFile(M.get(), *BCOSs[0]);
-    codegen(M.get(), *OSs[0], TheTarget, CPU, Features, Options, RM, CM, OL,
-            FileType);
+    codegen(M.get(), *OSs[0], TMFactory, FileType);
     return M;
   }
 
   // Create ThreadPool in nested scope so that threads will be joined
   // on destruction.
   {
     ThreadPool CodegenThreadPool(OSs.size());
     int ThreadCount = 0;
SplitModule(
           if (!BCOSs.empty()) {
             BCOSs[ThreadCount]->write(BC.begin(), BC.size());
             BCOSs[ThreadCount]->flush();
           }
 
           llvm::raw_pwrite_stream *ThreadOS = OSs[ThreadCount++];
           // Enqueue the task
           CodegenThreadPool.async(
-              [TheTarget, CPU, Features, Options, RM, CM, OL, FileType,
-               ThreadOS](const SmallVector<char, 0> &BC) {
+              [&TMFactory, FileType, ThreadOS](const SmallVector<char, 0> &BC) {
                 LLVMContext Ctx;
                 ErrorOr<std::unique_ptr<Module>> MOrErr = parseBitcodeFile(
                     MemoryBufferRef(StringRef(BC.data(), BC.size()),
                                     "<split-module>"),
                     Ctx);
                 if (!MOrErr)
                   report_fatal_error("Failed to read bitcode");
                 std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get());
 
-                codegen(MPartInCtx.get(), *ThreadOS, TheTarget, CPU, Features,
-                        Options, RM, CM, OL, FileType);
+                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 {};
 }