[ThinLTO] Use bulk importing in llvm-link

[tejohnson]

Summary

Along with bulk importing we can avoid the use of post-pass metadata
linking.

[mehdi_amini]

LGTM.
See inline comments.

tools/llvm-link/llvm-link.cpp

130

It can be used with the FunctionImporter

The function importer is not involved here right? Comment may be updated.

259

Do we actually want bulk import for llvm-link? Maybe doing one import at a time is a "feature", since it not testing the same thing.
It's probably good for now...

[tejohnson]

tools/llvm-link/llvm-link.cpp

130

Will fix. This was pretty much just cloned with minor tweaks to the arguments from what is in the function importer pass, since we want to do the same thing here.

259

Yes, we do need to do the bulk importing to get the metadata importing to work as expected without post-pass metadata linking.

tools/llvm-link/llvm-link.cpp

 static std::unique_ptr<Module> loadFile(const char *argv0,
                                         const std::string &FN,
                                         LLVMContext &Context,
                                         bool MaterializeMetadata = true) {
   SMDiagnostic Err;
   if (Verbose) errs() << "Loading '" << FN << "'\n";
   std::unique_ptr<Module> Result =
       getLazyIRFileModule(FN, Err, Context, !MaterializeMetadata);
-  if (!Result)
+  if (!Result) {
     Err.print(argv0, errs());
+    return nullptr;
+  }
 
   if (MaterializeMetadata) {
     Result->materializeMetadata();
     UpgradeDebugInfo(*Result);
   }
 
   return Result;
 }
 
