Use lazy-loading of Metadata in MetadataLoader when importing is enabled (NFC)

[mehdi_amini]

Summary

This is a relatively simple scheme: we use the index emitted in the
bitcode to avoid loading all the global metadata. Instead we load
the index with their position in the bitcode so that we can load each
of them individually. Materializing the global metadata block in this
condition only triggers loading the named metadata, and the ones
referenced from there (transitively). When materializing a function,
metadatas from the global block are loaded lazily as they are
referenced.

Two main current limitations are:

1. Global values other than functions are not materialized on demand,

so we need to eagerly load METADATA_GLOBAL_DECL_ATTACHMENT records
(and their transitive dependencies).

2. When we load a single metadata, we don't recurse on the operands,

instead we use a placeholder or a temporary metadata. Unfortunately
tepmorary nodes are very expensive. This is why we don't have it
always enabled and only for importing.

These two limitations can be lifted in a subsequent improvement if
needed.

With this change, the total link time of opt with ThinLTO and Debug
Info enabled is going down from 282s to 224s (~20%).

[mehdi_amini]

Improve comments here and there + very minor improvements

[mehdi_amini]

More cleanup and comments!

[mehdi_amini]

Fix one last comment!

[mehdi_amini]

This looks in good shape now!

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

469 ↗

(On Diff #82490)

This is a refactoring, since there are two places it needs to be called now.

[tejohnson]

I need to read it through again, but a few comments/questions based on my first pass.

Nit: since "data" is plural I also tend to think of "metadata" as being plural (as well as singular as "data" is commonly used as a singular today). So personally I prefer metadata over metadatas. I noticed a lot of places in the patch where "metadatas" is used.

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

95 ↗

(On Diff #82645)

Can there be a test added based on these strings and the option added below to verify e.g. that the number of records loaded is lower in the lazy-loading case?

414 ↗

(On Diff #82645)

s/associated to/associated with/

624 ↗

(On Diff #82645)

s/expliticitly/explicitly/

625 ↗

(On Diff #82645)

Should NumMDRecordLoaded be incremented in this case too? Maybe it would be better to CSE the increment of that down into a single place at the end of the BitstreamEntry::Record block (i.e. after the switch on Code, so after any case that doesn't cause an early return)?

671 ↗

(On Diff #82645)

s/give on/give up on/

1501 ↗

(On Diff #82645)

This is largely the same as the code in parseMetadataStrings. Could it be merged with a callback to do the right thing with each resulting StringRef?

Also, since we are essentially always parsing all the strings, what is the cost of always adding them to the MetadataList as MDStrings? Is this a memory savings?

llvm/lib/Transforms/IPO/FunctionImport.cpp

700 ↗

(On Diff #82645)

s/materialize/materializing/ and maybe "all the required globals"?

701 ↗

(On Diff #82645)

Maybe "all the required metadata" (since not all metadata are loaded now)?

[mehdi_amini]

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

625 ↗

(On Diff #82645)

Good point!

1501 ↗

(On Diff #82645)

We'd need to hit the context to add the string to a map and unique them there at the same time. I figured that if we load a single function for importing we likely want a tiny portion of the MDString.

[mehdi_amini]

Address comments:

• typos
• refactor MDString indexing using a callback to differentiate between parsing and indexing.
• Add a test showing the lazy-loading in action using the statistics.

[mehdi_amini]

Test cleanup

[mehdi_amini]

Post-break Ping :)

[tejohnson]

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

468 ↗

(On Diff #82690)

Use "///" instead for doxygen

615 ↗

(On Diff #82690)

Do we know if it is distinct or not?

771 ↗

(On Diff #82690)

I had to parse this condition a few times - might be worth a comment. Basically it is looking to see if we have left a temporary, in which case we need to parse the metadata record, otherwise we have a valid node of some kind, right?

782 ↗

(On Diff #82690)

Maybe add "/* ModuleLevel */" in front of "false" param. But why is this hardcoded to false? Aren't we module level if we are lazy loading?

787 ↗

(On Diff #82690)

From the name and the param, it sounds like this routine is only flushing the placeholders. Suggest renaming to something like resolveForwardRefsAndPlaceholders?

[mehdi_amini]

Address comments

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

615 ↗

(On Diff #82690)

The "IsDistinct" used to decide for placeholder is referring to the parent of the metadata we're about to load. NamedMD are always "distinct" AFAIK.
The issue however is that placeholder works only with MDNode which NamedMD are not. So while conceptually possible it requires some deeper change (I'll update the comment, I wrote it before figuring out the limitation)

771 ↗

(On Diff #82690)

That was the intent but I don't think it is useful, it is supposed to be already checked at call sites. So I replaced it with an assertion.

782 ↗

(On Diff #82690)

I removed the "ModuleLevel" parameter from parseOneMetadata, it was more useful originally but no longer.

[tejohnson]

LGTM

llvm/trunk/lib/Bitcode/Reader/MetadataLoader.cpp

 //===- MetadataLoader.cpp - Internal BitcodeReader implementation ---------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "MetadataLoader.h"
 #include "ValueList.h"
 
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/BitstreamReader.h"
 #include "llvm/Bitcode/LLVMBitCodes.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"

 #include <string>
 #include <system_error>
 #include <tuple>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
+#define DEBUG_TYPE "bitcode-reader"
+
+STATISTIC(NumMDStringLoaded, "Number of MDStrings loaded");
+STATISTIC(NumMDNodeTemporary, "Number of MDNode::Temporary created");
+STATISTIC(NumMDRecordLoaded, "Number of Metadata records loaded");
+
 /// Flag whether we need to import full type definitions for ThinLTO.
 /// Currently needed for Darwin and LLDB.
 static cl::opt<bool> ImportFullTypeDefinitions(
     "import-full-type-definitions", cl::init(false), cl::Hidden,
     cl::desc("Import full type definitions for ThinLTO."));
 
+static cl::opt<bool> DisableLazyLoading(
+    "disable-ondemand-mds-loading", cl::init(false), cl::Hidden,
+    cl::desc("Force disable the lazy-loading on-demand of metadata when "
+             "loading bitcode for importing."));
+
 namespace {
 
 static int64_t unrotateSign(uint64_t U) { return U & 1 ? ~(U >> 1) : U >> 1; }
 
 class BitcodeReaderMetadataList {
   /// Array of metadata references.
   ///
   /// Don't use std::vector here.  Some versions of libc++ copy (instead of
public:
   /// Gives the same result as \a lookup(), unless \a MDNode::isResolved()
   /// would give \c false.
   Metadata *getMetadataIfResolved(unsigned Idx);
 
   MDNode *getMDNodeFwdRefOrNull(unsigned Idx);
   void assignValue(Metadata *MD, unsigned Idx);
   void tryToResolveCycles();
   bool hasFwdRefs() const { return !ForwardReference.empty(); }
+  int getNextFwdRef() {
+    assert(hasFwdRefs());
+    return *ForwardReference.begin();
+  }
 
   /// Upgrade a type that had an MDString reference.
   void addTypeRef(MDString &UUID, DICompositeType &CT);
 
   /// Upgrade a type that had an MDString reference.
   Metadata *upgradeTypeRef(Metadata *MaybeUUID);
 
   /// Upgrade a type ref array that may have MDString references.
Metadata *BitcodeReaderMetadataList::getMetadataFwdRef(unsigned Idx) {
 
   if (Metadata *MD = MetadataPtrs[Idx])
     return MD;
 
   // Track forward refs to be resolved later.
   ForwardReference.insert(Idx);
 
   // Create and return a placeholder, which will later be RAUW'd.
+  ++NumMDNodeTemporary;
   Metadata *MD = MDNode::getTemporary(Context, None).release();
   MetadataPtrs[Idx].reset(MD);
   return MD;
 }
 
 Metadata *BitcodeReaderMetadataList::getMetadataIfResolved(unsigned Idx) {
   Metadata *MD = lookup(Idx);
   if (auto *N = dyn_cast_or_null<MDNode>(MD))

 namespace {
 
 class PlaceholderQueue {
   // Placeholders would thrash around when moved, so store in a std::deque
   // instead of some sort of vector.
   std::deque<DistinctMDOperandPlaceholder> PHs;
 
 public:
+  bool empty() { return PHs.empty(); }
   DistinctMDOperandPlaceholder &getPlaceholderOp(unsigned ID);
   void flush(BitcodeReaderMetadataList &MetadataList);
+
+  /// Return the list of temporaries nodes in the queue, these need to be
+  /// loaded before we can flush the queue.
+  void getTemporaries(BitcodeReaderMetadataList &MetadataList,
+                      DenseSet<unsigned> &Temporaries) {
+    for (auto &PH : PHs) {
+      auto ID = PH.getID();
+      auto *MD = MetadataList.lookup(ID);
+      if (!MD) {
+        Temporaries.insert(ID);
+        continue;
+      }
+      auto *N = dyn_cast_or_null<MDNode>(MD);
+      if (N && N->isTemporary())
+        Temporaries.insert(ID);
+    }
+  }
 };
 
 } // end anonymous namespace
 
 DistinctMDOperandPlaceholder &PlaceholderQueue::getPlaceholderOp(unsigned ID) {
   PHs.emplace_back(ID);
   return PHs.back();
 }

 class MetadataLoader::MetadataLoaderImpl {
   BitcodeReaderMetadataList MetadataList;
   BitcodeReaderValueList &ValueList;
   BitstreamCursor &Stream;
   LLVMContext &Context;
   Module &TheModule;
   std::function<Type *(unsigned)> getTypeByID;
 
+  /// Cursor associated with the lazy-loading of Metadata. This is the easy way
+  /// to keep around the right "context" (Abbrev list) to be able to jump in
+  /// the middle of the metadata block and load any record.
+  BitstreamCursor IndexCursor;
+
+  /// Index that keeps track of MDString values.
+  std::vector<StringRef> MDStringRef;
+
+  /// On-demand loading of a single MDString. Requires the index above to be
+  /// populated.
+  MDString *lazyLoadOneMDString(unsigned Idx);
+
+  /// Index that keeps track of where to find a metadata record in the stream.
+  std::vector<uint64_t> GlobalMetadataBitPosIndex;
+
+  /// Populate the index above to enable lazily loading of metadata, and load
+  /// the named metadata as well as the transitively referenced global
+  /// Metadata.
+  Expected<bool> lazyLoadModuleMetadataBlock(PlaceholderQueue &Placeholders);
+
+  /// On-demand loading of a single metadata. Requires the index above to be
+  /// populated.
+  void lazyLoadOneMetadata(unsigned Idx, PlaceholderQueue &Placeholders);
+
   // Keep mapping of seens pair of old-style CU <-> SP, and update pointers to
   // point from SP to CU after a block is completly parsed.
   std::vector<std::pair<DICompileUnit *, Metadata *>> CUSubprograms;
 
   /// Functions that need to be matched with subprograms when upgrading old
   /// metadata.
   SmallDenseMap<Function *, DISubprogram *, 16> FunctionsWithSPs;
 
   // Map the bitcode's custom MDKind ID to the Module's MDKind ID.
   DenseMap<unsigned, unsigned> MDKindMap;
 
   bool StripTBAA = false;
   bool HasSeenOldLoopTags = false;
 
   /// True if metadata is being parsed for a module being ThinLTO imported.
   bool IsImporting = false;
 
   Error parseOneMetadata(SmallVectorImpl<uint64_t> &Record, unsigned Code,
                          PlaceholderQueue &Placeholders, StringRef Blob,
-                         bool ModuleLevel, unsigned &NextMetadataNo);
-  Error parseMetadataStrings(ArrayRef<uint64_t> Record, StringRef Blob,
-                             unsigned &NextMetadataNo);
+  Error parseMetadataStrings(ArrayRef<uint64_t> Record, StringRef Blob,
+                             std::function<void(StringRef)> CallBack);
   Error parseGlobalObjectAttachment(GlobalObject &GO,
                                     ArrayRef<uint64_t> Record);
   Error parseMetadataKindRecord(SmallVectorImpl<uint64_t> &Record);
 
+  void resolveForwardRefsAndPlaceholders(PlaceholderQueue &Placeholders);
+
+  /// Upgrade old-style CU <-> SP pointers to point from SP to CU.
+  void upgradeCUSubprograms() {
+    for (auto CU_SP : CUSubprograms)
+      if (auto *SPs = dyn_cast_or_null<MDTuple>(CU_SP.second))
+        for (auto &Op : SPs->operands())
+          if (auto *SP = dyn_cast_or_null<MDNode>(Op))
+            SP->replaceOperandWith(7, CU_SP.first);
+    CUSubprograms.clear();
+  }
+
 public:
   MetadataLoaderImpl(BitstreamCursor &Stream, Module &TheModule,
                      BitcodeReaderValueList &ValueList,
                      std::function<Type *(unsigned)> getTypeByID,
                      bool IsImporting)
       : MetadataList(TheModule.getContext()), ValueList(ValueList),
         Stream(Stream), Context(TheModule.getContext()), TheModule(TheModule),
         getTypeByID(getTypeByID), IsImporting(IsImporting) {}
public:
   void shrinkTo(unsigned N) { MetadataList.shrinkTo(N); }
 };
 
 Error error(const Twine &Message) {
   return make_error<StringError>(
       Message, make_error_code(BitcodeError::CorruptedBitcode));
 }
 
+Expected<bool> MetadataLoader::MetadataLoaderImpl::lazyLoadModuleMetadataBlock(
+    PlaceholderQueue &Placeholders) {
+  IndexCursor = Stream;
+  SmallVector<uint64_t, 64> Record;
+  // Get the abbrevs, and preload record positions to make them lazy-loadable.
+  while (true) {
+    BitstreamEntry Entry = IndexCursor.advanceSkippingSubblocks(
+        BitstreamCursor::AF_DontPopBlockAtEnd);
+    switch (Entry.Kind) {
+    case BitstreamEntry::SubBlock: // Handled for us already.
+    case BitstreamEntry::Error:
+      return error("Malformed block");
+    case BitstreamEntry::EndBlock: {
+      return true;
+    }
+    case BitstreamEntry::Record: {
+      // The interesting case.
+      ++NumMDRecordLoaded;
+      uint64_t CurrentPos = IndexCursor.GetCurrentBitNo();
+      auto Code = IndexCursor.skipRecord(Entry.ID);
+      switch (Code) {
+      case bitc::METADATA_STRINGS: {
+        // Rewind and parse the strings.
+        IndexCursor.JumpToBit(CurrentPos);
+        StringRef Blob;
+        Record.clear();
+        IndexCursor.readRecord(Entry.ID, Record, &Blob);
+        unsigned NumStrings = Record[0];
+        MDStringRef.reserve(NumStrings);
+        auto IndexNextMDString = [&](StringRef Str) {
+          MDStringRef.push_back(Str);
+        };
+        if (auto Err = parseMetadataStrings(Record, Blob, IndexNextMDString))
+          return std::move(Err);
+        break;
+      }
+      case bitc::METADATA_INDEX_OFFSET: {
+        // This is the offset to the index, when we see this we skip all the
+        // records and load only an index to these.
+        IndexCursor.JumpToBit(CurrentPos);
+        Record.clear();
+        IndexCursor.readRecord(Entry.ID, Record);
+        if (Record.size() != 2)
+          return error("Invalid record");
+        auto Offset = Record[0] + (Record[1] << 32);
+        auto BeginPos = IndexCursor.GetCurrentBitNo();
+        IndexCursor.JumpToBit(BeginPos + Offset);
+        Entry = IndexCursor.advanceSkippingSubblocks(
+            BitstreamCursor::AF_DontPopBlockAtEnd);
+        assert(Entry.Kind == BitstreamEntry::Record &&
+               "Corrupted bitcode: Expected `Record` when trying to find the "
+               "Metadata index");
+        Record.clear();
+        auto Code = IndexCursor.readRecord(Entry.ID, Record);
+        (void)Code;
+        assert(Code == bitc::METADATA_INDEX && "Corrupted bitcode: Expected "
+                                               "`METADATA_INDEX` when trying "
+                                               "to find the Metadata index");
+
+        // Delta unpack
+        auto CurrentValue = BeginPos;
+        GlobalMetadataBitPosIndex.reserve(Record.size());
+        for (auto &Elt : Record) {
+          CurrentValue += Elt;
+          GlobalMetadataBitPosIndex.push_back(CurrentValue);
+        }
+        break;
+      }
+      case bitc::METADATA_INDEX:
+        // We don't expect to get there, the Index is loaded when we encounter
+        // the offset.
+        return error("Corrupted Metadata block");
+      case bitc::METADATA_NAME: {
+        // Named metadata need to be materialized now and aren't deferred.
+        IndexCursor.JumpToBit(CurrentPos);
+        Record.clear();
+        unsigned Code = IndexCursor.readRecord(Entry.ID, Record);
+        assert(Code == bitc::METADATA_NAME);
+
+        // Read name of the named metadata.
+        SmallString<8> Name(Record.begin(), Record.end());
+        Code = IndexCursor.ReadCode();
+
+        // Named Metadata comes in two parts, we expect the name to be followed
+        // by the node
+        Record.clear();
+        unsigned NextBitCode = IndexCursor.readRecord(Code, Record);
+        assert(NextBitCode == bitc::METADATA_NAMED_NODE);
+        (void)NextBitCode;
+
+        // Read named metadata elements.
+        unsigned Size = Record.size();
+        NamedMDNode *NMD = TheModule.getOrInsertNamedMetadata(Name);
+        for (unsigned i = 0; i != Size; ++i) {
+          // FIXME: We could use a placeholder here, however NamedMDNode are
+          // taking MDNode as operand and not using the Metadata infrastructure.
+          // It is acknowledged by 'TODO: Inherit from Metadata' in the
+          // NamedMDNode class definition.
+          MDNode *MD = MetadataList.getMDNodeFwdRefOrNull(Record[i]);
+          assert(MD && "Invalid record");
+          NMD->addOperand(MD);
+        }
+        break;
+      }
+      case bitc::METADATA_GLOBAL_DECL_ATTACHMENT: {
+        // FIXME: we need to do this early because we don't materialize global
+        // value explicitly.
+        IndexCursor.JumpToBit(CurrentPos);
+        Record.clear();
+        IndexCursor.readRecord(Entry.ID, Record);
+        if (Record.size() % 2 == 0)
+          return error("Invalid record");
+        unsigned ValueID = Record[0];
+        if (ValueID >= ValueList.size())
+          return error("Invalid record");
+        if (auto *GO = dyn_cast<GlobalObject>(ValueList[ValueID]))
+          if (Error Err = parseGlobalObjectAttachment(
+                  *GO, ArrayRef<uint64_t>(Record).slice(1)))
+            return std::move(Err);
+        break;
+      }
+      case bitc::METADATA_KIND:
+      case bitc::METADATA_STRING_OLD:
+      case bitc::METADATA_OLD_FN_NODE:
+      case bitc::METADATA_OLD_NODE:
+      case bitc::METADATA_VALUE:
+      case bitc::METADATA_DISTINCT_NODE:
+      case bitc::METADATA_NODE:
+      case bitc::METADATA_LOCATION:
+      case bitc::METADATA_GENERIC_DEBUG:
+      case bitc::METADATA_SUBRANGE:
+      case bitc::METADATA_ENUMERATOR:
+      case bitc::METADATA_BASIC_TYPE:
+      case bitc::METADATA_DERIVED_TYPE:
+      case bitc::METADATA_COMPOSITE_TYPE:
+      case bitc::METADATA_SUBROUTINE_TYPE:
+      case bitc::METADATA_MODULE:
+      case bitc::METADATA_FILE:
+      case bitc::METADATA_COMPILE_UNIT:
+      case bitc::METADATA_SUBPROGRAM:
+      case bitc::METADATA_LEXICAL_BLOCK:
+      case bitc::METADATA_LEXICAL_BLOCK_FILE:
+      case bitc::METADATA_NAMESPACE:
+      case bitc::METADATA_MACRO:
+      case bitc::METADATA_MACRO_FILE:
+      case bitc::METADATA_TEMPLATE_TYPE:
+      case bitc::METADATA_TEMPLATE_VALUE:
+      case bitc::METADATA_GLOBAL_VAR:
+      case bitc::METADATA_LOCAL_VAR:
+      case bitc::METADATA_EXPRESSION:
+      case bitc::METADATA_OBJC_PROPERTY:
+      case bitc::METADATA_IMPORTED_ENTITY:
+      case bitc::METADATA_GLOBAL_VAR_EXPR:
+        // We don't expect to see any of these, if we see one, give up on
+        // lazy-loading and fallback.
+        MDStringRef.clear();
+        GlobalMetadataBitPosIndex.clear();
+        return false;
+      }
+      break;
+    }
+    }
+  }
+}
+
 /// Parse a METADATA_BLOCK. If ModuleLevel is true then we are parsing
 /// module level metadata.
 Error MetadataLoader::MetadataLoaderImpl::parseMetadata(bool ModuleLevel) {
   if (!ModuleLevel && MetadataList.hasFwdRefs())
     return error("Invalid metadata: fwd refs into function blocks");
 
+  // Record the entry position so that we can jump back here and efficiently
+  // skip the whole block in case we lazy-load.
+  auto EntryPos = Stream.GetCurrentBitNo();
+
   if (Stream.EnterSubBlock(bitc::METADATA_BLOCK_ID))
     return error("Invalid record");
 
-  unsigned NextMetadataNo = MetadataList.size();
   SmallVector<uint64_t, 64> Record;
-
   PlaceholderQueue Placeholders;
 
+  // We lazy-load module-level metadata: we build an index for each record, and
+  // then load individual record as needed, starting with the named metadata.
+  if (ModuleLevel && IsImporting && MetadataList.empty() &&
+      !DisableLazyLoading) {
+    auto SuccessOrErr = lazyLoadModuleMetadataBlock(Placeholders);
+    if (!SuccessOrErr)
+      return SuccessOrErr.takeError();
+    if (SuccessOrErr.get()) {
+      // An index was successfully created and we will be able to load metadata
+      // on-demand.
+      MetadataList.resize(MDStringRef.size() +
+                          GlobalMetadataBitPosIndex.size());
+
+      // Reading the named metadata created forward references and/or
+      // placeholders, that we flush here.
+      resolveForwardRefsAndPlaceholders(Placeholders);
+      upgradeCUSubprograms();
+      // Return at the beginning of the block, since it is easy to skip it
+      // entirely from there.
+      Stream.ReadBlockEnd(); // Pop the abbrev block context.
+      Stream.JumpToBit(EntryPos);
+      if (Stream.SkipBlock())
+        return error("Invalid record");
+      return Error::success();
+    }
+    // Couldn't load an index, fallback to loading all the block "old-style".
+  }
+
+  unsigned NextMetadataNo = MetadataList.size();
+
   // Read all the records.
   while (true) {
     BitstreamEntry Entry = Stream.advanceSkippingSubblocks();
 
     switch (Entry.Kind) {
     case BitstreamEntry::SubBlock: // Handled for us already.
     case BitstreamEntry::Error:
       return error("Malformed block");
     case BitstreamEntry::EndBlock:
-      // Upgrade old-style CU <-> SP pointers to point from SP to CU.
-      for (auto CU_SP : CUSubprograms)
-        if (auto *SPs = dyn_cast_or_null<MDTuple>(CU_SP.second))
-          for (auto &Op : SPs->operands())
-            if (auto *SP = dyn_cast_or_null<MDNode>(Op))
-              SP->replaceOperandWith(7, CU_SP.first);
-      CUSubprograms.clear();
-
-      MetadataList.tryToResolveCycles();
-      Placeholders.flush(MetadataList);
+      resolveForwardRefsAndPlaceholders(Placeholders);
+      upgradeCUSubprograms();
       return Error::success();
     case BitstreamEntry::Record:
       // The interesting case.
       break;
     }
 
     // Read a record.
     Record.clear();
     StringRef Blob;
+    ++NumMDRecordLoaded;
     unsigned Code = Stream.readRecord(Entry.ID, Record, &Blob);
-    if (Error Err = parseOneMetadata(Record, Code, Placeholders, Blob,
-                                     ModuleLevel, NextMetadataNo))
+    if (Error Err =
+            parseOneMetadata(Record, Code, Placeholders, Blob, NextMetadataNo))
       return Err;
   }
 }
 
+MDString *MetadataLoader::MetadataLoaderImpl::lazyLoadOneMDString(unsigned ID) {
+  ++NumMDStringLoaded;
+  if (Metadata *MD = MetadataList.lookup(ID))
+    return cast<MDString>(MD);
+  auto MDS = MDString::get(Context, MDStringRef[ID]);
+  MetadataList.assignValue(MDS, ID);
+  return MDS;
+}
+
+void MetadataLoader::MetadataLoaderImpl::lazyLoadOneMetadata(
+    unsigned ID, PlaceholderQueue &Placeholders) {
+  assert(ID < (MDStringRef.size()) + GlobalMetadataBitPosIndex.size());
+  assert(ID >= MDStringRef.size() && "Unexpected lazy-loading of MDString");
+#ifndef NDEBUG
+  // Lookup first if the metadata hasn't already been loaded.
+  if (auto *MD = MetadataList.lookup(ID)) {
+    auto *N = dyn_cast_or_null<MDNode>(MD);
+    assert(N && N->isTemporary() && "Lazy loading an already loaded metadata");
+  }
+#endif
+  SmallVector<uint64_t, 64> Record;
+  StringRef Blob;
+  IndexCursor.JumpToBit(GlobalMetadataBitPosIndex[ID - MDStringRef.size()]);
+  auto Entry = IndexCursor.advanceSkippingSubblocks();
+  ++NumMDRecordLoaded;
+  unsigned Code = IndexCursor.readRecord(Entry.ID, Record, &Blob);
+  if (Error Err = parseOneMetadata(Record, Code, Placeholders, Blob, ID))
+    report_fatal_error("Can't lazyload MD");
+}
+
+/// Ensure that all forward-references and placeholders are resolved.
+/// Iteratively lazy-loading metadata on-demand if needed.
+void MetadataLoader::MetadataLoaderImpl::resolveForwardRefsAndPlaceholders(
+    PlaceholderQueue &Placeholders) {
+  DenseSet<unsigned> Temporaries;
+  while (1) {
+    // Populate Temporaries with the placeholders that haven't been loaded yet.
+    Placeholders.getTemporaries(MetadataList, Temporaries);
+
+    // If we don't have any temporary, or FwdReference, we're done!
+    if (Temporaries.empty() && !MetadataList.hasFwdRefs())
+      break;
+
+    // First, load all the temporaries. This can add new placeholders or
+    // forward references.
+    for (auto ID : Temporaries)
+      lazyLoadOneMetadata(ID, Placeholders);
+    Temporaries.clear();
+
+    // Second, load the forward-references. This can also add new placeholders
+    // or forward references.
+    while (MetadataList.hasFwdRefs())
+      lazyLoadOneMetadata(MetadataList.getNextFwdRef(), Placeholders);
+  }
+  // At this point we don't have any forward reference remaining, or temporary
+  // that haven't been loaded. We can safely drop RAUW support and mark cycles
+  // as resolved.
+  MetadataList.tryToResolveCycles();
+
+  // Finally, everything is in place, we can replace the placeholders operands
+  // with the final node they refer to.
+  Placeholders.flush(MetadataList);
+}
+
 Error MetadataLoader::MetadataLoaderImpl::parseOneMetadata(
     SmallVectorImpl<uint64_t> &Record, unsigned Code,
-    PlaceholderQueue &Placeholders, StringRef Blob, bool ModuleLevel,
-    unsigned &NextMetadataNo) {
+    PlaceholderQueue &Placeholders, StringRef Blob, unsigned &NextMetadataNo) {
 
   bool IsDistinct = false;
   auto getMD = [&](unsigned ID) -> Metadata * {
+    if (ID < MDStringRef.size())
+      return lazyLoadOneMDString(ID);
     if (!IsDistinct)
       return MetadataList.getMetadataFwdRef(ID);
     if (auto *MD = MetadataList.getMetadataIfResolved(ID))
       return MD;
     return &Placeholders.getPlaceholderOp(ID);
   };
   auto getMDOrNull = [&](unsigned ID) -> Metadata * {
     if (ID)
       return getMD(ID - 1);
     return nullptr;
   };
   auto getMDOrNullWithoutPlaceholders = [&](unsigned ID) -> Metadata * {
     if (ID)
       return MetadataList.getMetadataFwdRef(ID - 1);
     return nullptr;
   };
   auto getMDString = [&](unsigned ID) -> MDString * {
     // This requires that the ID is not really a forward reference.  In
     // particular, the MDString must already have been resolved.
-    return cast_or_null<MDString>(getMDOrNull(ID));
+    auto MDS = getMDOrNull(ID);
+    return cast_or_null<MDString>(MDS);
   };
 
   // Support for old type refs.
   auto getDITypeRefOrNull = [&](unsigned ID) {
     return MetadataList.upgradeTypeRef(getMDOrNull(ID));
   };
 
 #define GET_OR_DISTINCT(CLASS, ARGS)                                           \
   (IsDistinct ? CLASS::getDistinct ARGS : CLASS::get ARGS)
 
   switch (Code) {
   default: // Default behavior: ignore.
     break;
   case bitc::METADATA_NAME: {
     // Read name of the named metadata.
     SmallString<8> Name(Record.begin(), Record.end());
     Record.clear();
     Code = Stream.ReadCode();
 
+    ++NumMDRecordLoaded;
     unsigned NextBitCode = Stream.readRecord(Code, Record);
     if (NextBitCode != bitc::METADATA_NAMED_NODE)
       return error("METADATA_NAME not followed by METADATA_NAMED_NODE");
 
     // Read named metadata elements.
     unsigned Size = Record.size();
     NamedMDNode *NMD = TheModule.getOrInsertNamedMetadata(Name);
     for (unsigned i = 0; i != Size; ++i) {
MetadataList.assignValue(
         NextMetadataNo++);
     break;
   }
   case bitc::METADATA_STRING_OLD: {
     std::string String(Record.begin(), Record.end());
 
     // Test for upgrading !llvm.loop.
     HasSeenOldLoopTags |= mayBeOldLoopAttachmentTag(String);
-
+    ++NumMDStringLoaded;
     Metadata *MD = MDString::get(Context, String);
     MetadataList.assignValue(MD, NextMetadataNo++);
     break;
   }
-  case bitc::METADATA_STRINGS:
-    if (Error Err = parseMetadataStrings(Record, Blob, NextMetadataNo))
+  case bitc::METADATA_STRINGS: {
+    auto CreateNextMDString = [&](StringRef Str) {
+      ++NumMDStringLoaded;
+      MetadataList.assignValue(MDString::get(Context, Str), NextMetadataNo++);
+    };
+    if (Error Err = parseMetadataStrings(Record, Blob, CreateNextMDString))
       return Err;
     break;
+  }
   case bitc::METADATA_GLOBAL_DECL_ATTACHMENT: {
     if (Record.size() % 2 == 0)
       return error("Invalid record");
     unsigned ValueID = Record[0];
     if (ValueID >= ValueList.size())
       return error("Invalid record");
     if (auto *GO = dyn_cast<GlobalObject>(ValueList[ValueID]))
       if (Error Err = parseGlobalObjectAttachment(
               *GO, ArrayRef<uint64_t>(Record).slice(1)))
         return Err;
     break;
   }
   case bitc::METADATA_KIND: {
     // Support older bitcode files that had METADATA_KIND records in a
     // block with METADATA_BLOCK_ID.
     if (Error Err = parseMetadataKindRecord(Record))
       return Err;
     break;
   }
   }
-#undef GET_OR_DISTINCT
   return Error::success();
+#undef GET_OR_DISTINCT
 }
 
 Error MetadataLoader::MetadataLoaderImpl::parseMetadataStrings(
-    ArrayRef<uint64_t> Record, StringRef Blob, unsigned &NextMetadataNo) {
+    ArrayRef<uint64_t> Record, StringRef Blob,
+    std::function<void(StringRef)> CallBack) {
   // All the MDStrings in the block are emitted together in a single
   // record.  The strings are concatenated and stored in a blob along with
   // their sizes.
   if (Record.size() != 2)
     return error("Invalid record: metadata strings layout");
 
   unsigned NumStrings = Record[0];
   unsigned StringsOffset = Record[1];
Error MetadataLoader::MetadataLoaderImpl::parseMetadataStrings(
   do {
     if (R.AtEndOfStream())
       return error("Invalid record: metadata strings bad length");
 
     unsigned Size = R.ReadVBR(6);
     if (Strings.size() < Size)
       return error("Invalid record: metadata strings truncated chars");
 
-    MetadataList.assignValue(MDString::get(Context, Strings.slice(0, Size)),
-                             NextMetadataNo++);
+    CallBack(Strings.slice(0, Size));
     Strings = Strings.drop_front(Size);
   } while (--NumStrings);
 
   return Error::success();
 }
 
 Error MetadataLoader::MetadataLoaderImpl::parseGlobalObjectAttachment(
     GlobalObject &GO, ArrayRef<uint64_t> Record) {

 /// Parse metadata attachments.
 Error MetadataLoader::MetadataLoaderImpl::parseMetadataAttachment(
     Function &F, const SmallVectorImpl<Instruction *> &InstructionList) {
   if (Stream.EnterSubBlock(bitc::METADATA_ATTACHMENT_ID))
     return error("Invalid record");
 
   SmallVector<uint64_t, 64> Record;
 
+  PlaceholderQueue Placeholders;
+
   while (true) {
     BitstreamEntry Entry = Stream.advanceSkippingSubblocks();
 
     switch (Entry.Kind) {
     case BitstreamEntry::SubBlock: // Handled for us already.
     case BitstreamEntry::Error:
       return error("Malformed block");
     case BitstreamEntry::EndBlock:
+      resolveForwardRefsAndPlaceholders(Placeholders);
       return Error::success();
     case BitstreamEntry::Record:
       // The interesting case.
       break;
     }
 
     // Read a metadata attachment record.
     Record.clear();
+    ++NumMDRecordLoaded;
     switch (Stream.readRecord(Entry.ID, Record)) {
     default: // Default behavior: ignore.
       break;
     case bitc::METADATA_ATTACHMENT: {
       unsigned RecordLength = Record.size();
       if (Record.empty())
         return error("Invalid record");
       if (RecordLength % 2 == 0) {
         // A function attachment.
         if (Error Err = parseGlobalObjectAttachment(F, Record))
           return Err;
         continue;
       }
 
       // An instruction attachment.
       Instruction *Inst = InstructionList[Record[0]];
       for (unsigned i = 1; i != RecordLength; i = i + 2) {
         unsigned Kind = Record[i];
         DenseMap<unsigned, unsigned>::iterator I = MDKindMap.find(Kind);
         if (I == MDKindMap.end())
           return error("Invalid ID");
         if (I->second == LLVMContext::MD_tbaa && StripTBAA)
           continue;
 
-        Metadata *Node = MetadataList.getMetadataFwdRef(Record[i + 1]);
+        auto Idx = Record[i + 1];
+        if (Idx < (MDStringRef.size() + GlobalMetadataBitPosIndex.size()) &&
+            !MetadataList.lookup(Idx))
+          // Load the attachment if it is in the lazy-loadable range and hasn't
+          // been loaded yet.
+          lazyLoadOneMetadata(Idx, Placeholders);
+
+        Metadata *Node = MetadataList.getMetadataFwdRef(Idx);
         if (isa<LocalAsMetadata>(Node))
           // Drop the attachment.  This used to be legal, but there's no
           // upgrade path.
           break;
         MDNode *MD = dyn_cast_or_null<MDNode>(Node);
         if (!MD)
           return error("Invalid metadata attachment");
 
case BitstreamEntry::EndBlock:
       return Error::success();
     case BitstreamEntry::Record:
       // The interesting case.
       break;
     }
 
     // Read a record.
     Record.clear();
+    ++NumMDRecordLoaded;
     unsigned Code = Stream.readRecord(Entry.ID, Record);
     switch (Code) {
     default: // Default behavior: ignore.
       break;
     case bitc::METADATA_KIND: {
       if (Error Err = parseMetadataKindRecord(Record))
         return Err;
       break;

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

for (auto &Name : ModuleNameOrderedList) {
     std::unique_ptr<Module> SrcModule = std::move(*SrcModuleOrErr);
     assert(&DestModule.getContext() == &SrcModule->getContext() &&
            "Context mismatch");
 
     // If modules were created with lazy metadata loading, materialize it
     // now, before linking it (otherwise this will be a noop).
     if (Error Err = SrcModule->materializeMetadata())
       return std::move(Err);
-    UpgradeDebugInfo(*SrcModule);
 
     auto &ImportGUIDs = FunctionsToImportPerModule->second;
     // Find the globals to import
     DenseSet<const GlobalValue *> GlobalsToImport;
     for (Function &F : *SrcModule) {
       if (!F.hasName())
         continue;
       auto GUID = F.getGUID();
#endif
           return std::move(Err);
         GlobalsToImport.insert(GO);
         if (Error Err = GA.materialize())
           return std::move(Err);
         GlobalsToImport.insert(&GA);
       }
     }
 
+    // Upgrade debug info after we're done materializing all the globals and we
+    // have loaded all the required metadata!
+    UpgradeDebugInfo(*SrcModule);
+
     // Link in the specified functions.
     if (renameModuleForThinLTO(*SrcModule, Index, &GlobalsToImport))
       return true;
 
     if (PrintImports) {
       for (const auto *GV : GlobalsToImport)
         dbgs() << DestModule.getSourceFileName() << ": Import " << GV->getName()
                << " from " << SrcModule->getSourceFileName() << "\n";

llvm/trunk/test/ThinLTO/X86/Inputs/lazyload_metadata.ll

+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.11.0"
+
+declare void @globalfunc1()
+
+
+define i32 @main() {
+  call void @globalfunc1()
+  ret i32 0
+}
+
+

llvm/trunk/test/ThinLTO/X86/lazyload_metadata.ll

+; Do setup work for all below tests: generate bitcode and combined index
+; RUN: opt -module-summary %s -o %t.bc -bitcode-mdindex-threshold=0
+; RUN: opt -module-summary %p/Inputs/lazyload_metadata.ll -o %t2.bc -bitcode-mdindex-threshold=0
+; RUN: llvm-lto -thinlto-action=thinlink -o %t3.bc %t.bc %t2.bc
+
+
+; Check that importing @globalfunc1 does not trigger loading all the global
+; metadata for @globalfunc2 and @globalfunc3
+
+; RUN: llvm-lto -thinlto-action=import %t2.bc -thinlto-index=%t3.bc \
+; RUN:          -o /dev/null -stats \
+; RUN:  2>&1 | FileCheck %s -check-prefix=LAZY
+; LAZY: 49 bitcode-reader  - Number of Metadata records loaded
+; LAZY: 1 bitcode-reader  - Number of MDStrings loaded
+
+; RUN: llvm-lto -thinlto-action=import %t2.bc -thinlto-index=%t3.bc \
+; RUN:          -o /dev/null -disable-ondemand-mds-loading -stats \
+; RUN:  2>&1 | FileCheck %s -check-prefix=NOTLAZY
+; NOTLAZY: 58 bitcode-reader  - Number of Metadata records loaded
+; NOTLAZY: 8 bitcode-reader  - Number of MDStrings loaded
+
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.11.0"
+
+define void @globalfunc1(i32 %arg) {
+  %tmp = add i32 %arg, 0, !metadata !2
+  ret void
+}
+
+; We need two functions here that will both reference the same metadata.
+; This is to force the metadata to be emitted in the global metadata block and
+; not in the function specific metadata.
+; These function are not imported and so we don't want to load their metadata.
+
+define void @globalfunc2(i32 %arg) {
+  %tmp = add i32 %arg, 0, !metadata !1
+  ret void
+}
+
+define void @globalfunc3(i32 %arg) {
+  %tmp = add i32 %arg, 0, !metadata !1
+  ret void
+}
+
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"Hello World"}
+!3 = !{!"3"}
+!4 = !{!"4"}
+!5 = !{!"5"}
+!6 = !{!"6"}
+!7 = !{!"7"}
+!8 = !{!"8"}
+!9 = !{!"9"}