[ThinLTO] During module importing, close one source module before open another one for distributed mode.

[wmi]

Summary

Currently during module importing, ThinLTO opens all the source modules, collect functions to be imported and append them to the destination module, then leave all the modules open through out the lto backend pipeline. This patch refactors it in the way that one source module will be closed before another source module is opened. All the source modules will be closed after importing phase is done. It will save some amount of memory when there are many source modules to be imported.

Note that this patch only changes the distributed thinlto mode. For in process thinlto mode, one source module is shared acorss different thinlto backend threads so it is not changed in this patch.

[tejohnson]

LGTM with a couple of comment requests noted below.

clang/lib/CodeGen/BackendUtil.cpp

1583

Document constant parameter

llvm/include/llvm/LTO/LTOBackend.h

48–53

Can you add a note here about the behavior with regards to ModuleMap depending on whether it is null or not?

[wmi]

Address Teresa's comment.

clang/lib/CodeGen/BackendUtil.cpp

1583

Done.

llvm/include/llvm/LTO/LTOBackend.h

48–53

Add the comment.

[wenlei]

Does this help non-distributed ThinLTO as well? cc: @weiwang

[tejohnson]

In D99554#2674229, @wenlei wrote:

Does this help non-distributed ThinLTO as well? cc: @weiwang

No. See the note in Wei's description about this:

Note that this patch only changes the distributed thinlto mode. For in process thinlto mode, one source module is shared across different thinlto backend threads so it is not changed in this patch.

Specifically, for in-process ThinLTO all modules are kept open so that we don't need to reload a module for each thread that wants to import from it.

clang/lib/CodeGen/BackendUtil.cpp

static void runThinLTOBackend(
   CombinedIndex->collectDefinedGVSummariesPerModule(ModuleToDefinedGVSummaries);
 
   setCommandLineOpts(CGOpts);
 
   // 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.
   FunctionImporter::ImportMapTy ImportList;
-  std::vector<std::unique_ptr<llvm::MemoryBuffer>> OwnedImports;
-  MapVector<llvm::StringRef, llvm::BitcodeModule> ModuleMap;
-  if (!lto::loadReferencedModules(*M, *CombinedIndex, ImportList, ModuleMap,
-                                  OwnedImports))
+  if (!lto::initImportList(*M, *CombinedIndex, ImportList))
     return;
 
   auto AddStream = [&](size_t Task) {
     return std::make_unique<lto::NativeObjectStream>(std::move(OS));
   };
   lto::Config Conf;
   if (CGOpts.SaveTempsFilePrefix != "") {
     if (Error E = Conf.addSaveTemps(CGOpts.SaveTempsFilePrefix + ".",
case Backend_EmitBC:
     break;
   default:
     Conf.CGFileType = getCodeGenFileType(Action);
     break;
   }
   if (Error E =
           thinBackend(Conf, -1, AddStream, *M, *CombinedIndex, ImportList,
                       ModuleToDefinedGVSummaries[M->getModuleIdentifier()],
-                      ModuleMap, CGOpts.CmdArgs)) {
+                      /* ModuleMap */ nullptr, CGOpts.CmdArgs)) {
     handleAllErrors(std::move(E), [&](ErrorInfoBase &EIB) {
       errs() << "Error running ThinLTO backend: " << EIB.message() << '\n';
     });
   }
 }
 
 void clang::EmitBackendOutput(DiagnosticsEngine &Diags,
                               const HeaderSearchOptions &HeaderOpts,

clang/test/CodeGen/thinlto_backend.ll

 ; CHECK-IMPORT: define available_externally void @f2()
 ; RUN: llvm-nm %t3.o | FileCheck --check-prefix=CHECK-OBJ %s
 ; CHECK-OBJ: T f1
 ; CHECK-OBJ-NOT: U f2
 
 ; Ensure we get expected error for input files without summaries
 ; RUN: opt -o %t2.o %s
 ; RUN: %clang -target x86_64-unknown-linux-gnu -O2 -o %t3.o -x ir %t1.o -c -fthinlto-index=%t.thinlto.bc 2>&1 | FileCheck %s -check-prefix=CHECK-ERROR2
-; CHECK-ERROR2: Error loading imported file '{{.*}}': Could not find module summary
+; CHECK-ERROR2: Error loading imported file {{.*}}: Could not find module summary
 
 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 @f2()
 declare i8* @f3()
 
 define void @f1() {
   call void @f2()
   ; Make sure that the backend can handle undefined references.
   ; Do an indirect call so that the undefined ref shows up in the combined index.
   call void bitcast (i8*()* @f3 to void()*)()
   ret void
 }

llvm/include/llvm/LTO/LTOBackend.h

bool opt(const Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
          const std::vector<uint8_t> &CmdArgs);
 
 /// Runs a regular LTO backend. The regular LTO backend can also act as the
 /// regular LTO phase of ThinLTO, which may need to access the combined index.
 Error backend(const Config &C, AddStreamFn AddStream,
               unsigned ParallelCodeGenParallelismLevel, Module &M,
               ModuleSummaryIndex &CombinedIndex);
 
 /// Runs a ThinLTO backend.
+/// If \p ModuleMap is not nullptr, all the module files to be imported have
+/// already been mapped to memory and the corresponding BitcodeModule objects
+/// are saved in the ModuleMap. If \p ModuleMap is nullptr, module files will
+/// be mapped to memory on demand and at any given time during importing, only
+/// one source module will be kept open at the most.
 Error thinBackend(const Config &C, unsigned Task, AddStreamFn AddStream,
                   Module &M, const ModuleSummaryIndex &CombinedIndex,
                   const FunctionImporter::ImportMapTy &ImportList,
                   const GVSummaryMapTy &DefinedGlobals,
-                  MapVector<StringRef, BitcodeModule> &ModuleMap,
+                  MapVector<StringRef, BitcodeModule> *ModuleMap,
                   const std::vector<uint8_t> &CmdArgs = std::vector<uint8_t>());
 
 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
+/// Distributed ThinLTO: collect the referenced modules based on
+/// module summary and initialize ImportList. 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);
+bool initImportList(const Module &M, const ModuleSummaryIndex &CombinedIndex,
+                    FunctionImporter::ImportMapTy &ImportList);
 }
 }
 
 #endif

