Bitcode: Try to emit metadata in function blocks

Whenever metadata is only referenced by a single function, emit the
metadata just in that function block.  This should improve lazy-loading
by reducing the amount of metadata in the global block.

For now, this should catch all DILocations, and anything else that
happens to be referenced only by a single function.

It's also a first step toward a couple of possible future directions
(which this commit does *not* implement):

 1. Some debug info metadata is only referenced from compile units and
    individual functions.  If we can drop the link from the compile
    unit, this optimization will get more powerful.

 2. Any uniqued metadata that isn't referenced globally can in theory be
    emitted in every function block that references it (trading off
    bitcode size and full-parse time vs. lazy-load time).

Note: this assumes the new BitcodeReader error checking from r265223.
The metadata stored in function blocks gets purged after parsing each
function, which means unresolved forward references will get lost.
Since all the global metadata should have already been resolved by the
time we get to the function metadata blocks we just need to check for
that case.  (If for some reason we need to handle bitcode that fails the
checks in r265223, the fix is to store about-to-be-dropped unresolved
nodes in MetadataList::shrinkTo until they can be handled succesfully by
a future call to MetadataList::tryToResolveCycles.)

llvm-svn: 265226

1 files changed, 171 insertions, 30 deletions

diff --git a/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp b/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
index 5c5f05ca8f9..91b523f05bc 100644
--- a/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
+++ b/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
@@ -336,7 +336,7 @@ ValueEnumerator::ValueEnumerator(const Module &M,
     // Enumerate metadata attached to this function.
     F.getAllMetadata(MDs);
     for (const auto &I : MDs)
-      EnumerateMetadata(I.second);
+      EnumerateMetadata(&F, I.second);
 
     for (const BasicBlock &BB : F)
       for (const Instruction &I : BB) {
@@ -351,7 +351,7 @@ ValueEnumerator::ValueEnumerator(const Module &M,
           if (isa<LocalAsMetadata>(MD->getMetadata()))
             continue;
 
-          EnumerateMetadata(MD->getMetadata());
+          EnumerateMetadata(&F, MD->getMetadata());
         }
         EnumerateType(I.getType());
         if (const CallInst *CI = dyn_cast<CallInst>(&I))
@@ -363,12 +363,12 @@ ValueEnumerator::ValueEnumerator(const Module &M,
         MDs.clear();
         I.getAllMetadataOtherThanDebugLoc(MDs);
         for (unsigned i = 0, e = MDs.size(); i != e; ++i)
-          EnumerateMetadata(MDs[i].second);
+          EnumerateMetadata(&F, MDs[i].second);
 
         // Don't enumerate the location directly -- it has a special record
         // type -- but enumerate its operands.
         if (DILocation *L = I.getDebugLoc())
-          EnumerateMDNodeOperands(L);
+          EnumerateMDNodeOperands(&F, L);
       }
   }
 
@@ -447,8 +447,10 @@ void ValueEnumerator::print(raw_ostream &OS, const MetadataMapType &Map,
   OS << "Size: " << Map.size() << "\n";
   for (auto I = Map.begin(), E = Map.end(); I != E; ++I) {
     const Metadata *MD = I->first;
-    OS << "Metadata: slot = " << I->second << "\n";
+    OS << "Metadata: slot = " << I->second.ID << "\n";
+    OS << "Metadata: function = " << I->second.F << "\n";
     MD->print(OS);
+    OS << "\n";
   }
 }
 
@@ -500,22 +502,87 @@ void ValueEnumerator::EnumerateNamedMetadata(const Module &M) {
 
 void ValueEnumerator::EnumerateNamedMDNode(const NamedMDNode *MD) {
   for (unsigned i = 0, e = MD->getNumOperands(); i != e; ++i)
-    EnumerateMetadata(MD->getOperand(i));
+    EnumerateMetadata(nullptr, MD->getOperand(i));
+}
+
+unsigned ValueEnumerator::getMetadataFunctionID(const Function *F) const {
+  return F ? getValueID(F) + 1 : 0;
+}
+
+void ValueEnumerator::EnumerateMDNodeOperands(const Function *F,
+                                              const MDNode *N) {
+  EnumerateMDNodeOperands(getMetadataFunctionID(F), N);
+}
+
+void ValueEnumerator::EnumerateMetadata(const Function *F, const Metadata *MD) {
+  EnumerateMetadata(getMetadataFunctionID(F), MD);
+}
+
+void ValueEnumerator::EnumerateFunctionLocalMetadata(
+    const Function &F, const LocalAsMetadata *Local) {
+  EnumerateFunctionLocalMetadata(getMetadataFunctionID(&F), Local);
 }
 
 /// EnumerateMDNodeOperands - Enumerate all non-function-local values
 /// and types referenced by the given MDNode.
-void ValueEnumerator::EnumerateMDNodeOperands(const MDNode *N) {
+void ValueEnumerator::EnumerateMDNodeOperands(unsigned F, const MDNode *N) {
   for (unsigned i = 0, e = N->getNumOperands(); i != e; ++i) {
     Metadata *MD = N->getOperand(i);
     if (!MD)
       continue;
     assert(!isa<LocalAsMetadata>(MD) && "MDNodes cannot be function-local");
-    EnumerateMetadata(MD);
+    EnumerateMetadata(F, MD);
+  }
+}
+
+bool ValueEnumerator::insertMetadata(unsigned F, const Metadata *MD) {
+  auto Insertion = MetadataMap.insert(std::make_pair(MD, MDIndex(F)));
+  if (Insertion.second)
+    return true;
+
+  // Check whether F is a different function.
+  MDIndex &Entry = Insertion.first->second;
+  if (!Entry.hasDifferentFunction(F))
+    return false;
+
+  // Since MD was tagged from a different function entry point then it must
+  // already have an ID.
+  assert(Entry.ID && "Expected metadata to already be indexed");
+  Entry.F = 0;
+
+  // Drop the function from transitive operands.
+  if (auto *N = dyn_cast<MDNode>(MD))
+    dropFunctionFromOps(*N);
+
+  return false;
+}
+
+void ValueEnumerator::dropFunctionFromOps(const MDNode &N) {
+  SmallVector<const MDNode *, 64> WorkList;
+  WorkList.push_back(&N);
+  while (!WorkList.empty()) {
+    for (const Metadata *Op : WorkList.pop_back_val()->operands()) {
+      if (!Op)
+        continue;
+
+      // All transitive operands of N should already have IDs.  This should be
+      // a second traversal.
+      auto &Entry = MetadataMap[Op];
+      assert(Entry.ID && "Expected metadata to already be indexed");
+
+      // Nothing to do if this operand isn't tagged.
+      if (!Entry.F)
+        continue;
+
+      // Drop the tag, and if it's a node (with potential operands), queue it.
+      Entry.F = 0;
+      if (auto *OpN = dyn_cast<MDNode>(Op))
+        WorkList.push_back(OpN);
+    }
   }
 }
 
-void ValueEnumerator::EnumerateMetadata(const Metadata *MD) {
+void ValueEnumerator::EnumerateMetadata(unsigned F, const Metadata *MD) {
   assert(
       (isa<MDNode>(MD) || isa<MDString>(MD) || isa<ConstantAsMetadata>(MD)) &&
       "Invalid metadata kind");
@@ -524,49 +591,120 @@ void ValueEnumerator::EnumerateMetadata(const Metadata *MD) {
   // EnumerateMDNodeOperands() from re-visiting MD in a cyclic graph.
   //
   // Return early if there's already an ID.
-  if (!MetadataMap.insert(std::make_pair(MD, 0)).second)
+  if (!insertMetadata(F, MD))
     return;
 
   // Visit operands first to minimize RAUW.
   if (auto *N = dyn_cast<MDNode>(MD))
-    EnumerateMDNodeOperands(N);
+    EnumerateMDNodeOperands(F, N);
   else if (auto *C = dyn_cast<ConstantAsMetadata>(MD))
     EnumerateValue(C->getValue());
-  else
-    ++NumMDStrings;
 
-  // Replace the dummy ID inserted above with the correct one.  MetadataMap may
-  // have changed by inserting operands, so we need a fresh lookup here.
+  // Save the metadata.
   MDs.push_back(MD);
-  MetadataMap[MD] = MDs.size();
+  MetadataMap[MD].ID = MDs.size();
 }
 
 /// EnumerateFunctionLocalMetadataa - Incorporate function-local metadata
 /// information reachable from the metadata.
 void ValueEnumerator::EnumerateFunctionLocalMetadata(
-    const LocalAsMetadata *Local) {
+    unsigned F, const LocalAsMetadata *Local) {
+  assert(F && "Expected a function");
+
   // Check to see if it's already in!
-  unsigned &MetadataID = MetadataMap[Local];
-  if (MetadataID)
+  MDIndex &Index = MetadataMap[Local];
+  if (Index.ID) {
+    assert(Index.F == F && "Expected the same function");
     return;
+  }
 
   MDs.push_back(Local);
-  MetadataID = MDs.size();
+  Index.F = F;
+  Index.ID = MDs.size();
 
   EnumerateValue(Local->getValue());
 }
 
 void ValueEnumerator::organizeMetadata() {
-  if (!NumMDStrings)
+  assert(MetadataMap.size() == MDs.size() &&
+         "Metadata map and vector out of sync");
+
+  if (MDs.empty())
+    return;
+
+  // Copy out the index information from MetadataMap in order to choose a new
+  // order.
+  SmallVector<MDIndex, 64> Order;
+  Order.reserve(MetadataMap.size());
+  for (const Metadata *MD : MDs)
+    Order.push_back(MetadataMap.lookup(MD));
+
+  // Partition:
+  //   - by function, then
+  //   - by isa<MDString>
+  // and then sort by the original/current ID.  Since the IDs are guaranteed to
+  // be unique, the result of std::sort will be deterministic.  There's no need
+  // for std::stable_sort.
+  std::sort(Order.begin(), Order.end(), [this](MDIndex LHS, MDIndex RHS) {
+    return std::make_tuple(LHS.F, !isa<MDString>(LHS.get(MDs)), LHS.ID) <
+           std::make_tuple(RHS.F, !isa<MDString>(RHS.get(MDs)), RHS.ID);
+  });
+
+  // Return early if nothing is moving to functions and there are no strings.
+  if (!Order.back().F && !isa<MDString>(Order.front().get(MDs)))
+    return;
+
+  // Rebuild MDs, index the metadata ranges for each function in FunctionMDs,
+  // and fix up MetadataMap.
+  std::vector<const Metadata *> OldMDs = std::move(MDs);
+  MDs.reserve(OldMDs.size());
+  for (unsigned I = 0, E = Order.size(); I != E && !Order[I].F; ++I) {
+    auto *MD = Order[I].get(OldMDs);
+    MDs.push_back(MD);
+    MetadataMap[MD].ID = I + 1;
+    if (isa<MDString>(MD))
+      ++NumMDStrings;
+  }
+
+  // Return early if there's nothing for the functions.
+  if (MDs.size() == Order.size())
     return;
 
-  // Put the strings first.
-  std::stable_partition(MDs.begin(), MDs.end(),
-                        [](const Metadata *MD) { return isa<MDString>(MD); });
+  // Build the function metadata ranges.
+  MDRange R;
+  FunctionMDs.reserve(OldMDs.size());
+  unsigned PrevF = 0;
+  for (unsigned I = MDs.size(), E = Order.size(), ID = MDs.size(); I != E;
+       ++I) {
+    unsigned F = Order[I].F;
+    if (!PrevF) {
+      PrevF = F;
+    } else if (PrevF != F) {
+      R.Last = FunctionMDs.size();
+      std::swap(R, FunctionMDInfo[PrevF]);
+      R.First = FunctionMDs.size();
+
+      ID = MDs.size();
+      PrevF = F;
+    }
 
-  // Renumber.
-  for (unsigned I = 0, E = MDs.size(); I != E; ++I)
-    MetadataMap[MDs[I]] = I + 1;
+    auto *MD = Order[I].get(OldMDs);
+    FunctionMDs.push_back(MD);
+    MetadataMap[MD].ID = ++ID;
+    if (isa<MDString>(MD))
+      ++R.NumStrings;
+  }
+  R.Last = FunctionMDs.size();
+  FunctionMDInfo[PrevF] = R;
+}
+
+void ValueEnumerator::incorporateFunctionMetadata(const Function &F) {
+  NumModuleMDs = MDs.size();
+
+  auto R = FunctionMDInfo.lookup(getValueID(&F) + 1);
+  NumMDStrings = R.NumStrings;
+  MDs.insert(MDs.end(), FunctionMDs.begin() + R.First,
+             FunctionMDs.begin() + R.Last);
 }
 
 void ValueEnumerator::EnumerateValue(const Value *V) {
@@ -708,8 +846,10 @@ void ValueEnumerator::EnumerateAttributes(AttributeSet PAL) {
 void ValueEnumerator::incorporateFunction(const Function &F) {
   InstructionCount = 0;
   NumModuleValues = Values.size();
-  NumModuleMDs = MDs.size();
-  NumMDStrings = 0;
+
+  // Add global metadata to the function block.  This doesn't include
+  // LocalAsMetadata.
+  incorporateFunctionMetadata(F);
 
   // Adding function arguments to the value table.
   for (const auto &I : F.args())
@@ -755,7 +895,7 @@ void ValueEnumerator::incorporateFunction(const Function &F) {
 
   // Add all of the function-local metadata.
   for (unsigned i = 0, e = FnLocalMDVector.size(); i != e; ++i)
-    EnumerateFunctionLocalMetadata(FnLocalMDVector[i]);
+    EnumerateFunctionLocalMetadata(F, FnLocalMDVector[i]);
 }
 
 void ValueEnumerator::purgeFunction() {
@@ -770,6 +910,7 @@ void ValueEnumerator::purgeFunction() {
   Values.resize(NumModuleValues);
   MDs.resize(NumModuleMDs);
   BasicBlocks.clear();
+  NumMDStrings = 0;
 }
 
 static void IncorporateFunctionInfoGlobalBBIDs(const Function *F,