FunctionImporter: implement bulk function importing for efficiency

[mehdi_amini]

Summary

The current importing scheme is processing one function at a time,
loading the source Module, linking the function in the destination
module, and destroying the source Module before repeating with the
next function to import (potentially from the same Module).

Ideally we would keep the source Module alive and import the next
Function needed from this Module. Unfortunately this is not possible
because the linker does not leave it in a usable state.

However we can do better by first computing the list of all candidates
per Module, and only then load the source Module and import all the
function we need for it.

The trick to process callees is to materialize function in the source
module when building the list of function to import, and inspect them
in their source module, collecting the list of callees for each
callee.

When we move the the actual import, we will import from each source
module exactly once. Each source module is loaded exactly once.
The only drawback it that it requires to have all the lazy-loaded
source Module in memory at the same time.

Currently this patch already improves considerably the link time,
a multithreaded link of llvm-dis on my laptop was:

real  1m12.175s  user  6m32.430s sys  0m10.529s

and is now:

real  0m40.697s  user  2m10.237s sys  0m4.375s

Note: this is the full link time (linker+Import+Optimizer+CodeGen)

[mehdi_amini]

Supersede D15176

Note: this is still a wip (there are a few functions that are no longer imported, I need to figure out why) but the bulk of the changes should be there.

[tejohnson]

With this approach we consider a function for importing exactly once, and I'm concerned that this will lose opportunities where we might decide to import it after other imports are done (e.g. exposing more hot calls to it). Perhaps when we decide not to import it in the worklist iteration in GetImportList it should also be removed from the CalledFunctions set so that it can be reconsidered when new calls to it are seen.

What was the memory overhead difference with and without this patch? I guess having the materialized functions on all the source modules before importing is ok since they will all eventually be moved into the dest module during linking anyway. But there will still be overhead of keeping all the source modules which have some other overhead.

lib/Transforms/IPO/FunctionImport.cpp

248

How do you ensure that F is materialized in the lazy-loaded SrcModule? With the current code at HEAD it is the dest copy we walk, which was materialized in the source and then copied over during module linking.

Another issue with using the source copy here - the local functions haven't been renamed/promoted, so there may be apparent duplicates in CalledFunctions due to same-named local functions from other modules.

[mehdi_amini]

Fix materialization and renaming for internal functions

[mehdi_amini]

Use an ordered std::map for reproducibility

[tejohnson]

Thanks for working on this. Some more comments/questions below.

lib/Transforms/IPO/FunctionImport.cpp

87

Remove

101

Maybe "...internal function calls imported from..."? I.e. at this point we have only imported the call.

129

We reach here if the function is already in the CalledFunctions set, so do we want to issue a message that we are ignoring it? It seems misleading, i.e. it sounds like we decided not to import.

141

Why is this under #ifndef NDEBUG?

281–310

I don't see anything called ProcessImportWorklist? Do you mean GetImportList? Each iteration of what?

285

Comment reads funny: "helps required by"

[mehdi_amini]

Update following Teresa's comments

[mehdi_amini]

Rebase on trunk

[mehdi_amini]

Remove spurious NDEBUG

[tejohnson]

Looks like it needs another rebase given your most recent changes, but otherwise lgtm. Thanks!

[mehdi_amini]

r255100

[tejohnson]

One small post-commit suggestion for a TODO.

lib/Transforms/IPO/FunctionImport.cpp

130

We will never reconsider a function after initially deciding it isn't profitable to import. Currently the only profitability decision is based on the number of instructions, which won't change with importing. But perhaps we should add a TODO down where we have those checks to remove any callees from the CalledFunctions list if they fail a profitability check that may be affected by additional importing decisions.