+namespace {
+
+/// Helper to load on demand a Module from file and cache it for subsequent
+/// queries. It can be used with the FunctionImporter.
+class ModuleLazyLoaderCache {
+  /// Cache of lazily loaded module for import.
+  StringMap<std::unique_ptr<Module>> ModuleMap;
+
+  /// Retrieve a Module from the cache or lazily load it on demand.
+  std::function<std::unique_ptr<Module>(const char *argv0,
+                                        const std::string &FileName)>
+      createLazyModule;
+
+public:
+  /// Create the loader, Module will be initialized in \p Context.
+  ModuleLazyLoaderCache(std::function<std::unique_ptr<Module>(
+                            const char *argv0, const std::string &FileName)>
+                            createLazyModule)
+      : createLazyModule(createLazyModule) {}
+
+  /// Retrieve a Module from the cache or lazily load it on demand.
+  Module &operator()(const char *argv0, const std::string &FileName);
+
+  std::unique_ptr<Module> takeModule(const std::string &FileName) {
+    auto I = ModuleMap.find(FileName);
+    assert(I != ModuleMap.end());
+    std::unique_ptr<Module> Ret = std::move(I->second);
+    ModuleMap.erase(I);
+    return Ret;
+  }
+};
+
+// Get a Module for \p FileName from the cache, or load it lazily.
+Module &ModuleLazyLoaderCache::operator()(const char *argv0,
+                                          const std::string &Identifier) {
+  auto &Module = ModuleMap[Identifier];
+  if (!Module)
+    Module = createLazyModule(argv0, Identifier);
+  return *Module;
+}
+} // anonymous namespace
+
 static void diagnosticHandler(const DiagnosticInfo &DI) {
   unsigned Severity = DI.getSeverity();
   switch (Severity) {
   case DS_Error:
     errs() << "ERROR: ";
     break;
   case DS_Warning:
     if (SuppressWarnings)

 
 static void diagnosticHandlerWithContext(const DiagnosticInfo &DI, void *C) {
   diagnosticHandler(DI);
 }
 
 /// Import any functions requested via the -import option.
 static bool importFunctions(const char *argv0, LLVMContext &Context,
                             Linker &L) {
-  StringMap<std::unique_ptr<DenseMap<unsigned, MDNode *>>>
-      ModuleToTempMDValsMap;
+  if (SummaryIndex.empty())
+    return true;
+  ErrorOr<std::unique_ptr<ModuleSummaryIndex>> IndexOrErr =
+      llvm::getModuleSummaryIndexForFile(SummaryIndex, diagnosticHandler);
+  std::error_code EC = IndexOrErr.getError();
+  if (EC) {
+    errs() << EC.message() << '\n';
+    return false;
+  }
+  auto Index = std::move(IndexOrErr.get());
+
+  // Map of Module -> List of globals to import from the Module
+  std::map<StringRef, DenseSet<const GlobalValue *>> ModuleToGlobalsToImportMap;
+  auto ModuleLoader = [&Context](const char *argv0,
+                                 const std::string &Identifier) {
+    return loadFile(argv0, Identifier, Context, false);
+  };
+  ModuleLazyLoaderCache ModuleLoaderCache(ModuleLoader);
   for (const auto &Import : Imports) {
     // Identify the requested function and its bitcode source file.
     size_t Idx = Import.find(':');
     if (Idx == std::string::npos) {
       errs() << "Import parameter bad format: " << Import << "\n";
       return false;
     }
     std::string FunctionName = Import.substr(0, Idx);
     std::string FileName = Import.substr(Idx + 1, std::string::npos);
 
     // Load the specified source module.
-    std::unique_ptr<Module> M = loadFile(argv0, FileName, Context, false);
-    if (!M.get()) {
-      errs() << argv0 << ": error loading file '" << FileName << "'\n";
-      return false;
-    }
+    auto &SrcModule = ModuleLoaderCache(argv0, FileName);
 
-    if (verifyModule(*M, &errs())) {
+    if (verifyModule(SrcModule, &errs())) {
       errs() << argv0 << ": " << FileName
              << ": error: input module is broken!\n";
       return false;
     }
 
-    Function *F = M->getFunction(FunctionName);
+    Function *F = SrcModule.getFunction(FunctionName);
     if (!F) {
       errs() << "Ignoring import request for non-existent function "
              << FunctionName << " from " << FileName << "\n";
       continue;
     }
     // We cannot import weak_any functions without possibly affecting the
     // order they are seen and selected by the linker, changing program
     // semantics.
     if (F->hasWeakAnyLinkage()) {
       errs() << "Ignoring import request for weak-any function " << FunctionName
              << " from " << FileName << "\n";
       continue;
     }
 
     if (Verbose)
       errs() << "Importing " << FunctionName << " from " << FileName << "\n";
 
-    // Link in the specified function.
-    DenseSet<const GlobalValue *> GlobalsToImport;
-    GlobalsToImport.insert(F);
+    auto &Entry = ModuleToGlobalsToImportMap[SrcModule.getModuleIdentifier()];
+    Entry.insert(F);
 
-    if (!SummaryIndex.empty()) {
-      ErrorOr<std::unique_ptr<ModuleSummaryIndex>> IndexOrErr =
-          llvm::getModuleSummaryIndexForFile(SummaryIndex, diagnosticHandler);
-      std::error_code EC = IndexOrErr.getError();
-      if (EC) {
-        errs() << EC.message() << '\n';
-        return false;
-      }
-      auto Index = std::move(IndexOrErr.get());
+    F->materialize();
+
+  // Do the actual import of globals now, one Module at a time
+  for (auto &GlobalsToImportPerModule : ModuleToGlobalsToImportMap) {
+    // Get the module for the import
+    auto &GlobalsToImport = GlobalsToImportPerModule.second;
+    std::unique_ptr<Module> SrcModule =
+        ModuleLoaderCache.takeModule(GlobalsToImportPerModule.first);
+    assert(&Context == &SrcModule->getContext() && "Context mismatch");
+
+    // If modules were created with lazy metadata loading, materialize it
+    // now, before linking it (otherwise this will be a noop).
+    SrcModule->materializeMetadata();
+    UpgradeDebugInfo(*SrcModule);
 
-      // Linkage Promotion and renaming
-      if (renameModuleForThinLTO(*M, *Index, &GlobalsToImport))
-        return true;
-    }
-
-    // Save the mapping of value ids to temporary metadata created when
-    // importing this function. If we have already imported from this module,
-    // add new temporary metadata to the existing mapping.
-    auto &TempMDVals = ModuleToTempMDValsMap[FileName];
-    if (!TempMDVals)
-      TempMDVals = llvm::make_unique<DenseMap<unsigned, MDNode *>>();
-
-    if (L.linkInModule(std::move(M), Linker::Flags::None, &GlobalsToImport,
-                       TempMDVals.get()))
-      return false;
-  }
-
-  // Now link in metadata for all modules from which we imported functions.
-  for (StringMapEntry<std::unique_ptr<DenseMap<unsigned, MDNode *>>> &SME :
-       ModuleToTempMDValsMap) {
-    // Load the specified source module.
-    std::unique_ptr<Module> M = loadFile(argv0, SME.getKey(), Context, true);
-    if (!M.get()) {
-      errs() << argv0 << ": error loading file '" << SME.getKey() << "'\n";
-      return false;
-    }
+    if (renameModuleForThinLTO(*SrcModule, *Index, &GlobalsToImport))
 
-    if (verifyModule(*M, &errs())) {
-      errs() << argv0 << ": " << SME.getKey()
-             << ": error: input module is broken!\n";
+    if (L.linkInModule(std::move(SrcModule), Linker::Flags::None,
+                       &GlobalsToImport))
       return false;
-    }
 
-    // Link in all necessary metadata from this module.
-    if (L.linkInMetadata(std::move(M), SME.getValue().get()))
-      return false;
-  }
   return true;
 }
 
 static bool linkFiles(const char *argv0, LLVMContext &Context, Linker &L,
                       const cl::list<std::string> &Files,
                       unsigned Flags) {
   // Filter out flags that don't apply to the first file we load.
   unsigned ApplicableFlags = Flags & Linker::Flags::OverrideFromSrc;