[ThinLTO][ELF] Add --thinlto-emit-index-files option

Northbadge · Northbadge · commit 22f1273357cf · 2022-06-23T12:35:42.000-07:00
Allows ThinLTO indices to be written to disk on-the-fly/as-part-of “normal” linker execution. Previously ThinLTO indices could be written via --thinlto-index-only but that would cause the linker to exit early. For MLGO specifically, this enables saving the ThinLTO index files without having to restart the linker to collect data only available at later stages (i.e. output of --save-temps) of the linker's execution. Note, this option does not currently work with: --thinlto-object-suffix-replace, as this is intended to be used to consume minimized IR bitcode files while --thinlto-emit-index-files is intended to be run together with InProcessThinLTO (which cannot parse minimized IR). --thinlto-prefix-replace support is left unimplemented but can be implemented if needed Differential Revision: https://reviews.llvm.org/D127777
diff --git a/lld/ELF/Config.h b/lld/ELF/Config.h
@@ -221,6 +221,7 @@ struct Configuration {
   bool target1Rel;
   bool trace;
   bool thinLTOEmitImportsFiles;
+  bool thinLTOEmitIndexFiles;
   bool thinLTOIndexOnly;
   bool timeTraceEnabled;
   bool tocOptimize;
diff --git a/lld/ELF/Driver.cpp b/lld/ELF/Driver.cpp
@@ -1175,13 +1175,24 @@ static void readConfigs(opt::InputArgList &args) {
       parseCachePruningPolicy(args.getLastArgValue(OPT_thinlto_cache_policy)),
       "--thinlto-cache-policy: invalid cache policy");
   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->thinLTOModulesToCompile =
       args::getStrings(args, OPT_thinlto_single_module_eq);
   config->timeTraceEnabled = args.hasArg(OPT_time_trace);
diff --git a/lld/ELF/LTO.cpp b/lld/ELF/LTO.cpp
@@ -191,15 +191,17 @@ BitcodeCompiler::BitcodeCompiler() {
 
   // Initialize ltoObj.
   lto::ThinBackend backend;
+  auto onIndexWrite = [&](StringRef s) { thinIndices.erase(s); };
   if (config->thinLTOIndexOnly) {
-    auto onIndexWrite = [&](StringRef s) { thinIndices.erase(s); };
     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));
+        llvm::heavyweight_hardware_concurrency(config->thinLTOJobs),
+        onIndexWrite, config->thinLTOEmitIndexFiles,
+        config->thinLTOEmitImportsFiles);
   }
 
   ltoObj = std::make_unique<lto::LTO>(createConfig(), backend,
@@ -224,7 +226,7 @@ void BitcodeCompiler::add(BitcodeFile &f) {
   lto::InputFile &obj = *f.obj;
   bool isExec = !config->shared && !config->relocatable;
 
-  if (config->thinLTOIndexOnly)
+  if (config->thinLTOEmitIndexFiles)
     thinIndices.insert(obj.getName());
 
   ArrayRef<Symbol *> syms = f.getSymbols();
@@ -339,9 +341,10 @@ std::vector<InputFile *> BitcodeCompiler::compile() {
     }
   }
 
-  if (config->thinLTOIndexOnly) {
+  if (config->thinLTOEmitIndexFiles)
     thinLTOCreateEmptyIndexFiles();
 
+  if (config->thinLTOIndexOnly) {
     if (!config->ltoObjPath.empty())
       saveBuffer(buf[0], config->ltoObjPath);
 
diff --git a/lld/ELF/Options.td b/lld/ELF/Options.td
@@ -596,6 +596,7 @@ def thinlto_cache_dir: JJ<"thinlto-cache-dir=">,
   HelpText<"Path to ThinLTO cached object file directory">;
 defm thinlto_cache_policy: EEq<"thinlto-cache-policy", "Pruning policy for the ThinLTO cache">;
 def thinlto_emit_imports_files: FF<"thinlto-emit-imports-files">;
+def thinlto_emit_index_files: FF<"thinlto-emit-index-files">;
 def thinlto_index_only: FF<"thinlto-index-only">;
 def thinlto_index_only_eq: JJ<"thinlto-index-only=">;
 def thinlto_jobs: JJ<"thinlto-jobs=">,
diff --git a/lld/test/ELF/lto/thinlto-emit-index.ll b/lld/test/ELF/lto/thinlto-emit-index.ll
@@ -0,0 +1,107 @@
+; REQUIRES: x86
+
+;; Mostly copied/updated from thinlto-index-only.ll
+;; First ensure that the ThinLTO handling in lld handles
+;; bitcode without summary sections gracefully and generates index file.
+; RUN: rm -rf %t.dir && mkdir %t.dir && cd %t.dir
+; RUN: llvm-as %s -o 1.o
+; RUN: llvm-as %p/Inputs/thinlto.ll -o 2.o
+; RUN: ld.lld --thinlto-emit-index-files -shared 1.o 2.o -o 3
+; RUN: ls 2.o.thinlto.bc
+; RUN: ls 3
+; RUN: ld.lld -shared 1.o 2.o -o 3
+; RUN: llvm-nm 3 | FileCheck %s --check-prefix=NM
+
+;; Basic ThinLTO tests.
+; RUN: opt -module-summary %s -o 1.o
+; RUN: opt -module-summary %p/Inputs/thinlto.ll -o 2.o
+; RUN: opt -module-summary %p/Inputs/thinlto_empty.ll -o 3.o
+
+;; Ensure lld generates an index and also a binary if requested.
+; RUN: ld.lld --thinlto-emit-index-files -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: ls 4
+
+;; Ensure lld generates an index and not a binary if both emit-index and index-only are present.
+; RUN: ld.lld --thinlto-emit-index-files --thinlto-index-only -shared 1.o 2.o -o 5
+; RUN: not ls 5
+
+;; Ensure lld generates an index even if the file is wrapped in --start-lib/--end-lib
+; RUN: rm -f 2.o.thinlto.bc
+; RUN: ld.lld --thinlto-emit-index-files -shared 1.o 3.o --start-lib 2.o --end-lib -o 6
+; RUN: llvm-dis < 2.o.thinlto.bc | grep -q '\^0 = module:'
+; RUN: ls 6
+
+;; Test that LLD generates an empty index even for lazy object file that is not added to link.
+;; Test that LLD also generates empty imports file with the --thinlto-emit-imports-files option.
+; RUN: rm -f 1.o.thinlto.bc 1.o.imports
+; RUN: ld.lld --thinlto-emit-index-files -shared 2.o --start-lib 1.o --end-lib \
+; RUN: --thinlto-emit-imports-files -o 7
+; RUN: ls 7
+; RUN: ls 1.o.thinlto.bc
+; RUN: ls 1.o.imports
+
+;; Ensure LLD generates an empty index for each bitcode file even if all bitcode files are lazy.
+; RUN: rm -f 1.o.thinlto.bc
+; RUN: llvm-mc -filetype=obj -triple=x86_64-unknown-linux-gnu /dev/null -o dummy.o
+; RUN: ld.lld --thinlto-emit-index-files -shared dummy.o --start-lib 1.o --end-lib -o 8
+; RUN: ls 8
+; RUN: ls 1.o.thinlto.bc
+
+;; Test that LLD errors out when run with suffix replacement, or prefix replacement
+; RUN: not ld.lld --thinlto-emit-index-files -shared 2.o --start-lib 1.o --end-lib \
+; RUN: --thinlto-prefix-replace="abc;xyz" 2>&1 | FileCheck %s --check-prefix=ERR1
+; ERR1: --thinlto-prefix-replace is not supported with --thinlto-emit-index-files
+
+; RUN: not ld.lld --thinlto-emit-index-files -shared 2.o --start-lib 1.o --end-lib \
+; RUN: --thinlto-object-suffix-replace="abc;xyz" 2>&1 | FileCheck %s --check-prefix=ERR2
+; ERR2: --thinlto-object-suffix-replace is not supported with --thinlto-emit-index-files
+
+;; But not when passed with index only as well
+; RUN: ld.lld --thinlto-emit-index-files -shared 2.o --start-lib 1.o --end-lib \
+; RUN: --thinlto-prefix-replace="abc;xyz" --thinlto-index-only
+
+; RUN: ld.lld --thinlto-emit-index-files -shared 2.o --start-lib 1.o --end-lib \
+; RUN: --thinlto-object-suffix-replace="abc;xyz" --thinlto-index-only
+
+; NM: T f
+
+;; The backend index for this module contains summaries from itself and
+;; Inputs/thinlto.ll, as it imports from the latter.
+; BACKEND1: <MODULE_STRTAB_BLOCK
+; BACKEND1-NEXT: <ENTRY {{.*}} record string = '1.o'
+; BACKEND1-NEXT: <ENTRY {{.*}} record string = '2.o'
+; BACKEND1-NEXT: </MODULE_STRTAB_BLOCK
+; BACKEND1: <GLOBALVAL_SUMMARY_BLOCK
+; BACKEND1: <VERSION
+; BACKEND1: <FLAGS
+; BACKEND1: <VALUE_GUID op0={{1|2}} op1={{-3706093650706652785|-5300342847281564238}}
+; BACKEND1: <VALUE_GUID op0={{1|2}} op1={{-3706093650706652785|-5300342847281564238}}
+; BACKEND1: <COMBINED
+; BACKEND1: <COMBINED
+; BACKEND1: </GLOBALVAL_SUMMARY_BLOCK
+
+;; The backend index for Input/thinlto.ll contains summaries from itself only,
+;; as it does not import anything.
+; BACKEND2: <MODULE_STRTAB_BLOCK
+; BACKEND2-NEXT: <ENTRY {{.*}} record string = '2.o'
+; BACKEND2-NEXT: </MODULE_STRTAB_BLOCK
+; BACKEND2-NEXT: <GLOBALVAL_SUMMARY_BLOCK
+; BACKEND2-NEXT: <VERSION
+; BACKEND2-NEXT: <FLAGS
+; BACKEND2-NEXT: <VALUE_GUID op0=1 op1=-5300342847281564238
+; BACKEND2-NEXT: <COMBINED
+; BACKEND2-NEXT: <BLOCK_COUNT op0=2/>
+; BACKEND2-NEXT: </GLOBALVAL_SUMMARY_BLOCK
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @g(...)
+
+define void @f() {
+entry:
+  call void (...) @g()
+  ret void
+}
diff --git a/llvm/include/llvm/LTO/LTO.h b/llvm/include/llvm/LTO/LTO.h
@@ -197,7 +197,17 @@ using ThinBackend = std::function<std::unique_ptr<ThinBackendProc>(
 
 /// This ThinBackend runs the individual backend jobs in-process.
 /// The default value means to use one job per hardware core (not hyper-thread).
-ThinBackend createInProcessThinBackend(ThreadPoolStrategy Parallelism);
+/// OnWrite is callback which receives module identifier and notifies LTO user
+/// that index file for the module (and optionally imports file) was created.
+/// ShouldEmitIndexFiles being true will write sharded ThinLTO index files
+/// to the same path as the input module, with suffix ".thinlto.bc"
+/// ShouldEmitImportsFiles is true it also writes a list of imported files to a
+/// similar path with ".imports" appended instead.
+using IndexWriteCallback = std::function<void(const std::string &)>;
+ThinBackend createInProcessThinBackend(ThreadPoolStrategy Parallelism,
+                                       IndexWriteCallback OnWrite = nullptr,
+                                       bool ShouldEmitIndexFiles = false,
+                                       bool ShouldEmitImportsFiles = false);
 
 /// This ThinBackend writes individual module indexes to files, instead of
 /// running the individual backend jobs. This backend is for distributed builds
@@ -212,7 +222,6 @@ ThinBackend createInProcessThinBackend(ThreadPoolStrategy Parallelism);
 /// the final ThinLTO linking. Can be nullptr.
 /// OnWrite is callback which receives module identifier and notifies LTO user
 /// that index file for the module (and optionally imports file) was created.
-using IndexWriteCallback = std::function<void(const std::string &)>;
 ThinBackend createWriteIndexesThinBackend(std::string OldPrefix,
                                           std::string NewPrefix,
                                           bool ShouldEmitImportsFiles,
diff --git a/llvm/lib/LTO/LTO.cpp b/llvm/lib/LTO/LTO.cpp
@@ -1161,12 +1161,16 @@ class lto::ThinBackendProc {
   const Config &Conf;
   ModuleSummaryIndex &CombinedIndex;
   const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries;
+  lto::IndexWriteCallback OnWrite;
+  bool ShouldEmitImportsFiles;
 
 public:
   ThinBackendProc(const Config &Conf, ModuleSummaryIndex &CombinedIndex,
-                  const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries)
+                  const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
+                  lto::IndexWriteCallback OnWrite, bool ShouldEmitImportsFiles)
       : Conf(Conf), CombinedIndex(CombinedIndex),
-        ModuleToDefinedGVSummaries(ModuleToDefinedGVSummaries) {}
+        ModuleToDefinedGVSummaries(ModuleToDefinedGVSummaries),
+        OnWrite(OnWrite), ShouldEmitImportsFiles(ShouldEmitImportsFiles) {}
 
   virtual ~ThinBackendProc() = default;
   virtual Error start(
@@ -1177,6 +1181,30 @@ class lto::ThinBackendProc {
       MapVector<StringRef, BitcodeModule> &ModuleMap) = 0;
   virtual Error wait() = 0;
   virtual unsigned getThreadCount() = 0;
+
+  // Write sharded indices and (optionally) imports to disk
+  Error emitFiles(const FunctionImporter::ImportMapTy &ImportList,
+                  llvm::StringRef ModulePath,
+                  const std::string &NewModulePath) {
+    std::map<std::string, GVSummaryMapTy> ModuleToSummariesForIndex;
+    std::error_code EC;
+    gatherImportedSummariesForModule(ModulePath, ModuleToDefinedGVSummaries,
+                                     ImportList, ModuleToSummariesForIndex);
+
+    raw_fd_ostream OS(NewModulePath + ".thinlto.bc", EC,
+                      sys::fs::OpenFlags::OF_None);
+    if (EC)
+      return errorCodeToError(EC);
+    writeIndexToFile(CombinedIndex, OS, &ModuleToSummariesForIndex);
+
+    if (ShouldEmitImportsFiles) {
+      EC = EmitImportsFiles(ModulePath, NewModulePath + ".imports",
+                            ModuleToSummariesForIndex);
+      if (EC)
+        return errorCodeToError(EC);
+    }
+    return Error::success();
+  }
 };
 
 namespace {
@@ -1190,15 +1218,19 @@ class InProcessThinBackend : public ThinBackendProc {
   Optional<Error> Err;
   std::mutex ErrMu;
 
+  bool ShouldEmitIndexFiles;
+
 public:
   InProcessThinBackend(
       const Config &Conf, ModuleSummaryIndex &CombinedIndex,
       ThreadPoolStrategy ThinLTOParallelism,
       const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-      AddStreamFn AddStream, FileCache Cache)
-      : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries),
+      AddStreamFn AddStream, FileCache Cache, lto::IndexWriteCallback OnWrite,
+      bool ShouldEmitIndexFiles, bool ShouldEmitImportsFiles)
+      : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries,
+                        OnWrite, ShouldEmitImportsFiles),
         BackendThreadPool(ThinLTOParallelism), AddStream(std::move(AddStream)),
-        Cache(std::move(Cache)) {
+        Cache(std::move(Cache)), ShouldEmitIndexFiles(ShouldEmitIndexFiles) {
     for (auto &Name : CombinedIndex.cfiFunctionDefs())
       CfiFunctionDefs.insert(
           GlobalValue::getGUID(GlobalValue::dropLLVMManglingEscape(Name)));
@@ -1227,6 +1259,11 @@ class InProcessThinBackend : public ThinBackendProc {
 
     auto ModuleID = BM.getModuleIdentifier();
 
+    if (ShouldEmitIndexFiles) {
+      if (auto E = emitFiles(ImportList, ModuleID, ModuleID.str()))
+        return E;
+    }
+
     if (!Cache || !CombinedIndex.modulePaths().count(ModuleID) ||
         all_of(CombinedIndex.getModuleHash(ModuleID),
                [](uint32_t V) { return V == 0; }))
@@ -1285,6 +1322,9 @@ class InProcessThinBackend : public ThinBackendProc {
         },
         BM, std::ref(CombinedIndex), std::ref(ImportList), std::ref(ExportList),
         std::ref(ResolvedODR), std::ref(DefinedGlobals), std::ref(ModuleMap));
+
+    if (OnWrite)
+      OnWrite(std::string(ModulePath));
     return Error::success();
   }
 
@@ -1302,13 +1342,16 @@ class InProcessThinBackend : public ThinBackendProc {
 };
 } // end anonymous namespace
 
-ThinBackend lto::createInProcessThinBackend(ThreadPoolStrategy Parallelism) {
+ThinBackend lto::createInProcessThinBackend(ThreadPoolStrategy Parallelism,
+                                            lto::IndexWriteCallback OnWrite,
+                                            bool ShouldEmitIndexFiles,
+                                            bool ShouldEmitImportsFiles) {
   return [=](const Config &Conf, ModuleSummaryIndex &CombinedIndex,
              const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
              AddStreamFn AddStream, FileCache Cache) {
     return std::make_unique<InProcessThinBackend>(
         Conf, CombinedIndex, Parallelism, ModuleToDefinedGVSummaries, AddStream,
-        Cache);
+        Cache, OnWrite, ShouldEmitIndexFiles, ShouldEmitImportsFiles);
   };
 }
 
@@ -1335,20 +1378,18 @@ std::string lto::getThinLTOOutputFile(const std::string &Path,
 namespace {
 class WriteIndexesThinBackend : public ThinBackendProc {
   std::string OldPrefix, NewPrefix;
-  bool ShouldEmitImportsFiles;
   raw_fd_ostream *LinkedObjectsFile;
-  lto::IndexWriteCallback OnWrite;
 
 public:
   WriteIndexesThinBackend(
       const Config &Conf, ModuleSummaryIndex &CombinedIndex,
       const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
       std::string OldPrefix, std::string NewPrefix, bool ShouldEmitImportsFiles,
       raw_fd_ostream *LinkedObjectsFile, lto::IndexWriteCallback OnWrite)
-      : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries),
+      : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries,
+                        OnWrite, ShouldEmitImportsFiles),
         OldPrefix(OldPrefix), NewPrefix(NewPrefix),
-        ShouldEmitImportsFiles(ShouldEmitImportsFiles),
-        LinkedObjectsFile(LinkedObjectsFile), OnWrite(OnWrite) {}
+        LinkedObjectsFile(LinkedObjectsFile) {}
 
   Error start(
       unsigned Task, BitcodeModule BM,
@@ -1363,23 +1404,8 @@ class WriteIndexesThinBackend : public ThinBackendProc {
     if (LinkedObjectsFile)
       *LinkedObjectsFile << NewModulePath << '\n';
 
-    std::map<std::string, GVSummaryMapTy> ModuleToSummariesForIndex;
-    gatherImportedSummariesForModule(ModulePath, ModuleToDefinedGVSummaries,
-                                     ImportList, ModuleToSummariesForIndex);
-
-    std::error_code EC;
-    raw_fd_ostream OS(NewModulePath + ".thinlto.bc", EC,
-                      sys::fs::OpenFlags::OF_None);
-    if (EC)
-      return errorCodeToError(EC);
-    writeIndexToFile(CombinedIndex, OS, &ModuleToSummariesForIndex);
-
-    if (ShouldEmitImportsFiles) {
-      EC = EmitImportsFiles(ModulePath, NewModulePath + ".imports",
-                            ModuleToSummariesForIndex);
-      if (EC)
-        return errorCodeToError(EC);
-    }
+    if (auto E = emitFiles(ImportList, ModulePath, NewModulePath))
+      return E;
 
     if (OnWrite)
       OnWrite(std::string(ModulePath));
