[ThinLTO] Option to bypass function importing.

This completes the circle, complementing -lto-embed-bitcode
(specifically, post-merge-pre-opt). Using -thinlto-assume-merged skips
function importing. The index file is still needed for the other data it
contains.

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

include/llvm/LTO/LTOBackend.h

@@ -49,6 +49,22 @@ Error thinBackend(const Config &C, unsigned Task, AddStreamFn AddStream,
 
 Error finalizeOptimizationRemarks(
     std::unique_ptr<ToolOutputFile> DiagOutputFile);
+
+/// Returns the BitcodeModule that is ThinLTO.
+BitcodeModule *findThinLTOModule(MutableArrayRef<BitcodeModule> BMs);
+
+/// Variant of the above.
+Expected<BitcodeModule> findThinLTOModule(MemoryBufferRef MBRef);
+
+/// Distributed ThinLTO: load the referenced modules, keeping their buffers
+/// alive in the provided OwnedImportLifetimeManager. Returns false if the
+/// operation failed.
+bool loadReferencedModules(
+    const Module &M, const ModuleSummaryIndex &CombinedIndex,
+    FunctionImporter::ImportMapTy &ImportList,
+    MapVector<llvm::StringRef, llvm::BitcodeModule> &ModuleMap,
+    std::vector<std::unique_ptr<llvm::MemoryBuffer>>
+        &OwnedImportsLifetimeManager);
 }
 }
 

lib/LTO/LTOBackend.cpp

@@ -69,6 +69,11 @@ static cl::opt<LTOBitcodeEmbedding> EmbedBitcode(
                           "Embed post merge, but before optimizations")),
     cl::desc("Embed LLVM bitcode in object files produced by LTO"));
 
+static cl::opt<bool> ThinLTOAssumeMerged(
+    "thinlto-assume-merged", cl::init(false),
+    cl::desc("Assume the input has already undergone ThinLTO function "
+             "importing and the other pre-optimization pipeline changes."));
+
 LLVM_ATTRIBUTE_NORETURN static void reportOpenError(StringRef Path, Twine Msg) {
   errs() << "failed to open " << Path << ": " << Msg << '\n';
   errs().flush();
@@ -583,6 +588,21 @@ Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   if (Conf.PreOptModuleHook && !Conf.PreOptModuleHook(Task, Mod))
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 
+  auto OptimizeAndCodegen =
+      [&](Module &Mod, TargetMachine *TM,
+          std::unique_ptr<ToolOutputFile> DiagnosticOutputFile) {
+        if (!opt(Conf, TM, Task, Mod, /*IsThinLTO=*/true,
+                 /*ExportSummary=*/nullptr, /*ImportSummary=*/&CombinedIndex,
+                 CmdArgs))
+          return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
+
+        codegen(Conf, TM, AddStream, Task, Mod, CombinedIndex);
+        return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
+      };
+
+  if (ThinLTOAssumeMerged)
+    return OptimizeAndCodegen(Mod, TM.get(), std::move(DiagnosticOutputFile));
+
   // When linking an ELF shared object, dso_local should be dropped. We
   // conservatively do this for -fpic.
   bool ClearDSOLocalOnDeclarations =
@@ -623,11 +643,81 @@ Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   if (Conf.PostImportModuleHook && !Conf.PostImportModuleHook(Task, Mod))
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 
-  if (!opt(Conf, TM.get(), Task, Mod, /*IsThinLTO=*/true,
+  return OptimizeAndCodegen(Mod, TM.get(), std::move(DiagnosticOutputFile));
-           /*ExportSummary=*/nullptr, /*ImportSummary=*/&CombinedIndex,
-           CmdArgs))
-    return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
-
-  codegen(Conf, TM.get(), AddStream, Task, Mod, CombinedIndex);
-  return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 }
+
+BitcodeModule *lto::findThinLTOModule(MutableArrayRef<BitcodeModule> BMs) {
+  if (ThinLTOAssumeMerged && BMs.size() == 1)
+    return BMs.begin();
+
+  for (BitcodeModule &BM : BMs) {
+    Expected<BitcodeLTOInfo> LTOInfo = BM.getLTOInfo();
+    if (LTOInfo && LTOInfo->IsThinLTO)
+      return &BM;
+  }
+  return nullptr;
+}
+
+Expected<BitcodeModule> lto::findThinLTOModule(MemoryBufferRef MBRef) {
+  Expected<std::vector<BitcodeModule>> BMsOrErr = getBitcodeModuleList(MBRef);
+  if (!BMsOrErr)
+    return BMsOrErr.takeError();
+
+  // The bitcode file may contain multiple modules, we want the one that is
+  // marked as being the ThinLTO module.
+  if (const BitcodeModule *Bm = lto::findThinLTOModule(*BMsOrErr))
+    return *Bm;
+
+  return make_error<StringError>("Could not find module summary",
+                                 inconvertibleErrorCode());
+}
+
+bool lto::loadReferencedModules(
+    const Module &M, const ModuleSummaryIndex &CombinedIndex,
+    FunctionImporter::ImportMapTy &ImportList,
+    MapVector<llvm::StringRef, llvm::BitcodeModule> &ModuleMap,
+    std::vector<std::unique_ptr<llvm::MemoryBuffer>>
+        &OwnedImportsLifetimeManager) {
+  if (ThinLTOAssumeMerged)
+    return true;
+  // We can simply import the values mentioned in the combined index, since
+  // we should only invoke this using the individual indexes written out
+  // via a WriteIndexesThinBackend.
+  for (const auto &GlobalList : CombinedIndex) {
+    // Ignore entries for undefined references.
+    if (GlobalList.second.SummaryList.empty())
+      continue;
+
+    auto GUID = GlobalList.first;
+    for (const auto &Summary : GlobalList.second.SummaryList) {
+      // Skip the summaries for the importing module. These are included to
+      // e.g. record required linkage changes.
+      if (Summary->modulePath() == M.getModuleIdentifier())
+        continue;
+      // Add an entry to provoke importing by thinBackend.
+      ImportList[Summary->modulePath()].insert(GUID);
+    }
+  }
+
+  for (auto &I : ImportList) {
+    ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> MBOrErr =
+        llvm::MemoryBuffer::getFile(I.first());
+    if (!MBOrErr) {
+      errs() << "Error loading imported file '" << I.first()
+             << "': " << MBOrErr.getError().message() << "\n";
+      return false;
+    }
+
+    Expected<BitcodeModule> BMOrErr = findThinLTOModule(**MBOrErr);
+    if (!BMOrErr) {
+      handleAllErrors(BMOrErr.takeError(), [&](ErrorInfoBase &EIB) {
+        errs() << "Error loading imported file '" << I.first()
+               << "': " << EIB.message() << '\n';
+      });
+      return false;
+    }
+    ModuleMap.insert({I.first(), *BMOrErr});
+    OwnedImportsLifetimeManager.push_back(std::move(*MBOrErr));
+  }
+  return true;
+}