[LTO] Refactor the ThinLTO distributed backend

[mehdi_amini]

Summary

Delegate the way the information is communicated to the build
system to the client.

[tejohnson]

It isn't clear to me why we would want this functionality to live in the clients, rather than providing a set of configurable functionality that different build systems might need inside LTO itself. It means possibly implementing the same functionality required by a single build system in potentially all clients, since the build system is not dependent on a single client linker.

include/llvm/LTO/LTO.h

262

Parameter is called AddModule elsewhere.

tools/gold/gold-plugin.cpp

754

This may be null

[mehdi_amini]

Address comments

include/llvm/LTO/LTO.h

typedef std::function<std::unique_ptr<ThinBackendProc>(
     Config &C, ModuleSummaryIndex &CombinedIndex,
     StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
     AddOutputFn AddOutput)>
     ThinBackend;
 
 /// This ThinBackend runs the individual backend jobs in-process.
 ThinBackend createInProcessThinBackend(unsigned ParallelismLevel);
 
+/// Defines a callback called for each module for the distributed ThinLTO
+/// backend. The client implements this callback to transmit the information to
+/// the distributed build system.
+typedef std::function<void(
+    unsigned Task, StringRef ModuleIdentifier,
+    const ModuleSummaryIndex &CombinedIndex,
+    const FunctionImporter::ImportMapTy &ImportList,
+    const std::map<std::string, GVSummaryMapTy> &ModuleToSummariesForIndex)>
+    WriteIndexFn;
+
 /// This ThinBackend writes individual module indexes to files, instead of
 /// running the individual backend jobs. This backend is for distributed builds
 /// where separate processes will invoke the real backends.
-///
-/// To find the path to write the index to, the backend checks if the path has a
-/// prefix of OldPrefix; if so, it replaces that prefix with NewPrefix. It then
-/// appends ".thinlto.bc" and writes the index to that path. If
-/// ShouldEmitImportsFiles is true it also writes a list of imported files to a
-/// similar path with ".imports" appended instead.
-ThinBackend createWriteIndexesThinBackend(std::string OldPrefix,
-                                          std::string NewPrefix,
-                                          bool ShouldEmitImportsFiles,
-                                          std::string LinkedObjectsFile);
+ThinBackend createWriteIndexesThinBackend(WriteIndexFn PerModuleCallback);
 
 /// This class implements a resolution-based interface to LLVM's LTO
 /// functionality. It supports regular LTO, parallel LTO code generation and
 /// ThinLTO. You can use it from a linker in the following way:
 /// - Set hooks and code generation options (see lto::Config struct defined in
 ///   Config.h), and use the lto::Config object to create an lto::LTO object.
 /// - Create lto::InputFile objects using lto::InputFile::create(), then use
 ///   the symbols() function to enumerate its symbols and compute a resolution

lib/LTO/LTO.cpp

return [=](Config &Conf, ModuleSummaryIndex &CombinedIndex,
              AddOutputFn AddOutput) {
     return llvm::make_unique<InProcessThinBackend>(
         Conf, CombinedIndex, ParallelismLevel, ModuleToDefinedGVSummaries,
         AddOutput);
   };
 }
 
 class WriteIndexesThinBackend : public ThinBackendProc {
-  std::string OldPrefix, NewPrefix;
-  bool ShouldEmitImportsFiles;
-
-  std::string LinkedObjectsFileName;
-  std::unique_ptr<llvm::raw_fd_ostream> LinkedObjectsFile;
+  WriteIndexFn PerModuleCallback;
 
 public:
-  WriteIndexesThinBackend(Config &Conf, ModuleSummaryIndex &CombinedIndex,
-                          StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-                          std::string OldPrefix, std::string NewPrefix,
-                          bool ShouldEmitImportsFiles,
-                          std::string LinkedObjectsFileName)
+  WriteIndexesThinBackend(
+      Config &Conf, ModuleSummaryIndex &CombinedIndex,
+      StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries, // FIXME const
+      WriteIndexFn PerModuleCallback)
       : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries),
