[ThinLTO] Avoid temporaries when loading global decl attachment metadata

[tejohnson]

Summary

When performing ThinLTO importing, the metadata loader attempts to lazy
load, by building an index. However, module level global decl attachment
metadata was being parsed early while building the index, since the
associated (module level) global values aren't materialized on demand.
This results in the creation of forward reference temporary metadatas,
which are expensive.

Normally, these module level global values don't have much attached
metadata. However, in the case of -fwhole-program-vtables (e.g. for
whole program devirtualization), the vtables may have many attached type
metadatas. This was resulting in very slow performance when performing
ThinLTO importing with the default lazy loading.

This patch restructures the handling of these global decl attachment
records, delaying their parsing until after the lazy loading index has
been built. Then the parser can use the interface that loads from the
index, which resolves forward references immediately instead of creating
expensive temporaries.

For one ThinLTO backend that imports from modules containing huge
numbers of vtables and associated types, I measured the following
compile times for the metadata materialization during function
importing, rounded to nearest second:

No -fwhole-program-vtables:

Lazy loading on (head):  1s
Lazy loading off (head): 3s
Lazy loading on (patch): 1s

With -fwhole-program-vtables:

Lazy loading on (head):  440s
Lazy loading off (head): 4s
Lazy loading on (patch): 2s

[tejohnson]

I realized after intentionally suppressing the new method that there was no regression test that broke. It turns out that the lazy loading index is only emitted if there is a certain number of metadatas, which the tests that exercise this code apparently don't have. Therefore, I added an option to drop the threshold for one of the ThinLTO WPD tests. This exposed a bug: Because parseGlobalObjectAttachment now resolves forward references via the index, the index cursor has moved when it returns, and we weren't parsing all of the global decl attachment metadatas. Fixed by saving and restoring the current position around that call.

My compile time test didn't expose this because it was NDEBUG. I confirmed that it now parses all the expected metadata records. The metadata materialization time under -fwhole-program-vtables goes from 1s up to 2s. Will update the summary accordingly.

[tejohnson]

Add test options that exercise code, fix exposed bug.

[evgeny777]

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

894

So this is not called during ThinLTO import, is it?

2115

Why is this change needed?

[tejohnson]

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

894

It is only called during ThinLTO importing. See the if conditions for the block of code at its callsite (line 989)

2115

Because getMDNodeFwdRefOrNull unconditionally creates a forward reference with temporary when the requested metadata is not yet loaded. Whereas getMetadataFwdRefOrLoad will use the lazy loading index built by lazyLoadModuleMetadataBlock to perform the load rather than create a temporary. With the code before this patch, that could not be done as the global decl attachment metadatas were being parsed also by lazyLoadModuleMetadataBlock, so the index wasn't available. And without making the change to this callsite, the patch doesn't have any measureable effect since we won't use the index even when we have it (which I discovered when I first made the change to restructure into loadGlobalDeclAttachments but didn't change this callsite).

[evgeny777]

LGTM with nits

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

916

nit: no need for {} here

929

nit: this is hard to read. May be change to somthing like this?

if (MaybeCode.get() != bitc::METADATA_GLOBAL_DECL_ATTACHMENT) {
      assert(NumGlobalDeclAttachSkipped == NumGlobalDeclAttachParsed);
      return true;
}

// Parse stuff here

2115

Thanks, now clear

[tejohnson]

Address comments

[MaskRay]

This caused a -fsanitize=cfi issue in thinltoBackend. I reverted the patch temporarily and shared a detailed reproduce with Teresa

% clang ...
fatal error: error in backend: Invalid abbrev number

[tejohnson]

PTAL

I have a fix for the issue that required a revert. We shouldn't be using the lazy loading IndexCursor to load the global decl attachment metadata. That cursor is initialized with the necessary abbrev ids from the module level metadata block at the end of the lazy loading index setup, so that you can use those abbrev ids when jumping into random locations, as noted in the comment above its declaration:

/// 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.

This is only exposed in limited situations, since there are almost no abbrev ids used in the metadata block (basically only debug location records), and you need a debug location that used in certain ways and shared by multiple functions so that it ends up in the module level metadata block for lazy loading. I beefed up the test case to contain metadata on the imported module that exposes the failure without this fix.

[tejohnson]

Fix bug and add test case.

[tejohnson]

Ping - could someone take a quick look at the update and approve the recommit?

[tejohnson]

ping, @evgeny777 can you take a quick look at the change. It's a pretty minor fix.

[evgeny777]

LGTM

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

class MetadataLoader::MetadataLoaderImpl {
 