llvm/lib/LTO/LTO.cpp

Error runThinLTOBackendThread(
       MapVector<StringRef, BitcodeModule> &ModuleMap) {
     auto RunThinBackend = [&](AddStreamFn AddStream) {
       LTOLLVMContext BackendContext(Conf);
       Expected<std::unique_ptr<Module>> MOrErr = BM.parseModule(BackendContext);
       if (!MOrErr)
         return MOrErr.takeError();
 
       return thinBackend(Conf, Task, AddStream, **MOrErr, CombinedIndex,
-                         ImportList, DefinedGlobals, ModuleMap);
+                         ImportList, DefinedGlobals, &ModuleMap);
     };
 
     auto ModuleID = BM.getModuleIdentifier();
 
     if (!Cache || !CombinedIndex.modulePaths().count(ModuleID) ||
         all_of(CombinedIndex.getModuleHash(ModuleID),
                [](uint32_t V) { return V == 0; }))
       // Cache disabled or no entry for this module in the combined index or

llvm/lib/LTO/LTOBackend.cpp

if (GV->use_empty())
       GV->eraseFromParent();
   }
 }
 
 Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
                        Module &Mod, const ModuleSummaryIndex &CombinedIndex,
                        const FunctionImporter::ImportMapTy &ImportList,
                        const GVSummaryMapTy &DefinedGlobals,
-                       MapVector<StringRef, BitcodeModule> &ModuleMap,
+                       MapVector<StringRef, BitcodeModule> *ModuleMap,
                        const std::vector<uint8_t> &CmdArgs) {
   Expected<const Target *> TOrErr = initAndLookupTarget(Conf, Mod);
   if (!TOrErr)
     return TOrErr.takeError();
 
   std::unique_ptr<TargetMachine> TM = createTargetMachine(Conf, *TOrErr, Mod);
 
   // Setup optimization remarks.
Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
 
   if (Conf.PostInternalizeModuleHook &&
       !Conf.PostInternalizeModuleHook(Task, Mod))
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 
   auto ModuleLoader = [&](StringRef Identifier) {
     assert(Mod.getContext().isODRUniquingDebugTypes() &&
            "ODR Type uniquing should be enabled on the context");
-    auto I = ModuleMap.find(Identifier);
-    assert(I != ModuleMap.end());
+    if (ModuleMap) {
+      auto I = ModuleMap->find(Identifier);
+      assert(I != ModuleMap->end());
       return I->second.getLazyModule(Mod.getContext(),
                                      /*ShouldLazyLoadMetadata=*/true,
                                      /*IsImporting*/ true);
+    }
+
+    ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> MBOrErr =
+        llvm::MemoryBuffer::getFile(Identifier);
+    if (!MBOrErr)
+      return Expected<std::unique_ptr<llvm::Module>>(make_error<StringError>(
+          Twine("Error loading imported file ") + Identifier + " : ",
+          MBOrErr.getError()));
+
+    Expected<BitcodeModule> BMOrErr = findThinLTOModule(**MBOrErr);
+    if (!BMOrErr)
+      return Expected<std::unique_ptr<llvm::Module>>(make_error<StringError>(
+          Twine("Error loading imported file ") + Identifier + " : " +
+              toString(BMOrErr.takeError()),
+          inconvertibleErrorCode()));
+
+    Expected<std::unique_ptr<Module>> MOrErr =
+        BMOrErr->getLazyModule(Mod.getContext(),
+                               /*ShouldLazyLoadMetadata=*/true,
+                               /*IsImporting*/ true);
+    if (MOrErr)
+      (*MOrErr)->setOwnedMemoryBuffer(std::move(*MBOrErr));
+    return MOrErr;
   };
 
   FunctionImporter Importer(CombinedIndex, ModuleLoader,
                             ClearDSOLocalOnDeclarations);
   if (Error Err = Importer.importFunctions(Mod, ImportList).takeError())
     return Err;
 
   if (Conf.PostImportModuleHook && !Conf.PostImportModuleHook(Task, Mod))
Expected<BitcodeModule> lto::findThinLTOModule(MemoryBufferRef MBRef) {
   // 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) {
+bool lto::initImportList(const Module &M,
+                         const ModuleSummaryIndex &CombinedIndex,
+                         FunctionImporter::ImportMapTy &ImportList) {
   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;
 }