[ThinLTO] Use bulk importing in llvm-link

Summary:
Use bulk importing so we can avoid the use of post-pass metadata
linking. Cloned the ModuleLazyLoaderCache from the FunctionImport pass
to facilitate this.

Reviewers: joker.eph

Subscribers: dexonsmith, llvm-commits, joker.eph

Differential Revision: http://reviews.llvm.org/D18455

llvm-svn: 264326

1 files changed, 89 insertions, 56 deletions

diff --git a/llvm/tools/llvm-link/llvm-link.cpp b/llvm/tools/llvm-link/llvm-link.cpp
index b70ec87f889..729f8fd6449 100644
--- a/llvm/tools/llvm-link/llvm-link.cpp
+++ b/llvm/tools/llvm-link/llvm-link.cpp
@@ -111,8 +111,10 @@ static std::unique_ptr<Module> loadFile(const char *argv0,
   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();
@@ -122,6 +124,48 @@ static std::unique_ptr<Module> loadFile(const char *argv0,
   return Result;
 }
 
+namespace {
+
+/// Helper to load on demand a Module from file and cache it for subsequent
+/// queries during function importing.
+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) {
@@ -150,8 +194,24 @@ static void diagnosticHandlerWithContext(const DiagnosticInfo &DI, void *C) {
 /// 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(':');
@@ -163,19 +223,15 @@ static bool importFunctions(const char *argv0, LLVMContext &Context,
     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";
@@ -193,57 +249,34 @@ static bool importFunctions(const char *argv0, LLVMContext &Context,
     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());
-
-      // 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;
+    F->materialize();
   }
 
-  // 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 (verifyModule(*M, &errs())) {
-      errs() << argv0 << ": " << SME.getKey()
-             << ": error: input module is broken!\n";
-      return false;
-    }
-
-    // Link in all necessary metadata from this module.
-    if (L.linkInMetadata(std::move(M), SME.getValue().get()))
+  // 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(*SrcModule, *Index, &GlobalsToImport))
+      return true;
+
+    if (L.linkInModule(std::move(SrcModule), Linker::Flags::None,
+                       &GlobalsToImport))
       return false;
   }
+
   return true;
 }