   /// 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;
 
+  /// Cursor position of the start of the global decl attachments, to enable
+  /// loading using the index built for lazy loading, instead of forward
+  /// references.
+  uint64_t GlobalDeclAttachmentPos = 0;
+
+#ifndef NDEBUG
+  /// Sanity check that we end up parsing all of the global decl attachments.
+  unsigned NumGlobalDeclAttachSkipped = 0;
+  unsigned NumGlobalDeclAttachParsed = 0;
+#endif
+
+  /// Load the global decl attachments, using the index built for lazy loading.
+  Expected<bool> loadGlobalDeclAttachments();
+
   /// 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();
 
   /// On-demand loading of a single metadata. Requires the index above to be
   /// populated.
   void lazyLoadOneMetadata(unsigned Idx, PlaceholderQueue &Placeholders);
public:
   void shrinkTo(unsigned N) { MetadataList.shrinkTo(N); }
   void upgradeDebugIntrinsics(Function &F) { upgradeDeclareExpressions(F); }
 };
 
 Expected<bool>
 MetadataLoader::MetadataLoaderImpl::lazyLoadModuleMetadataBlock() {
   IndexCursor = Stream;
   SmallVector<uint64_t, 64> Record;
+  GlobalDeclAttachmentPos = 0;
   // Get the abbrevs, and preload record positions to make them lazy-loadable.
   while (true) {
+    uint64_t SavedPos = IndexCursor.GetCurrentBitNo();
     Expected<BitstreamEntry> MaybeEntry = IndexCursor.advanceSkippingSubblocks(
         BitstreamCursor::AF_DontPopBlockAtEnd);
     if (!MaybeEntry)
       return MaybeEntry.takeError();
     BitstreamEntry Entry = MaybeEntry.get();
 
     switch (Entry.Kind) {
     case BitstreamEntry::SubBlock: // Handled for us already.
case BitstreamEntry::Record: {
           // NamedMDNode class definition.
           MDNode *MD = MetadataList.getMDNodeFwdRefOrNull(Record[i]);
           assert(MD && "Invalid metadata: expect fwd ref to MDNode");
           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.
-        if (Error Err = IndexCursor.JumpToBit(CurrentPos))
-          return std::move(Err);
-        Record.clear();
-        if (Expected<unsigned> MaybeRecord =
-                IndexCursor.readRecord(Entry.ID, Record))
-          ;
-        else
-          return MaybeRecord.takeError();
-        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);
+        if (!GlobalDeclAttachmentPos)
+          GlobalDeclAttachmentPos = SavedPos;
+#ifndef NDEBUG
+        NumGlobalDeclAttachSkipped++;
+#endif
         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:
#endif
         return false;
       }
       break;
     }
     }
   }
 }
 
