[lld/mac] Add support for distributed ThinLTO

Adds support for the following flags:

* --thinlto-index-only, --thinlto-index-only=
* --thinlto-emit-imports-files
* --thinlto-emit-index-files
* --thinlto-object-suffix-replace=
* --thinlto-prefix-replace=

See https://blog.llvm.org/2016/06/thinlto-scalable-and-incremental-lto.html
for some words on --thinlto-index-only.

I don't really need the other flags, but they were in the vicinity
and _someone_ might need them, so I figured I'd add them too.

`-object_path_lto` now sets `c.AlwaysEmitRegularLTOObj` as in the other ports,
which means it can now only point to a filename for non-thin LTO.
I think that was the intent of D129705 anyways, so update
test/MachO/lto-object-path.ll to use a non-thin bitcode file for that test.

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

Author: nico

lld/ELF/InputFiles.cpp

@@ -1773,9 +1773,7 @@ bool InputFile::shouldExtractForCommon(StringRef name) {
 }
 
 std::string elf::replaceThinLTOSuffix(StringRef path) {
-  StringRef suffix = config->thinLTOObjectSuffixReplace.first;
-  StringRef repl = config->thinLTOObjectSuffixReplace.second;
-
+  auto [suffix, repl] = config->thinLTOObjectSuffixReplace;
   if (path.consume_back(suffix))
     return (path + repl).str();
   return std::string(path);

lld/MachO/Config.h

@@ -133,6 +133,9 @@ struct Configuration {
   bool emitEncryptionInfo = false;
   bool emitInitOffsets = false;
   bool emitChainedFixups = false;
+  bool thinLTOEmitImportsFiles;
+  bool thinLTOEmitIndexFiles;
+  bool thinLTOIndexOnly;
   bool timeTraceEnabled = false;
   bool dataConst = false;
   bool dedupStrings = true;
@@ -164,6 +167,9 @@ struct Configuration {
   uint32_t ltoo = 2;
   llvm::CachePruningPolicy thinLTOCachePolicy;
   llvm::StringRef thinLTOCacheDir;
+  llvm::StringRef thinLTOIndexOnlyArg;
+  std::pair<llvm::StringRef, llvm::StringRef> thinLTOObjectSuffixReplace;
+  std::pair<llvm::StringRef, llvm::StringRef> thinLTOPrefixReplace;
   bool deadStripDylibs = false;
   bool demangle = false;
   bool deadStrip = false;

lld/MachO/Driver.cpp

@@ -858,6 +858,20 @@ static const char *getReproduceOption(InputArgList &args) {
   return getenv("LLD_REPRODUCE");
 }
 
+// Parse options of the form "old;new".
+static std::pair<StringRef, StringRef> getOldNewOptions(opt::InputArgList &args,
+                                                        unsigned id) {
+  auto *arg = args.getLastArg(id);
+  if (!arg)
+    return {"", ""};
+
+  StringRef s = arg->getValue();
+  std::pair<StringRef, StringRef> ret = s.split(';');
+  if (ret.second.empty())
+    error(arg->getSpelling() + " expects 'old;new' format, but got " + s);
+  return ret;
+}
+
 static void parseClangOption(StringRef opt, const Twine &msg) {
   std::string err;
   raw_string_ostream os(err);
@@ -1549,6 +1563,25 @@ bool macho::link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
     error("--lto-O: invalid optimization level: " + Twine(config->ltoo));
   config->thinLTOCacheDir = args.getLastArgValue(OPT_cache_path_lto);
   config->thinLTOCachePolicy = getLTOCachePolicy(args);
+  config->thinLTOEmitImportsFiles = args.hasArg(OPT_thinlto_emit_imports_files);
+  config->thinLTOEmitIndexFiles = args.hasArg(OPT_thinlto_emit_index_files) ||
+                                  args.hasArg(OPT_thinlto_index_only) ||
+                                  args.hasArg(OPT_thinlto_index_only_eq);
+  config->thinLTOIndexOnly = args.hasArg(OPT_thinlto_index_only) ||
+                             args.hasArg(OPT_thinlto_index_only_eq);
+  config->thinLTOIndexOnlyArg = args.getLastArgValue(OPT_thinlto_index_only_eq);
+  config->thinLTOObjectSuffixReplace =
+      getOldNewOptions(args, OPT_thinlto_object_suffix_replace_eq);
+  config->thinLTOPrefixReplace =
+      getOldNewOptions(args, OPT_thinlto_prefix_replace_eq);
+  if (config->thinLTOEmitIndexFiles && !config->thinLTOIndexOnly) {
+    if (args.hasArg(OPT_thinlto_object_suffix_replace_eq))
+      error("--thinlto-object-suffix-replace is not supported with "
+            "--thinlto-emit-index-files");
+    else if (args.hasArg(OPT_thinlto_prefix_replace_eq))
+      error("--thinlto-prefix-replace is not supported with "
+            "--thinlto-emit-index-files");
+  }
   config->runtimePaths = args::getStrings(args, OPT_rpath);
   config->allLoad = args.hasFlag(OPT_all_load, OPT_noall_load, false);
   config->archMultiple = args.hasArg(OPT_arch_multiple);
@@ -1834,11 +1867,20 @@ bool macho::link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
     // explicitly exported. Do this before running LTO so that LTO can better
     // optimize.
     handleExplicitExports();
+
+    bool didCompileBitcodeFiles = compileBitcodeFiles();
+
+    // If --thinlto-index-only is given, we should create only "index
+    // files" and not object files. Index file creation is already done
+    // in compileBitcodeFiles, so we are done if that's the case.
+    if (config->thinLTOIndexOnly)
+      return errorCount() == 0;
+
     // LTO may emit a non-hidden (extern) object file symbol even if the
     // corresponding bitcode symbol is hidden. In particular, this happens for
     // cross-module references to hidden symbols under ThinLTO. Thus, if we
     // compiled any bitcode files, we must redo the symbol hiding.
-    if (compileBitcodeFiles())
+    if (didCompileBitcodeFiles)
       handleExplicitExports();
     replaceCommonSymbols();
 

lld/MachO/InputFiles.cpp

@@ -2206,6 +2206,9 @@ BitcodeFile::BitcodeFile(MemoryBufferRef mb, StringRef archiveName,
     : InputFile(BitcodeKind, mb, lazy), forceHidden(forceHidden) {
   this->archiveName = std::string(archiveName);
   std::string path = mb.getBufferIdentifier().str();
+  if (config->thinLTOIndexOnly)
+    path = replaceThinLTOSuffix(mb.getBufferIdentifier());
+
   // ThinLTO assumes that all MemoryBufferRefs given to it have a unique
   // name. If two members with the same name are provided, this causes a
   // collision and ThinLTO can't proceed.
@@ -2246,6 +2249,13 @@ void BitcodeFile::parseLazy() {
   }
 }
 
+std::string macho::replaceThinLTOSuffix(StringRef path) {
+  auto [suffix, repl] = config->thinLTOObjectSuffixReplace;
+  if (path.consume_back(suffix))
+    return (path + repl).str();
+  return std::string(path);
+}
+
 void macho::extract(InputFile &file, StringRef reason) {
   if (!file.lazy)
     return;

lld/MachO/InputFiles.h

@@ -357,6 +357,7 @@ std::vector<const CommandType *> findCommands(const void *anyHdr,
   return detail::findCommands<CommandType>(anyHdr, 0, types...);
 }
 
+std::string replaceThinLTOSuffix(StringRef path);
 } // namespace macho
 
 std::string toString(const macho::InputFile *file);

lld/MachO/LTO.cpp

@@ -17,6 +17,7 @@
 #include "lld/Common/CommonLinkerContext.h"
 #include "lld/Common/Strings.h"
 #include "lld/Common/TargetOptionsCommandFlags.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/LTO/Config.h"
 #include "llvm/LTO/LTO.h"
 #include "llvm/Support/Caching.h"
@@ -31,6 +32,25 @@ using namespace llvm;
 using namespace llvm::MachO;
 using namespace llvm::sys;
 
+// Creates an empty file to store a list of object files for final
+// linking of distributed ThinLTO.
+static std::unique_ptr<raw_fd_ostream> openFile(StringRef file) {
+  std::error_code ec;
+  auto ret =
+      std::make_unique<raw_fd_ostream>(file, ec, sys::fs::OpenFlags::OF_None);
+  if (ec) {
+    error("cannot open " + file + ": " + ec.message());
+    return nullptr;
+  }
+  return ret;
+}
+
+static std::string getThinLTOOutputFile(StringRef modulePath) {
+  return lto::getThinLTOOutputFile(
+      std::string(modulePath), std::string(config->thinLTOPrefixReplace.first),
+      std::string(config->thinLTOPrefixReplace.second));
+}
+
 static lto::Config createConfig() {
   lto::Config c;
   c.Options = initTargetOptionsFromCodeGenFlags();
@@ -44,6 +64,9 @@ static lto::Config createConfig() {
   c.PreCodeGenPassesHook = [](legacy::PassManager &pm) {
     pm.add(createObjCARCContractPass());
   };
+
+  c.AlwaysEmitRegularLTOObj = !config->ltoObjPath.empty();
+
   c.TimeTraceEnabled = config->timeTraceEnabled;
   c.TimeTraceGranularity = config->timeTraceGranularity;
   c.OptLevel = config->ltoo;
@@ -67,13 +90,35 @@ static void saveOrHardlinkBuffer(StringRef buffer, const Twine &path,
 }
 
 BitcodeCompiler::BitcodeCompiler() {
-  lto::ThinBackend backend = lto::createInProcessThinBackend(
-      heavyweight_hardware_concurrency(config->thinLTOJobs));
+  // Initialize indexFile.
+  if (!config->thinLTOIndexOnlyArg.empty())
+    indexFile = openFile(config->thinLTOIndexOnlyArg);
+
+  // Initialize ltoObj.
+  lto::ThinBackend backend;
+  auto onIndexWrite = [&](StringRef S) { thinIndices.erase(S); };
+  if (config->thinLTOIndexOnly) {
+    backend = lto::createWriteIndexesThinBackend(
+        std::string(config->thinLTOPrefixReplace.first),
+        std::string(config->thinLTOPrefixReplace.second),
+        config->thinLTOEmitImportsFiles, indexFile.get(), onIndexWrite);
+  } else {
+    backend = lto::createInProcessThinBackend(
+        llvm::heavyweight_hardware_concurrency(config->thinLTOJobs),
+        onIndexWrite, config->thinLTOEmitIndexFiles,
+        config->thinLTOEmitImportsFiles);
+  }
+
   ltoObj = std::make_unique<lto::LTO>(createConfig(), backend);
 }
 
 void BitcodeCompiler::add(BitcodeFile &f) {
-  ArrayRef<lto::InputFile::Symbol> objSyms = f.obj->symbols();
+  lto::InputFile &obj = *f.obj;
+
+  if (config->thinLTOEmitIndexFiles)
+    thinIndices.insert(obj.getName());
+
+  ArrayRef<lto::InputFile::Symbol> objSyms = obj.symbols();
   std::vector<lto::SymbolResolution> resols;
   resols.reserve(objSyms.size());
 
@@ -117,6 +162,37 @@ void BitcodeCompiler::add(BitcodeFile &f) {
   checkError(ltoObj->add(std::move(f.obj), resols));
 }
 
+// If LazyObjFile has not been added to link, emit empty index files.
+// This is needed because this is what GNU gold plugin does and we have a
+// distributed build system that depends on that behavior.
+static void thinLTOCreateEmptyIndexFiles() {
+  DenseSet<StringRef> linkedBitCodeFiles;
+  for (InputFile *file : inputFiles)
+    if (auto *f = dyn_cast<BitcodeFile>(file))
+      if (!f->lazy)
+        linkedBitCodeFiles.insert(f->getName());
+
+  for (InputFile *file : inputFiles) {
+    if (auto *f = dyn_cast<BitcodeFile>(file)) {
+      if (!f->lazy)
+        continue;
+      if (linkedBitCodeFiles.contains(f->getName()))
+        continue;
+      std::string path =
+          replaceThinLTOSuffix(getThinLTOOutputFile(f->obj->getName()));
+      std::unique_ptr<raw_fd_ostream> os = openFile(path + ".thinlto.bc");
+      if (!os)
+        continue;
+
+      ModuleSummaryIndex m(/*HaveGVs=*/false);
+      m.setSkipModuleByDistributedBackend();
+      writeIndexToFile(m, *os);
+      if (config->thinLTOEmitImportsFiles)
+        openFile(path + ".imports");
+    }
+  }
+}
+
 // Merge all the bitcode files we have seen, codegen the result
 // and return the resulting ObjectFile(s).
 std::vector<ObjFile *> BitcodeCompiler::compile() {
@@ -142,8 +218,16 @@ std::vector<ObjFile *> BitcodeCompiler::compile() {
       },
       cache));
 
-  if (!config->thinLTOCacheDir.empty())
-    pruneCache(config->thinLTOCacheDir, config->thinLTOCachePolicy, files);
+  // Emit empty index files for non-indexed files
+  for (StringRef s : thinIndices) {
+    std::string path = getThinLTOOutputFile(s);
+    openFile(path + ".thinlto.bc");
+    if (config->thinLTOEmitImportsFiles)
+      openFile(path + ".imports");
+  }
+
+  if (config->thinLTOEmitIndexFiles)
+    thinLTOCreateEmptyIndexFiles();
 
   // In ThinLTO mode, Clang passes a temporary directory in -object_path_lto,
   // while the argument is a single file in FullLTO mode.
@@ -162,6 +246,32 @@ std::vector<ObjFile *> BitcodeCompiler::compile() {
     }
   }
 
+  auto outputFilePath = [objPathIsDir](int i) {
+    SmallString<261> filePath("/tmp/lto.tmp");
+    if (!config->ltoObjPath.empty()) {
+      filePath = config->ltoObjPath;
+      if (objPathIsDir)
+        path::append(filePath, Twine(i) + "." +
+                                   getArchitectureName(config->arch()) +
+                                   ".lto.o");
+    }
+    return filePath;
+  };
+
+  // ThinLTO with index only option is required to generate only the index
+  // files. After that, we exit from linker and ThinLTO backend runs in a
+  // distributed environment.
+  if (config->thinLTOIndexOnly) {
+    if (!config->ltoObjPath.empty())
+      saveBuffer(buf[0], outputFilePath(0));
+    if (indexFile)
+      indexFile->close();
+    return {};
+  }
+
+  if (!config->thinLTOCacheDir.empty())
+    pruneCache(config->thinLTOCacheDir, config->thinLTOCachePolicy, files);
+
   std::vector<ObjFile *> ret;
   for (unsigned i = 0; i < maxTasks; ++i) {
     // Get the native object contents either from the cache or from memory.  Do
@@ -183,14 +293,9 @@ std::vector<ObjFile *> BitcodeCompiler::compile() {
       saveBuffer(objBuf,
                  config->outputFile + ((i == 0) ? "" : Twine(i)) + ".lto.o");
 
-    SmallString<261> filePath("/tmp/lto.tmp");
+    auto filePath = outputFilePath(i);
     uint32_t modTime = 0;
     if (!config->ltoObjPath.empty()) {
-      filePath = config->ltoObjPath;
-      if (objPathIsDir)
-        path::append(filePath, Twine(i) + "." +
-                                   getArchitectureName(config->arch()) +
-                                   ".lto.o");
       saveOrHardlinkBuffer(objBuf, filePath, cachePath);
       modTime = getModTime(filePath);
     }

lld/MachO/LTO.h

@@ -9,8 +9,11 @@
 #ifndef LLD_MACHO_LTO_H
 #define LLD_MACHO_LTO_H
 
+#include "lld/Common/LLVM.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/raw_ostream.h"
 #include <memory>
 #include <vector>
 
@@ -34,6 +37,8 @@ class BitcodeCompiler {
   std::unique_ptr<llvm::lto::LTO> ltoObj;
   std::vector<llvm::SmallString<0>> buf;
   std::vector<std::unique_ptr<llvm::MemoryBuffer>> files;
+  std::unique_ptr<llvm::raw_fd_ostream> indexFile;
+  llvm::DenseSet<StringRef> thinIndices;
 };
 
 } // namespace lld::macho

lld/MachO/Options.td

@@ -28,9 +28,22 @@ def color_diagnostics_eq: Joined<["--"], "color-diagnostics=">,
 def threads_eq : Joined<["--"], "threads=">,
     HelpText<"Number of threads. '1' disables multi-threading. By default all available hardware threads are used">,
     Group<grp_lld>;
+def thinlto_emit_imports_files: Flag<["--"], "thinlto-emit-imports-files">,
+    Group<grp_lld>;
+def thinlto_emit_index_files: Flag<["--"], "thinlto-emit-index-files">,
+    Group<grp_lld>;
+def thinlto_index_only: Flag<["--"], "thinlto-index-only">,
+    Group<grp_lld>;
+def thinlto_index_only_eq: Joined<["--"], "thinlto-index-only=">,
+    Group<grp_lld>;
 def thinlto_jobs_eq : Joined<["--"], "thinlto-jobs=">,
     HelpText<"Number of ThinLTO jobs. Default to --threads=">,
     Group<grp_lld>;
+def thinlto_object_suffix_replace_eq:
+    Joined<["--"], "thinlto-object-suffix-replace=">,
+    Group<grp_lld>;
+def thinlto_prefix_replace_eq: Joined<["--"], "thinlto-prefix-replace=">,
+    Group<grp_lld>;
 def reproduce: Separate<["--"], "reproduce">,
     Group<grp_lld>;
 def reproduce_eq: Joined<["--"], "reproduce=">,

lld/test/ELF/lto/thinlto-index-only.ll

@@ -18,18 +18,17 @@
 
 ;; Ensure lld doesn't generates index files when --thinlto-index-only is not enabled.
 ; RUN: rm -f 1.o.thinlto.bc 2.o.thinlto.bc
-; RUN: ld.lld -shared 1.o 2.o -o 4
+; RUN: ld.lld -shared 1.o 2.o -o 5
 ; RUN: not ls 1.o.thinlto.bc
 ; RUN: not ls 2.o.thinlto.bc
 
 ;; Ensure lld generates an index and not a binary if requested.
-; RUN: rm -f 4
 ; RUN: ld.lld --plugin-opt=thinlto-index-only -shared 1.o 2.o -o 4
 ; RUN: llvm-bcanalyzer -dump 1.o.thinlto.bc | FileCheck %s --check-prefix=BACKEND1
 ; RUN: llvm-bcanalyzer -dump 2.o.thinlto.bc | FileCheck %s --check-prefix=BACKEND2
 ; RUN: not test -e 4
 
-;; Ensure lld generates an index even if the file is wrapped in --start-lib/--end-lib
+;; Ensure lld generates an index even if the file is wrapped in --start-lib/--end-lib.
 ; RUN: rm -f 2.o.thinlto.bc 4
 ; RUN: ld.lld --plugin-opt=thinlto-index-only -shared 1.o 3.o --start-lib 2.o --end-lib -o 4
 ; RUN: llvm-dis < 2.o.thinlto.bc | grep -q '\^0 = module:'

lld/test/ELF/lto/thinlto-obj-path.ll

@@ -47,7 +47,7 @@
 
 ;; Ensure lld does not emit empty combined module in default.
 ; RUN: rm -fr objpath && mkdir objpath
-; RUN: ld.lld 1.bc 2.bc -o objpath/a.out --save-temps
+; RUN: ld.lld -shared 1.bc 2.bc -o objpath/a.out --save-temps
 ; RUN: ls objpath/a.out*.lto.* | count 2
 
 ; EMPTY:     file format elf64-x86-64