[ThinLTO] Do metadata linking during batch function importing

[tejohnson]

Summary

Since we are currently not doing incremental importing there is
no need to link metadata as a postpass. The module linker will
only link in the imported subroutines due to the functionality
added by r256003.

(Note that the metadata postpass linking functionalitiy is still
used by llvm-link, and may be needed here in the future if a more
incremental strategy is adopted.)

[mehdi_amini]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45576)

I'd question why did we load them with lazy metadata in the first place? If we have materialize the metadata anyway, we can as well load them up-front.

[tejohnson]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45576)

I think delaying it until here should save memory - otherwise we have all metadata for all source modules in memory at once. Materializing it here just before we link, after which the SrcModule is destroyed, should mean we only have one source module worth of metadata materialized and in memory at once (other than what is linked into the dest module, which will be a subset of the metadata, even smaller once I stop linking in the other things hanging off the DICompileUnit).

[mehdi_amini]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45576)

I'm fine with leaving the call to materializeMetadata as it will turn to be a no-op if they were loaded up-front.
Can you reflect this in the comment above? Just prepending an "If" at the beginning looks good to me.

[mehdi_amini]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45576)

Actually, I have another idea: why not load metadata on demand when materializing a function?
This seems quite easy to achieve, and would lead us to a state where we would have everything we need materialized, and nothing more.

[tejohnson]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45576)

That has the same disadvantages as not lazy loading in the first place - since we will materialize each function when we decide to import it (in order to look for external calls), we will end up materializing all metadata for all source modules we're importing from and holding them in memory at the same time.

[tejohnson]

• Update comments.

[mehdi_amini]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45592)

Not at the same scale: my point is to materialize only the metatadata referenced by the function being materialized, so it is not the same as materializing *all* the metadata: if we import one function from a module we would be able to materialize only the needed metadata.

[tejohnson]

lib/Transforms/IPO/FunctionImport.cpp

308 ↗

(On Diff #45592)

Oh I see. I thought you were suggesting that I move this call to materializeMetadata into the materialize() call for a function.

What you are suggesting would be interesting, but right now isn't supported as metadata materialization involves parsing all of the module level metadata. Presumably with the approach you are describing it would need to be parsed for each function materialized and somehow only the records of interest would be recorded? The metadata is correlated to references by the index into the MetadataList, which is populated while parsing. You won't even know which metadata values you need until you parse the function, at which point you have already created the forward references. Sorry if I am missing something.

[mehdi_amini]

What I have in mind is to rememberAndSkipMetadata() when parsing the module. But embedded the implementation in BitcodeReaderMetadataList so that whenever we call getValueFwdRef we don't create a placeholder (if the ID is one of the module-level constant) but parse the metadata on-demand. I should be able to craft something quickly.

[tejohnson]

In D16424#332826, @joker.eph wrote:

What I have in mind is to rememberAndSkipMetadata() when parsing the module. But embedded the implementation in BitcodeReaderMetadataList so that whenever we call getValueFwdRef we don't create a placeholder (if the ID is one of the module-level constant) but parse the metadata on-demand. I should be able to craft something quickly.

Note that the metadata value ids are assigned implicitly based on the order of the metadata records. Will you only parse the needed records, skipping the intermediate ones, while counting metadata value ids that would be assigned by skipped MD records that assign values (all but METADATA_NAME)? You'd need to go back and parse the skipped records on backwards references within the metadata. Otherwise, you'd have to parse all MD records up through the ones needed for the imported function (presumably leading to overhead from saving them, unless you can remove the ones that weren't actually needed by the imported function - but then they may need to be reparsed later for a different imported function).

In the meantime, should we put this patch in, since it is a simple change and should improve your situation give the current bulk importing? Most of this is needed regardless of whether you are able to parse just the needed metadata.

[mehdi_amini]

LGTM: strictly better than before.

llvm/trunk/lib/Linker/LinkModules.cpp

ModuleLinker(IRMover &Mover, Module &SrcM, unsigned Flags,
     assert((ImportIndex || !FunctionsToImport) &&
            "Expect a FunctionInfoIndex when importing");
     // If we have a FunctionInfoIndex but no function to import,
     // then this is the primary module being compiled in a ThinLTO
     // backend compilation, and we need to see if it has functions that
     // may be exported to another backend compilation.
     if (ImportIndex && !FunctionsToImport)
       HasExportedFunctions = ImportIndex->hasExportedFunctions(SrcM);
-    assert((ValIDToTempMDMap || !FunctionsToImport) &&
-           "Function importing must provide a ValIDToTempMDMap");
   }
 
   bool run();
 };
 
 /// Class to handle necessary GlobalValue changes required by ThinLTO including
 /// linkage changes and any necessary renaming.
 class ThinLTOGlobalProcessing {

llvm/trunk/lib/Transforms/IPO/FunctionImport.cpp

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"));
 
 // Load lazily a module from \p FileName in \p Context.
 static std::unique_ptr<Module> loadFile(const std::string &FileName,
                                         LLVMContext &Context) {
   SMDiagnostic Err;
   DEBUG(dbgs() << "Loading '" << FileName << "'\n");
-  // Metadata isn't loaded or linked until after all functions are
-  // imported, after which it will be materialized and linked.
+  // Metadata isn't loaded until functions are imported, to minimize
+  // the memory overhead.
   std::unique_ptr<Module> Result =
       getLazyIRFileModule(FileName, Err, Context,
                           /* ShouldLazyLoadMetadata = */ true);
   if (!Result) {
     Err.print("function-import", errs());
     return nullptr;
   }
 
bool FunctionImporter::importFunctions(Module &DestModule) {
 
   // 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");
 
-  StringMap<std::unique_ptr<DenseMap<unsigned, MDNode *>>>
-      ModuleToTempMDValsMap;
-
   // Do the actual import of functions now, one Module at a time
   for (auto &FunctionsToImportPerModule : ModuleToFunctionsToImportMap) {
     // Get the module for the import
     auto &FunctionsToImport = FunctionsToImportPerModule.second;
     std::unique_ptr<Module> SrcModule =
         ModuleLoaderCache.takeModule(FunctionsToImportPerModule.first);
     assert(&DestModule.getContext() == &SrcModule->getContext() &&
            "Context mismatch");
 
-    // 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[SrcModule->getModuleIdentifier()];
-    if (!TempMDVals)
-      TempMDVals = llvm::make_unique<DenseMap<unsigned, MDNode *>>();
+    // If modules were created with lazy metadata loading, materialize it
+    // now, before linking it (otherwise this will be a noop).
+    SrcModule->materializeMetadata();
+    UpgradeDebugInfo(*SrcModule);
 
     // Link in the specified functions.
     if (TheLinker.linkInModule(std::move(SrcModule), Linker::Flags::None,
-                               &Index, &FunctionsToImport, TempMDVals.get()))
+                               &Index, &FunctionsToImport))
       report_fatal_error("Function Import: link error");
 
     ImportedCount += FunctionsToImport.size();
   }
 
-  // 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.
-    auto &SrcModule = ModuleLoaderCache(SME.getKey());
-    // The modules were created with lazy metadata loading. Materialize it
-    // now, before linking it.
-    SrcModule.materializeMetadata();
-    UpgradeDebugInfo(SrcModule);
-
-    // Link in all necessary metadata from this module.
-    if (TheLinker.linkInMetadata(SrcModule, SME.getValue().get()))
-      return false;
-  }
-
   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>