Modify FunctionImport to take a callback to load modules

[mehdi_amini]

Summary

When linking static archive, there is no individual module files to
load. Instead they can be mmap'ed and could be initialized from a
buffer directly. The callback provide flexibility to override the
scheme for loading module from the summary.

[tejohnson]

loading module from the summary.

Change LGTM, I just didn't understand the above note about loading the module from the summary? Do you just mean from memory (not the function summary)?

[mehdi_amini]

When I wrote "when loading module from the summary" I meant "when loading module from the module identifier in the summary (which may not be a file despite the accessor being called modulePath())". Makes sense?

[mehdi_amini]

r254479

include/llvm/Transforms/IPO/FunctionImport.h

 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace llvm {
 class LLVMContext;
 class Module;
 class FunctionInfoIndex;
 
-/// The function importer is automatically importing function from other modules
-/// based on the provided summary informations.
-class FunctionImporter {
+/// 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;
 
-  /// The context that will be used for importing.
-  LLVMContext &Context;
+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.
-  Module &getOrLoadModule(StringRef FileName);
+  std::function<Module &(StringRef FileName)> getLazyModule;
 
 public:
   /// Create a Function Importer.
-  FunctionImporter(LLVMContext &Context, const FunctionInfoIndex &Index,
-                   DiagnosticHandlerFunction DiagnosticHandler)
-      : Context(Context), Index(Index), DiagnosticHandler(DiagnosticHandler) {}
+  FunctionImporter(const FunctionInfoIndex &Index,
+                   DiagnosticHandlerFunction DiagnosticHandler,
+                   std::function<Module &(StringRef FileName)> ModuleLoader)
+      : Index(Index), DiagnosticHandler(DiagnosticHandler),
+        getLazyModule(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

static std::unique_ptr<Module> loadFile(const std::string &FileName,
 
   Result->materializeMetadata();
   UpgradeDebugInfo(*Result);
 
   return Result;
 }
 
 // Get a Module for \p FileName from the cache, or load it lazily.
-Module &FunctionImporter::getOrLoadModule(StringRef FileName) {
+Module &ModuleLazyLoaderCache::operator()(StringRef FileName) {
   auto &Module = ModuleMap[FileName];
   if (!Module)
     Module = loadFile(FileName, Context);
   return *Module;
 }
 
 // Automatically import functions in Module \p M based on the summaries index.
 //
 // The current implementation imports every called functions that exists in the
 // summaries index.
 bool FunctionImporter::importFunctions(Module &M) {
-  assert(&Context == &M.getContext());
 
   bool Changed = false;
 
   /// First step is collecting the called functions and the one defined in this
   /// module.
   StringSet<> CalledFunctions;
   for (auto &F : M) {
     if (F.isDeclaration() || F.hasFnAttribute(Attribute::OptimizeNone))
while (!Worklist.empty()) {
     //
 
     // 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 = getOrLoadModule(FileName);
+    auto &Module = getLazyModule(FileName);
+    assert(&Module.getContext() == &M.getContext());
 
     // The function that we will import!
     GlobalValue *SGV = Module.getNamedValue(CalledFunctionName);
     Function *F = dyn_cast<Function>(SGV);
     if (!F && isa<GlobalAlias>(SGV)) {
       auto *SGA = dyn_cast<GlobalAlias>(SGV);
       F = dyn_cast<Function>(SGA->getBaseObject());
     }
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.
-    FunctionImporter Importer(M.getContext(), *Index, diagnosticHandler);
+    ModuleLazyLoaderCache Loader(M.getContext());
+    FunctionImporter Importer(*Index, diagnosticHandler,
+                              [&](StringRef Name)
+                                  -> Module &{ return Loader(Name); });
     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(); }
 }