include/llvm/Transforms/IPO/FunctionImport.h

 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace llvm {
 class LLVMContext;
 class Module;
 class FunctionInfoIndex;
 
-/// Helper to load on demand a Module from file and cache it for subsequent
-/// queries. It can be used with the FunctionImporter.
-class ModuleLazyLoaderCache {
-  /// The context that will be used for importing.
-  LLVMContext &Context;
-
-  /// Cache of lazily loaded module for import.
-  StringMap<std::unique_ptr<Module>> ModuleMap;
-
-public:
-  /// Create the loader, Module will be initialized in \p Context.
-  ModuleLazyLoaderCache(LLVMContext &Context) : Context(Context) {}
-
-  /// Retrieve a Module from the cache or lazily load it on demand.
-  Module &operator()(StringRef FileName);
-};
-
 /// The function importer is automatically importing function from other modules
 /// based on the provided summary informations.
 class FunctionImporter {
 
   /// The summaries index used to trigger importing.
   const FunctionInfoIndex &Index;
 
   /// Diagnostic will be sent to this handler.
   DiagnosticHandlerFunction DiagnosticHandler;
 
-  /// Retrieve a Module from the cache or lazily load it on demand.
-  std::function<Module &(StringRef FileName)> getLazyModule;
+  /// Factory function to load a Module for a given identifier
+  std::function<std::unique_ptr<Module>(StringRef Identifier)> ModuleLoader;
 
 public:
+
   /// Create a Function Importer.
   FunctionImporter(const FunctionInfoIndex &Index,
                    DiagnosticHandlerFunction DiagnosticHandler,
-                   std::function<Module &(StringRef FileName)> ModuleLoader)
+                   std::function<std::unique_ptr<Module>(StringRef Identifier)> ModuleLoader)
       : Index(Index), DiagnosticHandler(DiagnosticHandler),
-        getLazyModule(ModuleLoader) {}
+        ModuleLoader(ModuleLoader) {}
 
   /// Import functions in Module \p M based on the summary informations.
   bool importFunctions(Module &M);
 };
 }
 
 #endif // LLVM_FUNCTIONIMPORT_H

lib/Transforms/IPO/FunctionImport.cpp

 //===- FunctionImport.cpp - ThinLTO Summary-based Function Import ---------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements Function import based on summaries.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/FunctionImport.h"
 
 #include "llvm/ADT/StringSet.h"
+#include "llvm/Bitcode/ReaderWriter.h"
 #include "llvm/IR/AutoUpgrade.h"
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Object/FunctionIndexObjectFile.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/SourceMgr.h"
+
+#include <map>
+
 using namespace llvm;
 
 #define DEBUG_TYPE "function-import"
 
 /// Limit on instruction count of imported functions.
 static cl::opt<unsigned> ImportInstrLimit(
     "import-instr-limit", cl::init(100), cl::Hidden, cl::value_desc("N"),
     cl::desc("Only import functions with less than N instructions"));
static std::unique_ptr<Module> loadFile(const std::string &FileName,
   }
 
   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>(StringRef FileName)> createLazyModule;
+
+public:
+  /// Create the loader, Module will be initialized in \p Context.
+  ModuleLazyLoaderCache(std::function<
+      std::unique_ptr<Module>(StringRef FileName)> createLazyModule)
+      : createLazyModule(createLazyModule) {}
+
+  /// Retrieve a Module from the cache or lazily load it on demand.
+  Module &operator()(StringRef FileName);
+};
+
 // Get a Module for \p FileName from the cache, or load it lazily.
-Module &ModuleLazyLoaderCache::operator()(StringRef FileName) {
-  auto &Module = ModuleMap[FileName];
+Module &ModuleLazyLoaderCache::operator()(StringRef Identifier) {
+  auto &Module = ModuleMap[Identifier];
   if (!Module)
-    Module = loadFile(FileName, Context);
+    Module = createLazyModule(Identifier);
   return *Module;
 }
+}
+
+//// Get a Module for \p FileName from the cache, or load it lazily.
+// Module &ModuleLazyLoaderCache::operator()(StringRef FileName) {
+//  auto &Module = ModuleMap[FileName];
+//  if (!Module)
+//    Module = loadFile(FileName, Context);
+//  return *Module;
+//}
 
 /// Walk through the instructions in \p F looking for external
 /// calls not already in the \p CalledFunctions set. If any are
 /// found they are added to the \p Worklist for importing.