+// Load the global decl attachments after building the lazy loading index.
+// We don't load them "lazily" - all global decl attachments must be
+// parsed since they aren't materialized on demand. However, by delaying
+// their parsing until after the index is created, we can use the index
+// instead of creating temporaries.
+Expected<bool> MetadataLoader::MetadataLoaderImpl::loadGlobalDeclAttachments() {
+  // Nothing to do if we didn't find any of these metadata records.
+  if (!GlobalDeclAttachmentPos)
+    return true;
+  IndexCursor = Stream;
+  SmallVector<uint64_t, 64> Record;
+  // Jump to the position before the first global decl attachment, so we can
+  // scan for the first BitstreamEntry record.
+  if (Error Err = IndexCursor.JumpToBit(GlobalDeclAttachmentPos))
+    return std::move(Err);
+  while (true) {
+    Expected<BitstreamEntry> MaybeEntry = IndexCursor.advanceSkippingSubblocks(
+        BitstreamCursor::AF_DontPopBlockAtEnd);
+    if (!MaybeEntry)
+      return MaybeEntry.takeError();
+    BitstreamEntry Entry = MaybeEntry.get();
+
+    switch (Entry.Kind) {
+    case BitstreamEntry::SubBlock: // Handled for us already.
+    case BitstreamEntry::Error:
+      return error("Malformed block");
+    case BitstreamEntry::EndBlock: {
+      // Sanity check that we parsed them all.
+      assert(NumGlobalDeclAttachSkipped == NumGlobalDeclAttachParsed);
+      return true;
+    }
+    case BitstreamEntry::Record:
+      break;
+    }
+    uint64_t CurrentPos = IndexCursor.GetCurrentBitNo();
+    Expected<unsigned> MaybeCode = IndexCursor.skipRecord(Entry.ID);
+    if (!MaybeCode)
+      return MaybeCode.takeError();
+    unsigned Code = MaybeCode.get();
+    switch (Code) {
+    default:
+      // Anything other than a global decl attachment signals the end of
+      // these records. sanity check that we parsed them all.
+      assert(NumGlobalDeclAttachSkipped == NumGlobalDeclAttachParsed);
+      return true;
+    case bitc::METADATA_GLOBAL_DECL_ATTACHMENT: {
+#ifndef NDEBUG
+      NumGlobalDeclAttachParsed++;
+#endif
+      // FIXME: we need to do this early because we don't materialize global
+      // value explicitly.
+      if (Error Err = IndexCursor.JumpToBit(CurrentPos))
+        return std::move(Err);
+      Record.clear();
+      if (Expected<unsigned> MaybeRecord =
+              IndexCursor.readRecord(Entry.ID, Record))
+        ;
+      else
+        return MaybeRecord.takeError();
+      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])) {
+        // Need to save and restore the current parser since
+        // parseGlobalObjectAttachment will resolve all forward references which
+        // would require parsing from locations stored in the index.
+        CurrentPos = IndexCursor.GetCurrentBitNo();
+        if (Error Err = parseGlobalObjectAttachment(
+                *GO, ArrayRef<uint64_t>(Record).slice(1)))
+          return std::move(Err);
+        if (Error Err = IndexCursor.JumpToBit(CurrentPos))
+          return std::move(Err);
+      }
+      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.
if (ModuleLevel && IsImporting && MetadataList.empty() &&
     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());
 
+      // Now that we have built the index, load the global decl attachments
+      // that were deferred during that process. This avoids creating
+      // temporaries.
+      SuccessOrErr = loadGlobalDeclAttachments();
+      if (!SuccessOrErr)
+        return SuccessOrErr.takeError();
+      assert(SuccessOrErr.get());
+
       // Reading the named metadata created forward references and/or
       // placeholders, that we flush here.
       resolveForwardRefsAndPlaceholders(Placeholders);
       upgradeDebugInfo();
       // Return at the beginning of the block, since it is easy to skip it
       // entirely from there.
       Stream.ReadBlockEnd(); // Pop the abbrev block context.
       if (Error Err = IndexCursor.JumpToBit(EntryPos))

 
 Error MetadataLoader::MetadataLoaderImpl::parseGlobalObjectAttachment(
     GlobalObject &GO, ArrayRef<uint64_t> Record) {
   assert(Record.size() % 2 == 0);
   for (unsigned I = 0, E = Record.size(); I != E; I += 2) {
     auto K = MDKindMap.find(Record[I]);
     if (K == MDKindMap.end())
       return error("Invalid ID");
-    MDNode *MD = MetadataList.getMDNodeFwdRefOrNull(Record[I + 1]);
+    MDNode *MD =
+        dyn_cast_or_null<MDNode>(getMetadataFwdRefOrLoad(Record[I + 1]));
     if (!MD)
       return error("Invalid metadata attachment: expect fwd ref to MDNode");
     GO.addMetadata(K->second, *MD);
   }
   return Error::success();
 }
 
 /// Parse metadata attachments.

llvm/test/ThinLTO/X86/devirt2.ll

 ; generate summary information needed for index-based WPD.
 ; RUN: llvm-modextract -b -n=0 %t2.o -o %t2.o.0
 ; RUN: llvm-dis -o - %t2.o.0 | FileCheck %s --check-prefix=ENABLESPLITFLAG --implicit-check-not=vTableFuncs --implicit-check-not=typeidCompatibleVTable
 ; RUN: llvm-modextract -b -n=1 %t2.o -o %t2.o.1
 ; RUN: llvm-dis -o - %t2.o.1 | FileCheck %s --check-prefix=ENABLESPLITFLAG --implicit-check-not=vTableFuncs --implicit-check-not=typeidCompatibleVTable
 ; ENABLESPLITFLAG: !{i32 1, !"EnableSplitLTOUnit", i32 1}
 
 ; Generate unsplit module with summary for ThinLTO index-based WPD.
-; RUN: opt -thinlto-bc -o %t3.o %s
-; RUN: opt -thinlto-bc -o %t4.o %p/Inputs/devirt2.ll
+; Force generation of the bitcode index so that we also test lazy metadata
+; loader handling of the type metadata.
+; RUN: opt -bitcode-mdindex-threshold=0 -thinlto-bc -o %t3.o %s
+; RUN: opt -bitcode-mdindex-threshold=0 -thinlto-bc -o %t4.o %p/Inputs/devirt2.ll
 
 ; Check that we don't have module flag when splitting not enabled for ThinLTO,
 ; and that we generate summary information needed for index-based WPD.
 ; RUN: llvm-dis -o - %t4.o | FileCheck %s --check-prefix=NOENABLESPLITFLAG
 ; NOENABLESPLITFLAG-DAG: !{i32 1, !"EnableSplitLTOUnit", i32 0}
 ; NOENABLESPLITFLAG-DAG: [[An:\^[0-9]+]] = gv: (name: "_ZN1A1nEi"
 ; NOENABLESPLITFLAG-DAG: [[Bf:\^[0-9]+]] = gv: (name: "_ZN1B1fEi"
 ; NOENABLESPLITFLAG-DAG: [[Cf:\^[0-9]+]] = gv: (name: "_ZN1C1fEi"