-        OldPrefix(OldPrefix), NewPrefix(NewPrefix),
-        ShouldEmitImportsFiles(ShouldEmitImportsFiles),
-        LinkedObjectsFileName(LinkedObjectsFileName) {}
-
-  /// Given the original \p Path to an output file, replace any path
-  /// prefix matching \p OldPrefix with \p NewPrefix. Also, create the
-  /// resulting directory if it does not yet exist.
-  std::string getThinLTOOutputFile(const std::string &Path,
-                                   const std::string &OldPrefix,
-                                   const std::string &NewPrefix) {
-    if (OldPrefix.empty() && NewPrefix.empty())
-      return Path;
-    SmallString<128> NewPath(Path);
-    llvm::sys::path::replace_path_prefix(NewPath, OldPrefix, NewPrefix);
-    StringRef ParentPath = llvm::sys::path::parent_path(NewPath.str());
-    if (!ParentPath.empty()) {
-      // Make sure the new directory exists, creating it if necessary.
-      if (std::error_code EC = llvm::sys::fs::create_directories(ParentPath))
-        llvm::errs() << "warning: could not create directory '" << ParentPath
-                     << "': " << EC.message() << '\n';
-    }
-    return NewPath.str();
-  }
+        PerModuleCallback(std::move(PerModuleCallback)) {}
 
   Error start(unsigned Task, MemoryBufferRef MBRef,
               const FunctionImporter::ImportMapTy &ImportList,
               MapVector<StringRef, MemoryBufferRef> &ModuleMap) override {
     StringRef ModulePath = MBRef.getBufferIdentifier();
-    std::string NewModulePath =
-        getThinLTOOutputFile(ModulePath, OldPrefix, NewPrefix);
-
-    std::error_code EC;
-    if (!LinkedObjectsFileName.empty()) {
-      if (!LinkedObjectsFile) {
-        LinkedObjectsFile = llvm::make_unique<raw_fd_ostream>(
-            LinkedObjectsFileName, EC, sys::fs::OpenFlags::F_None);
-        if (EC)
-          return errorCodeToError(EC);
-      }
-      *LinkedObjectsFile << NewModulePath << '\n';
-    }
-
     std::map<std::string, GVSummaryMapTy> ModuleToSummariesForIndex;
     gatherImportedSummariesForModule(ModulePath, ModuleToDefinedGVSummaries,
                                      ImportList, ModuleToSummariesForIndex);
 
-    raw_fd_ostream OS(NewModulePath + ".thinlto.bc", EC,
-                      sys::fs::OpenFlags::F_None);
-    if (EC)
-      return errorCodeToError(EC);
-    WriteIndexToFile(CombinedIndex, OS, &ModuleToSummariesForIndex);
-
-    if (ShouldEmitImportsFiles)
-      return errorCodeToError(
-          EmitImportsFiles(ModulePath, NewModulePath + ".imports", ImportList));
+    PerModuleCallback(Task, ModulePath, CombinedIndex, ImportList,
+                      ModuleToSummariesForIndex);
     return Error();
   }
 
   Error wait() override { return Error(); }
 };
 
-ThinBackend lto::createWriteIndexesThinBackend(std::string OldPrefix,
-                                               std::string NewPrefix,
-                                               bool ShouldEmitImportsFiles,
-                                               std::string LinkedObjectsFile) {
-  return [=](Config &Conf, ModuleSummaryIndex &CombinedIndex,
-             StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-             AddOutputFn AddOutput) {
+ThinBackend lto::createWriteIndexesThinBackend(WriteIndexFn PerModuleCallback) {
+  return std::bind(
+      [](WriteIndexFn PerModuleCallback, Config &Conf,
+         ModuleSummaryIndex &CombinedIndex,
+         AddOutputFn AddOutput) -> std::unique_ptr<ThinBackendProc> {
         return llvm::make_unique<WriteIndexesThinBackend>(
-        Conf, CombinedIndex, ModuleToDefinedGVSummaries, OldPrefix, NewPrefix,
-        ShouldEmitImportsFiles, LinkedObjectsFile);
-  };
+            Conf, CombinedIndex, ModuleToDefinedGVSummaries,
+            std::move(PerModuleCallback));
+      },
+      std::move(PerModuleCallback), std::placeholders::_1,
+      std::placeholders::_2, std::placeholders::_3, std::placeholders::_4);
 }
 
 Error LTO::runThinLTO(AddOutputFn AddOutput) {
   if (ThinLTO.ModuleMap.empty())
     return Error();
 
   if (Conf.CombinedIndexHook && !Conf.CombinedIndexHook(ThinLTO.CombinedIndex))
     return Error();

tools/gold/gold-plugin.cpp

std::unique_ptr<raw_pwrite_stream> getStream() override {
     std::error_code EC = sys::fs::openFileForWrite(Path, FD, sys::fs::F_None);
     if (EC)
       message(LDPL_FATAL, "Could not open file: %s", EC.message().c_str());
     return llvm::make_unique<llvm::raw_fd_ostream>(FD, true);
   }
 };
 }
 