-static void findExternalCalls(const Function &F, StringSet<> &CalledFunctions,
+static void findExternalCalls(const Module &DestModule, Function &F,
+                              const FunctionInfoIndex &Index,
+                              StringSet<> &CalledFunctions,
                               SmallVector<StringRef, 64> &Worklist) {
+  // We need to suffix internal function imported from other modules, prepare
+  // the suffix ahead of time.
+  StringRef Suffix;
+  if (F.getParent() != &DestModule)
+    Suffix =
+        (Twine(".llvm.") +
+         Twine(Index.getModuleId(F.getParent()->getModuleIdentifier()))).str();
+
   for (auto &BB : F) {
     for (auto &I : BB) {
       if (isa<CallInst>(I)) {
         auto CalledFunction = cast<CallInst>(I).getCalledFunction();
         // Insert any new external calls that have not already been
         // added to set/worklist.
-        if (CalledFunction && CalledFunction->hasName() &&
-            CalledFunction->isDeclaration() &&
-            !CalledFunctions.count(CalledFunction->getName())) {
-          CalledFunctions.insert(CalledFunction->getName());
-          Worklist.push_back(CalledFunction->getName());
+        if (!CalledFunction || !CalledFunction->hasName())
+          continue;
+        // Ignore intrinsics early
+        if (CalledFunction->isIntrinsic()) {
+          assert(CalledFunction->getIntrinsicID() != 0);
+          continue;
+        }
+        auto ImportedName = CalledFunction->getName();
+        // Rename internal functions
+        if (CalledFunction->hasInternalLinkage()) {
+          ImportedName = (ImportedName + Suffix).str();
+        }
+        auto It = CalledFunctions.insert(ImportedName);
+        if (!It.second) {
+          DEBUG(dbgs() << DestModule.getModuleIdentifier() << " Ignoring "
+                       << ImportedName << "\n");
+          continue;
+        }
+#ifndef NDEBUG
+        // Ignore functions already present in the destination module
+        auto *SrcGV = DestModule.getNamedValue(ImportedName);
+        if (SrcGV) {
+          assert(isa<Function>(SrcGV) && "Name collision during import");
+          if (!cast<Function>(SrcGV)->isDeclaration()) {
+            DEBUG(dbgs() << DestModule.getModuleIdentifier() << "Ignoring "
+                         << ImportedName << " already in DestinationModule\n");
+            continue;
+          }
         }
+#endif // NDEBUG
+
+        Worklist.push_back(It.first->getKey());
+        DEBUG(dbgs() << DestModule.getModuleIdentifier()
+                     << " Adding callee for : " << ImportedName << " : "
+                     << F.getName() << "\n");
       }
     }
   }
 }
 
 // Helper function: given a worklist and an index, will process all the worklist
 // and import them based on the summary information
-static unsigned ProcessImportWorklist(
+static void GetImportList(
     Module &DestModule, SmallVector<StringRef, 64> &Worklist,
-    StringSet<> &CalledFunctions, Linker &TheLinker,
-    const FunctionInfoIndex &Index,
-    std::function<Module &(StringRef FileName)> &LazyModuleLoader) {
-  unsigned ImportCount = 0;
+    StringSet<> &CalledFunctions,
+    std::map<StringRef, std::pair<Module *, DenseSet<const GlobalValue *>>> &
+        ModuleToFunctionsToImportMap,
+    const FunctionInfoIndex &Index, ModuleLazyLoaderCache &ModuleLoaderCache) {
   while (!Worklist.empty()) {
     auto CalledFunctionName = Worklist.pop_back_val();
-    DEBUG(dbgs() << "Process import for " << CalledFunctionName << "\n");
+    DEBUG(dbgs() << DestModule.getModuleIdentifier() << "Process import for "
+                 << CalledFunctionName << "\n");
 
     // Try to get a summary for this function call.
     auto InfoList = Index.findFunctionInfoList(CalledFunctionName);
     if (InfoList == Index.end()) {
-      DEBUG(dbgs() << "No summary for " << CalledFunctionName
-                   << " Ignoring.\n");
+      DEBUG(dbgs() << DestModule.getModuleIdentifier() << "No summary for "
+                   << CalledFunctionName << " Ignoring.\n");
       continue;
     }
     assert(!InfoList->second.empty() && "No summary, error at import?");
 
     // Comdat can have multiple entries, FIXME: what do we do with them?
     auto &Info = InfoList->second[0];
     assert(Info && "Nullptr in list, error importing summaries?\n");
 
     auto *Summary = Info->functionSummary();
     if (!Summary) {
       // FIXME: in case we are lazyloading summaries, we can do it now.
-      DEBUG(dbgs() << "Missing summary for  " << CalledFunctionName
+      DEBUG(dbgs() << DestModule.getModuleIdentifier()
+                   << " Missing summary for  " << CalledFunctionName
                    << ", error at import?\n");
       llvm_unreachable("Missing summary");
     }
 
     if (Summary->instCount() > ImportInstrLimit) {
-      DEBUG(dbgs() << "Skip import of " << CalledFunctionName << " with "
-                   << Summary->instCount() << " instructions (limit "
-                   << ImportInstrLimit << ")\n");
+      DEBUG(dbgs() << DestModule.getModuleIdentifier() << " Skip import of "
+                   << CalledFunctionName << " with " << Summary->instCount()
+                   << " instructions (limit " << ImportInstrLimit << ")\n");
       continue;
     }
 
     // Get the module path from the summary.
-    auto FileName = Summary->modulePath();
-    DEBUG(dbgs() << "Importing " << CalledFunctionName << " from " << FileName
-                 << "\n");
-
-    // Get the module for the import (potentially from the cache).
-    auto &Module = LazyModuleLoader(FileName);
-    assert(&Module.getContext() == &DestModule.getContext());
+    auto ModuleIdentifier = Summary->modulePath();
+    DEBUG(dbgs() << DestModule.getModuleIdentifier() << " Importing "
+                 << CalledFunctionName << " from " << ModuleIdentifier << "\n");
+
+    auto &SrcModule = ModuleLoaderCache(ModuleIdentifier);
 
     // The function that we will import!
-    GlobalValue *SGV = Module.getNamedValue(CalledFunctionName);
-    StringRef ImportFunctionName = CalledFunctionName;
+    GlobalValue *SGV = SrcModule.getNamedValue(CalledFunctionName);
+
     if (!SGV) {
-      // Might be local in source Module, promoted/renamed in DestModule.
+      // The destination module is referencing function using their renamed name
+      // when importing a function that was originally local in the source
+      // module. The source module we have might not have been renamed so we try
+      // to remove the suffix added during the renaming to recover the original
+      // name in the source module.
       std::pair<StringRef, StringRef> Split =
           CalledFunctionName.split(".llvm.");
-      SGV = Module.getNamedValue(Split.first);
-#ifndef NDEBUG
-      // Assert that Split.second is module id
-      uint64_t ModuleId;
-      assert(!Split.second.getAsInteger(10, ModuleId));
-      assert(ModuleId == Index.getModuleId(FileName));
-#endif
+      SGV = SrcModule.getNamedValue(Split.first);
+      assert(SGV && "Can't find function to import in source module");
+    }
+    if (!SGV) {
+      report_fatal_error(Twine("Can't load function '") + CalledFunctionName +
+                         "' in Module '" + SrcModule.getModuleIdentifier() +
+                         "', error in the summary?\n");
     }
+
     Function *F = dyn_cast<Function>(SGV);
     if (!F && isa<GlobalAlias>(SGV)) {
       auto *SGA = dyn_cast<GlobalAlias>(SGV);
       F = dyn_cast<Function>(SGA->getBaseObject());
-      ImportFunctionName = F->getName();
-    }
-    if (!F) {
-      errs() << "Can't load function '" << CalledFunctionName << "' in Module '"
-             << FileName << "', error in the summary?\n";
-      llvm_unreachable("Can't load function in Module");
+      CalledFunctionName = F->getName();
     }
+    assert(F && "Imported Function is ... not a Function");
 
     // We cannot import weak_any functions/aliases without possibly affecting
     // the order they are seen and selected by the linker, changing program
     // semantics.
     if (SGV->hasWeakAnyLinkage()) {
-      DEBUG(dbgs() << "Ignoring import request for weak-any "
+      DEBUG(dbgs() << DestModule.getModuleIdentifier()
+                   << " Ignoring import request for weak-any "
                    << (isa<Function>(SGV) ? "function " : "alias ")
-                   << CalledFunctionName << " from " << FileName << "\n");
+                   << CalledFunctionName << " from "
+                   << SrcModule.getModuleIdentifier() << "\n");
       continue;
     }
 
-    // Link in the specified function.
-    DenseSet<const GlobalValue *> FunctionsToImport;
-    FunctionsToImport.insert(F);
-    if (TheLinker.linkInModule(Module, Linker::Flags::None, &Index,
-                               &FunctionsToImport))
-      report_fatal_error("Function Import: link error");
-
-    // Process the newly imported function and add callees to the worklist.
-    GlobalValue *NewGV = DestModule.getNamedValue(ImportFunctionName);
-    assert(NewGV);
-    Function *NewF = dyn_cast<Function>(NewGV);
-    assert(NewF);
-    findExternalCalls(*NewF, CalledFunctions, Worklist);
-    ++ImportCount;
+    // Add the function to the import list
+    auto &Entry = ModuleToFunctionsToImportMap[SrcModule.getModuleIdentifier()];
+    Entry.first = &SrcModule;
+    Entry.second.insert(F);
+
+    // Process the newly imported functions and add callees to the worklist.
+    F->materialize();
+    findExternalCalls(DestModule, *F, Index, CalledFunctions, Worklist);
   }
-  return ImportCount;
 }
 
 // Automatically import functions in Module \p DestModule based on the summaries
 // index.
 //
 // The current implementation imports every called functions that exists in the
 // summaries index.
 bool FunctionImporter::importFunctions(Module &DestModule) {
-  DEBUG(errs() << "Starting import for Module "
+  DEBUG(dbgs() << "Starting import for Module "
                << DestModule.getModuleIdentifier() << "\n");
   unsigned ImportedCount = 0;
 
   /// First step is collecting the called external functions.
   StringSet<> CalledFunctions;
   SmallVector<StringRef, 64> Worklist;
   for (auto &F : DestModule) {
     if (F.isDeclaration() || F.hasFnAttribute(Attribute::OptimizeNone))
       continue;
-    findExternalCalls(F, CalledFunctions, Worklist);
+    findExternalCalls(DestModule, F, Index, CalledFunctions, Worklist);
   }
   if (Worklist.empty())
     return false;
 
   /// Second step: for every call to an external function, try to import it.
 
   // Linker that will be used for importing function
   Linker TheLinker(DestModule, DiagnosticHandler);
 
-  ImportedCount += ProcessImportWorklist(DestModule, Worklist, CalledFunctions,
-                                         TheLinker, Index, getLazyModule);
+  /// For each iteration, ProcessImportWorklist() will empty the Worklist and
+  /// then fill it with the callees of the newly imported functions, that we'll
+  /// process again.
+  /// This workflow is there for efficiency, since batch-importing functions
+  /// in ProcessImportWorklist helps required by the ModuleLinker.
+
+  // Map of Module -> List of Function to import from the Module
+  std::map<StringRef, std::pair<Module *, DenseSet<const GlobalValue *>>>
+      ModuleToFunctionsToImportMap;
+
+  // Analyze the summaries and get the list of functions to import by
+  // populating ModuleToFunctionsToImportMap
+  ModuleLazyLoaderCache ModuleLoaderCache(ModuleLoader);
+  GetImportList(DestModule, Worklist, CalledFunctions,
+                ModuleToFunctionsToImportMap, Index, ModuleLoaderCache);
+  assert(Worklist.empty() && "Worklist hasn't been flushed in GetImportList");
+
+  // Do the actual import of functions now, one Module at a time
+  for (auto &FunctionsToImportPerModule : ModuleToFunctionsToImportMap) {
+    // Get the module for the import
+    // Can't cache it for now because the linker mess it up.
+    auto &FunctionsToImport = FunctionsToImportPerModule.second.second;
+    auto *SrcModule = FunctionsToImportPerModule.second.first;
+    assert(&DestModule.getContext() == &SrcModule->getContext() &&
+           "Context mismatch");
 
-  DEBUG(errs() << "Imported " << ImportedCount << " functions for Module "
+    // Link in the specified functions.
+    if (TheLinker.linkInModule(*SrcModule, Linker::Flags::None, &Index,
+                               &FunctionsToImport))
+      report_fatal_error("Function Import: link error");
+
+    ImportedCount += FunctionsToImport.size();
+  }
+  DEBUG(dbgs() << "Imported " << ImportedCount << " functions for Module "
                << DestModule.getModuleIdentifier() << "\n");
   return ImportedCount;
 }
 
 /// Summary file to use for function importing when using -function-import from
 /// the command line.
 static cl::opt<std::string>
     SummaryFile("summary-file",
bool runOnModule(Module &M) override {
     std::unique_ptr<FunctionInfoIndex> Index =
         getFunctionIndexForFile(SummaryFile, Error, diagnosticHandler);
     if (!Index) {
       errs() << "Error loading file '" << SummaryFile << "': " << Error << "\n";
       return false;
     }
 
     // Perform the import now.
-    ModuleLazyLoaderCache Loader(M.getContext());
-    FunctionImporter Importer(*Index, diagnosticHandler,
-                              [&](StringRef Name)
-                                  -> Module &{ return Loader(Name); });
+    auto ModuleLoader =
+        [&M](StringRef FileName) { return loadFile(FileName, M.getContext()); };
+    FunctionImporter Importer(*Index, diagnosticHandler, ModuleLoader);
     return Importer.importFunctions(M);
 
     return false;
   }
 };
 
 char FunctionImportPass::ID = 0;
 INITIALIZE_PASS_BEGIN(FunctionImportPass, "function-import",
                       "Summary Based Function Import", false, false)
 INITIALIZE_PASS_END(FunctionImportPass, "function-import",
                     "Summary Based Function Import", false, false)
 
 namespace llvm {
 Pass *createFunctionImportPass() { return new FunctionImportPass(); }
 }