Change ModuleLinker to take a set of GlobalValues to import instead of a single one

[mehdi_amini]

Summary

For efficiency reason, when importing multiple functions for the same Module,
we can avoid reparsing it every time.

[tejohnson]

lib/Transforms/IPO/FunctionImport.cpp

85

I don't see this map being populated anywhere?

153

Needs merge - I have already implemented this functionality in head.

161

This fundamentally changes the pass from doing an iterative worklist approach to a single pass of import per module, and makes it more difficult to go to a priority queue type approach. I don't think it will merge with the changes I made from head for adding newly imported external calls to the worklist actually.

We may not decide it is profitable to import every external call into a particular module at once. E.g. we may see a cold-ish call into an external module, and later after another import expose a hotter call that bumps its priority. I'm fine with passing in a list of functions to import in one go though - there may be multiple calls into that module that have similar priority, and they can be batch imported.

Perhaps add in the mechanics for doing a list of imports, but leave out the changes to this function for now (and just continue to pass in a single function) until we figure out how that should work.

[mehdi_amini]

Rebased. There is only the API change now, we can't yet take advantage of this.

[tejohnson]

LGTM, thanks

[• rafael]

LGTM

[mehdi_amini]

r254486

include/llvm/Linker/Linker.h

public:
   /// Passing OverrideSymbols as true will have symbols from Src
   /// shadow those in the Dest.
   /// For ThinLTO function importing/exporting the \p FunctionInfoIndex
   /// is passed. If a \p FuncToImport is provided, only that single
   /// function is imported from the source module.
   /// Returns true on error.
   bool linkInModule(Module *Src, unsigned Flags = Flags::None,
                     const FunctionInfoIndex *Index = nullptr,
-                    Function *FuncToImport = nullptr);
+                    DenseSet<const GlobalValue *> *FuncToImport = nullptr);
 
   /// \brief Set the composite to the passed-in module.
   void setModule(Module *Dst);
 
   static bool LinkModules(Module *Dest, Module *Src,
                           DiagnosticHandlerFunction DiagnosticHandler,
                           unsigned Flags = Flags::None);
 

lib/Linker/LinkModules.cpp

class ModuleLinker {
   unsigned Flags;
 
   /// Function index passed into ModuleLinker for using in function
   /// importing/exporting handling.
   const FunctionInfoIndex *ImportIndex;
 
   /// Function to import from source module, all other functions are
   /// imported as declarations instead of definitions.
-  Function *ImportFunction;
+  DenseSet<const GlobalValue *> *ImportFunction;
 
   /// Set to true if the given FunctionInfoIndex contains any functions
   /// from this source module, in which case we must conservatively assume
   /// that any of its functions may be imported into another module
   /// as part of a different backend compilation process.
   bool HasExportedFunctions;
 
   /// Set to true when all global value body linking is complete (including
   /// lazy linking). Used to prevent metadata linking from creating new
   /// references.
   bool DoneLinkingBodies;
 