+/// Given the original \p Path to an output file, replace any path
+/// prefix matching \p OldPrefix with \p NewPrefix. Also, create the
+/// resulting directory if it does not yet exist.
+static std::string getThinLTOOutputFile(const std::string &Path,
+                                        const std::string &OldPrefix,
+                                        const std::string &NewPrefix) {
+  if (OldPrefix.empty() && NewPrefix.empty())
+    return Path;
+  SmallString<128> NewPath(Path);
+  llvm::sys::path::replace_path_prefix(NewPath, OldPrefix, NewPrefix);
+  StringRef ParentPath = llvm::sys::path::parent_path(NewPath.str());
+  if (!ParentPath.empty()) {
+    // Make sure the new directory exists, creating it if necessary.
+    if (std::error_code EC = llvm::sys::fs::create_directories(ParentPath))
+      llvm::errs() << "warning: could not create directory '" << ParentPath
+                   << "': " << EC.message() << '\n';
+  }
+  return NewPath.str();
+}
+
 static std::unique_ptr<LTO> createLTO() {
   Config Conf;
   ThinBackend Backend;
   unsigned ParallelCodeGenParallelismLevel = 1;
 
   Conf.CPU = options::mcpu;
   Conf.Options = InitTargetOptionsFromCodeGenFlags();
 
if (options::Parallelism) {
     if (options::thinlto)
       Backend = createInProcessThinBackend(options::Parallelism);
     else
       ParallelCodeGenParallelismLevel = options::Parallelism;
   }
   if (options::thinlto_index_only) {
     std::string OldPrefix, NewPrefix;
     getThinLTOOldAndNewPrefix(OldPrefix, NewPrefix);
-    Backend = createWriteIndexesThinBackend(
-        OldPrefix, NewPrefix, options::thinlto_emit_imports_files,
-        options::thinlto_linked_objects_file);
+
+    std::unique_ptr<llvm::raw_fd_ostream> LinkedObjectsFile;
+    if (!options::thinlto_linked_objects_file.empty()) {
+      std::error_code EC;
+      LinkedObjectsFile = llvm::make_unique<raw_fd_ostream>(
+          options::thinlto_linked_objects_file, EC, sys::fs::OpenFlags::F_None);
+      check(EC, "Can't open ", options::thinlto_linked_objects_file);
+    }
+
+    Backend = createWriteIndexesThinBackend([&](
+        unsigned Task, StringRef ModulePath,
+        const ModuleSummaryIndex &CombinedIndex,
+        const FunctionImporter::ImportMapTy &ImportList,
+        const std::map<std::string, GVSummaryMapTy>
+            &ModuleToSummariesForIndex) {
+      std::string NewModulePath =
+          getThinLTOOutputFile(ModulePath, OldPrefix, NewPrefix);
+
+      // Add to the list of emitted objects
+      if (LinkedObjectsFile)
+        *LinkedObjectsFile << NewModulePath << '\n';
+      raw_fd_ostream OS(NewModulePath + ".thinlto.bc", EC,
+                        sys::fs::OpenFlags::F_None);
+      if (EC)
+        return errorCodeToError(EC);
+      WriteIndexToFile(CombinedIndex, OS, &ModuleToSummariesForIndex);
+
+      if (options::thinlto_emit_imports_files)
+        check(EmitImportsFiles(ModulePath, NewModulePath + ".imports",
+                               ImportList),
+              "Can't emit imports to " + options::thinlto_emit_imports_files);
+    });
   }
 
   Conf.OverrideTriple = options::triple;
   Conf.DefaultTriple = sys::getDefaultTargetTriple();
 
   Conf.DiagHandler = diagnosticHandler;
 
   switch (options::TheOutputType) {

tools/llvm-lto2/llvm-lto2.cpp

 // link-time optimization according to the provided symbol resolutions using the
 // resolution-based LTO interface, and outputs one or more object files.
 //
 // This program is intended to eventually replace llvm-lto which uses the legacy
 // LTO interface.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Bitcode/ReaderWriter.h"
 #include "llvm/LTO/LTO.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace llvm;
 using namespace lto;
 using namespace object;
 
int main(int argc, char **argv) {
   };
 
   if (SaveTemps)
     check(Conf.addSaveTemps(OutputFilename + "."),
           "Config::addSaveTemps failed");
 
   ThinBackend Backend;
   if (ThinLTODistributedIndexes)
-    Backend = createWriteIndexesThinBackend("", "", true, "");
-  else
+    Backend = createWriteIndexesThinBackend(
+        [&](unsigned Task, StringRef ModuleIdentifier,
+            const ModuleSummaryIndex &CombinedIndex,
+            const FunctionImporter::ImportMapTy &ImportList,
+            const std::map<std::string, GVSummaryMapTy>
+                &ModuleToSummariesForIndex) {
+          std::error_code EC;
+          auto IndexPath = (ModuleIdentifier + ".thinlto.bc").str();
+          raw_fd_ostream OS(IndexPath, EC, sys::fs::OpenFlags::F_None);
+          check(EC, "Creating " + IndexPath);
+          WriteIndexToFile(CombinedIndex, OS, &ModuleToSummariesForIndex);
+
+          auto ImportsPath = (ModuleIdentifier + ".imports").str();
+          EmitImportsFiles(ModuleIdentifier, ImportsPath, ImportList);
+          check(EC, "Emitting Imports list to " + ImportsPath);
+        });
+  else {
     Backend = createInProcessThinBackend(Threads);
+  }
   LTO Lto(std::move(Conf), std::move(Backend));
 
   bool HasErrors = false;
   for (std::string F : InputFilenames) {
     std::unique_ptr<MemoryBuffer> MB = check(MemoryBuffer::getFile(F), F);
     std::unique_ptr<InputFile> Input =
         check(InputFile::create(MB->getMemBufferRef()), F);