 public:
   ModuleLinker(Module *dstM, Linker::IdentifiedStructTypeSet &Set, Module *srcM,
                DiagnosticHandlerFunction DiagnosticHandler, unsigned Flags,
                const FunctionInfoIndex *Index = nullptr,
-               Function *FuncToImport = nullptr)
+               DenseSet<const GlobalValue *> *FuncToImport = nullptr)
       : DstM(dstM), SrcM(srcM), TypeMap(Set),
         ValMaterializer(TypeMap, DstM, LazilyLinkGlobalValues, this),
         DiagnosticHandler(DiagnosticHandler), Flags(Flags), ImportIndex(Index),
         ImportFunction(FuncToImport), HasExportedFunctions(false),
         DoneLinkingBodies(false) {
     assert((ImportIndex || !ImportFunction) &&
            "Expect a FunctionInfoIndex when importing");
     // If we have a FunctionInfoIndex but no function to import,
bool ModuleLinker::doImportAsDefinition(const GlobalValue *SGV) {
   // global variables with external linkage are transformed to
   // available_externally definitions, which are ultimately turned into
   // declarations after the EliminateAvailableExternally pass).
   if (isa<GlobalVariable>(SGV) && !SGV->isDeclaration() &&
       !SGV->hasWeakAnyLinkage())
     return true;
   // Only import the function requested for importing.
   auto *SF = dyn_cast<Function>(SGV);
-  if (SF && SF == ImportFunction)
+  if (SF && ImportFunction->count(SF))
     return true;
   // Otherwise no.
   return false;
 }
 
 bool ModuleLinker::doPromoteLocalToGlobal(const GlobalValue *SGV) {
   assert(SGV->hasLocalLinkage());
   // Both the imported references and the original local variable must
bool ModuleLinker::shouldLinkFromSource(bool &LinkFromSrc,
 
   bool SrcIsDeclaration = Src.isDeclarationForLinker();
   bool DestIsDeclaration = Dest.isDeclarationForLinker();
 
   if (isPerformingImport()) {
     if (isa<Function>(&Src)) {
       // For functions, LinkFromSrc iff this is the function requested
       // for importing. For variables, decide below normally.
-      LinkFromSrc = (&Src == ImportFunction);
+      LinkFromSrc = ImportFunction->count(&Src);
       return false;
     }
 
     // Check if this is an alias with an already existing definition
     // in Dest, which must have come from a prior importing pass from
     // the same Src module. Unlike imported function and variable
     // definitions, which are imported as available_externally and are
     // not definitions for the linker, that is not a valid linkage for
bool ModuleLinker::linkGlobalValueProto(GlobalValue *SGV) {
   GlobalValue *NewGV;
   if (!LinkFromSrc) {
     NewGV = DGV;
     // When linking from source we setVisibility from copyGlobalValueProto.
     setVisibility(NewGV, SGV, DGV);
   } else {
     // If the GV is to be lazily linked, don't create it just yet.
     // The ValueMaterializerTy will deal with creating it if it's used.
-    if (!DGV && !shouldOverrideFromSrc() && SGV != ImportFunction &&
+    if (!DGV && !shouldOverrideFromSrc() && ImportFunction->count(SGV) &&
         (SGV->hasLocalLinkage() || SGV->hasLinkOnceLinkage() ||
          SGV->hasAvailableExternallyLinkage())) {
       DoNotLinkFromSource.insert(SGV);
       return false;
     }
 
     // When we only want to link in unresolved dependencies, blacklist
     // the symbol unless unless DestM has a matching declaration (DGV).

 
 void Linker::deleteModule() {
   delete Composite;
   Composite = nullptr;
 }
 
 bool Linker::linkInModule(Module *Src, unsigned Flags,
                           const FunctionInfoIndex *Index,
-                          Function *FuncToImport) {
+                          DenseSet<const GlobalValue *> *FuncToImport) {
   ModuleLinker TheLinker(Composite, IdentifiedStructTypes, Src,
                          DiagnosticHandler, Flags, Index, FuncToImport);
   bool RetCode = TheLinker.run();
   Composite->dropTriviallyDeadConstantArrays();
   return RetCode;
 }
 
 void Linker::setModule(Module *Dst) {

lib/Transforms/IPO/FunctionImport.cpp

bool FunctionImporter::importFunctions(Module &M) {
   /// module.
   StringSet<> CalledFunctions;
   for (auto &F : M) {
     if (F.isDeclaration() || F.hasFnAttribute(Attribute::OptimizeNone))
       continue;
     for (auto &BB : F) {
       for (auto &I : BB) {
         if (isa<CallInst>(I)) {
-          DEBUG(dbgs() << "Found a call: '" << I << "'\n");
           auto CalledFunction = cast<CallInst>(I).getCalledFunction();
           if (CalledFunction && CalledFunction->hasName() &&
               CalledFunction->isDeclaration())
             CalledFunctions.insert(CalledFunction->getName());
         }
       }
     }
   }
 
-  /// Second step: for every call to an external function, try to import it.
-
-  // Linker that will be used for importing function
-  Linker L(&M, DiagnosticHandler);
+  /// Second step: for every call to an external function, add it to the import
+  /// list.
+  DenseMap<Module *, DenseSet<const GlobalValue *>>
+      ModuleToFunctionsToImportMap;
 
   /// Insert initial called function set in a worklist, so that we can add
   /// transively called functions when importing.
   SmallVector<StringRef, 64> Worklist;
   for (auto &CalledFunction : CalledFunctions)
     Worklist.push_back(CalledFunction.first());
 
   while (!Worklist.empty()) {
while (!Worklist.empty()) {
     // order they are seen and selected by the linker, changing program
     // semantics.
     if (F->hasWeakAnyLinkage()) {
       DEBUG(dbgs() << "Ignoring import request for weak-any function "
                    << CalledFunctionName << " from " << FileName << "\n");
       continue;
     }
 
-    // Link in the specified function.
-    if (L.linkInModule(&Module, Linker::Flags::None, &Index, F))
-      report_fatal_error("Function Import: link error");
+    // TODO: add callees from the newly imported function to the worklist.
+  }
 
-    // TODO: Process the newly imported function and add callees to the
-    // worklist.
+  // Link in the specified functions.
 
-    Changed = true;
+  // Linker that will be used for importing function
+  Linker L(&M, DiagnosticHandler);
+
+  for (auto It : ModuleToFunctionsToImportMap) {
+    auto *Module = It.first;
+    auto &FunctionsToImport = It.second;
+    if (L.linkInModule(Module, Linker::Flags::None, &Index, &FunctionsToImport))
+      report_fatal_error("Function Import: link error");
   }
+
   return Changed;
 }
 
 /// Summary file to use for function importing when using -function-import from
 /// the command line.
 static cl::opt<std::string>
     SummaryFile("summary-file",
                 cl::desc("The summary file to use for function importing."));

tools/llvm-link/llvm-link.cpp

if (!FunctionIndex.empty()) {
       if (EC) {
         errs() << EC.message() << '\n';
         return false;
       }
       Index = std::move(IndexOrErr.get());
     }
 
     // Link in the specified function.
-    if (L.linkInModule(M.get(), Linker::Flags::None, Index.get(), F))
+    DenseSet<const GlobalValue *> FunctionToImport;
+    FunctionToImport.insert(F);
+    if (L.linkInModule(M.get(), Linker::Flags::None, Index.get(),
+                       &FunctionToImport))
       return false;
   }
   return true;
 }
 
 static bool linkFiles(const char *argv0, LLVMContext &Context, Linker &L,
                       const cl::list<std::string> &Files,
                       unsigned Flags) {