[BitcodeReader] Fix O(N^2) in placeholder replacement algorithm.

[efriedma]

Summary

The key here is to make sure we never RAUW a constant that's used by other constants. Since Constants form a DAG, this is just a matter of constructing the replacements and RAUW'ing the constants in the right order. Most of the complexity here is computing that order.

One thing that's a bit unfortunate about this algorithm is that it creates a couple DenseMaps which have only entry per replaced constant. So it uses significantly more memory than the previous algorithm. It might be possible to optimize this to some extent, but hopefully this doesn't matter in practice.

Fixes https://bugs.llvm.org/show_bug.cgi?id=47395 .

[resistor]

If the maps are typically small, would SmallDenseMap make sense?

[efriedma]

resolveConstantForwardRefs() is called once per bitcode file. For bitcode file of any significant complexity, the list of placeholders probably won't fit in a SmallDenseMap.

[efriedma]

Ping

[mehdi_amini]

I tried to review, but I'm too unfamiliar with this code at this point and won't have enough time to page-in the existing one and review the new one.

llvm/lib/Bitcode/Reader/ValueList.cpp

157

Nit: can you run git clang-format?

159

I think this lambda could be documented, or alternatively maybe renamed: alone it does not "count" but "increment" users' count.

214

auto->Value ?

[spatel]

Adding more potential reviewers. I'm not sure if anyone is familiar with this code, but others definitely have a better understanding of the data structures than me.

Also note that we got confirmation in PR47395 that the compile time went from 12 hours (!) back to 3 minutes, so this patch does solve the motivating bug.

[dblaikie]

OK, having a go at reviewing this (though likely you're about in the best position to assess this, Eli)...

So, the original code:

For each constant placeholder
  While there are still uses of the constant placeholder
    if the use isn't a constant, just change it to use the real value
    if the use is a constant
      walk all the operands of the user & update any that refer to placeholders
      make a new constant with the new list of operands
  replace the remaining non-constant users

You mention:

The key here is to make sure we never RAUW a constant that's used by other constants.

Though my naive reading of the existing code looks to me like it checks all its users and fixes them up before RAUW the remaining... ah, but that's not done recursively, so if a user itself has users that are constants - then there's a RAUW of a constant that's used by other constants. OK

Right, so the goal is to start with constants that aren't used by any other constants - I guess that means starting with something in ResolveConstants that has a value 0/not present in NumRewrittenOperands?

So we start by walking all the constant users of placeholder constants - and then also (line 174) the constant users of any of those (so I guess they're not all using placeholders - some are already resolved?)

Then nah, this is where I start to not be able to keep it all in my head/clearly. Could you perhaps explain in more detail to summarize the algorithm?

[efriedma]

There are basically three steps here:

1. We find all the relevant constants, and do a topological sort of them. Think of all the Constants as a DAG, where the edges go from a Constant to its uses. We treat placeholders as if they're uses of the corresponding constant in the ResolveConstants mapping. The roots are placeholders that don't refer to another placeholder. The result is saved in RewrittenConstants.
2. We go through the constants in order, and compute the rewritten version of every constant.
3. We go through the constants in reverse order, and RAUW every constant to its rewritten version.

The key here is ensuring the RAUW operations in step 3 never rewrite a Constant.

Looking at this again, it's maybe a little trickier to follow than it needs to be. The topological sort and the forward iteration step are mixed together; I could separate them. And maybe there's a better way to write the topological sort.

(so I guess they're not all using placeholders - some are already resolved?)

We only use a placeholder in constants that refer to another constant that hasn't been constructed yet. (If the bitcode writer wrote out all the constants in sorted order, we wouldn't need placeholders at all.)

[efriedma]

Completely separate the topological sort from the other steps. Add more comments/asserts to clarify how the algorithm works.

[efriedma]

Ping

[dblaikie]

Think I've got it - sorry for the delay & hopefully my rambling notes (if you want to have a read) at least provide a little check on the work - but I hardly feel the most qualified in this area.

llvm/lib/Bitcode/Reader/ValueList.cpp

152–153

I'd consider moving this down closer to the scope it's used, perhaps?

159–172

So if I understand this correctly - we might have, say 3 forward references ("ResolveConstants" - would that variable be more aptly named "UnresolvedConstants"? (just clarifying my own understanding, not a request/need to change if it doesn't seem generally worthwhile)) - but there might be constants without any forward references but they do have references to constants that have forward references - and so on. So, say ResolveConstants is {A, B, C} but some already resolved constants are {X, Y, Z}. And the operands are Y(X(A(C),B)) and Z references nothing (or at least no other constants (that aren't also GlobalValues, I guess)).

Placeholders (in ResolveConstant) don't actually use anything, I guess? (because they're placeholders/haven't been filled out yet)

So the 169 for loop goes through {A, B, C}:
Constant: A -> NumRewrittenOperands = {X:1}, RewriteWorklist = {X}
Constant: B -> NumRewrittenOperands = {X:2}, RewriteWorklist = {X}
Constant: C -> NumRewrittenOperands = {X:2}, RewriteWorklist = {X}

Then the loop on 171:
Constant: X -> NumRewrittenOperands = {X:2, Y:1}, RewriteWorklist = {Y}
Constant: Y -> NumRewrittenOperands = {X:2, Y:1}, RewriteWorklist = {}

Looking further down, I see that we might find resolved constants during this walk. So, for instance when visiting the uses of C we might find A', the resolved version of A? I would've thought if that was the case then we'd end up double-counting X's NumRewrittenOperands, for instance - but I guess X is only using either A or A', not both. So A' gets built, but it isn't inserted yet... which is the whole point of all this, right right. So if we did come across A' like that, we'd put A' in NumRewrittenOperands:1, A' would go in the RewriteWorklist, but have no effect since no one is using it - they're all using A.

So let's say we endu p with NumRewrittenOperands = {X:2, Y:1, A':1}

196–206

So here we visit {A, B, C} again. Find their mapped equivalents, and build an ordered RewriteWorklist... (RewriteWorklist is the same variable, but is being used independently here, since it was cleared out by the loop at 171 above)

A, A' -> ConstToPlaceholders = {A':[A]}, NumRewrittenOperands = {X:2, Y:1, A':1}, RewriteWorklist = {}, ConstantMapping = {}
B, B' -> ConstToPlaceholders = {A':[A]}, NumRewrittenOperands = {X:1, Y:1, A':1}, RewriteWorklist = {}, ConstantMapping = {B:B'}, RewrittenConstants = {B}
C, C' -> ConstToPlaceholders = {A':[A]}, NumRewrittenOperands = {X:1, Y:1}, RewriteWorklist = {A'}, ConstantMapping = {B:B', C:C'}, RewrittenConstants = {B, C}

210–218

Constant: A', RewriteWorklist = {}

AddConstantToOrder(A') -> NumRewrittenOperands = {X:1, Y:1}, RewrittenConstants = {B, C, A'}
  AddConstantToOrder(A) -> NumRewrittenOperands = {Y:1}, RewrittenConstants = {B, C, A', A}, RewriteWorklist = {X}

Constant: X, RewriteWorklist = {}

AddConstantToOrder(X) -> NumRewrittenOperands = {Y:1}, RewrittenConstants = {B, C, A', A, X}, RewriteWorklist = {Y}

Constant: Y, RewriteWorklist = {}

AddConstantToOrder(Y) -> NumRewrittenOperands = {}, RewrittenConstants = {B, C, A', A, X, Y}, RewriteWorklist = {}

222–254

RewrittenConstants = {B, C, A', A, X, Y}, ConstantMapping = {B:B', C:C'},

Loop:
Constant: B -> Placeholder skip
Constant: C -> Placeholder skip
Constant: A' -> ConstantMapping = {B:B', C:C', A':A'', A:A''}
Constant: A -> Placeholder skip
Constant: X -> ConstantMapping = {B:B', C:C', A':A'', X, X'} (mapping A and B to make X'(A'', B'))
Constant: Y -> ConstantMapping = {B:B', C:C', A':A'', X, X', Y:Y'} (mapping X to make Y'(X'(A'', B')))

256–257

Ah, this is because if we RAUW, say, C first - it'd replace the use of C in A', but we're going to throw away A' (line 266) anyway in favor of A''?

llvm/lib/Bitcode/Reader/ValueList.cpp

 //===- ValueList.cpp - Internal BitcodeReader implementation --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ValueList.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"

 
 /// Once all constants are read, this method bulk resolves any forward
 /// references.  The idea behind this is that we sometimes get constants (such
 /// as large arrays) which reference *many* forward ref constants.  Replacing
 /// each of these causes a lot of thrashing when building/reuniquing the
 /// constant.  Instead of doing this, we look at all the uses and rewrite all
 /// the place holders at once for any constant that uses a placeholder.
 void BitcodeReaderValueList::resolveConstantForwardRefs() {
-  // Sort the values by-pointer so that they are efficient to look up with a
-  // binary search.
-  llvm::sort(ResolveConstants);
+  DenseMap<Constant *, Constant *> ConstantMapping;
+  ConstantMapping.reserve(ResolveConstants.size());
 
+  // For each constant, compute how many of its operands needs to be rewritten.
+  // We use these counts to drive the topological sort.
+  DenseMap<Constant *, unsigned> NumRewrittenOperands;
+  SmallVector<Constant *, 64> RewriteWorklist;
+  auto CountOperandsOfUsers = [&](Constant *C) {
+    for (Use &ConstantUse : C->uses()) {
+      auto *ConstantUser = dyn_cast<Constant>(ConstantUse.getUser());
+      if (ConstantUser && !isa<GlobalValue>(ConstantUser)) {
+        unsigned NumOperandsFound = ++NumRewrittenOperands[ConstantUser];
+        if (NumOperandsFound == 1)
+          RewriteWorklist.push_back(ConstantUser);
+      }
+    }
+  };
+  for (auto ResolveConstant : ResolveConstants)
+    CountOperandsOfUsers(ResolveConstant.first);
+  while (!RewriteWorklist.empty())
+    CountOperandsOfUsers(RewriteWorklist.pop_back_val());
+
+  // List of constants in sorted order.
+  SmallVector<Constant *, 64> RewrittenConstants;
+  // Temporary vector for operands.
   SmallVector<Constant *, 64> NewOps;
-
-  while (!ResolveConstants.empty()) {
-    Value *RealVal = operator[](ResolveConstants.back().second);
-    Constant *Placeholder = ResolveConstants.back().first;
-    ResolveConstants.pop_back();
-
-    // Loop over all users of the placeholder, updating them to reference the
-    // new value.  If they reference more than one placeholder, update them all
-    // at once.
-    while (!Placeholder->use_empty()) {
-      auto UI = Placeholder->user_begin();
-      User *U = *UI;
-
-      // If the using object isn't uniqued, just update the operands.  This
-      // handles instructions and initializers for global variables.
-      if (!isa<Constant>(U) || isa<GlobalValue>(U)) {
-        UI.getUse().set(RealVal);
-        continue;
-      }
+  // Mapping from constants which will be rewritten to placeholders for
+  // those constants.
+  DenseMap<Constant *, TinyPtrVector<Constant *>> ConstToPlaceholders;
+
+  // Add a constant to the sorted order, and add its uses to the worklist.
+  auto AddConstantToOrder = [&](Constant *C) {
+    RewrittenConstants.push_back(C);
+    for (auto &ConstantUse : C->uses()) {
+      auto *ConstantUser = dyn_cast<Constant>(ConstantUse.getUser());
+      if (ConstantUser && !isa<GlobalValue>(ConstantUser) &&
+          --NumRewrittenOperands[ConstantUser] == 0) {
+        RewriteWorklist.push_back(ConstantUser);
+      }
+  };
 
-      // Otherwise, we have a constant that uses the placeholder.  Replace that
-      // constant with a new constant that has *all* placeholder uses updated.
-      Constant *UserC = cast<Constant>(U);
-      for (User::op_iterator I = UserC->op_begin(), E = UserC->op_end(); I != E;
-           ++I) {
-        Value *NewOp;
-        if (!isa<ConstantPlaceHolder>(*I)) {
-          // Not a placeholder reference.
-          NewOp = *I;
-        } else if (*I == Placeholder) {
-          // Common case is that it just references this one placeholder.
-          NewOp = RealVal;
-        } else {
-          // Otherwise, look up the placeholder in ResolveConstants.
-          ResolveConstantsTy::iterator It = llvm::lower_bound(
-              ResolveConstants,
-              std::pair<Constant *, unsigned>(cast<Constant>(*I), 0));
-          assert(It != ResolveConstants.end() && It->first == *I);
-          NewOp = operator[](It->second);
-        }
+  // Initialize the worklist with the roots for the topological sort:
+  // placeholders which don't depend on another placeholder.
+  for (auto ResolveConstant : ResolveConstants) {
+    Constant *MappedC = cast<Constant>(operator[](ResolveConstant.second));
+    if (NumRewrittenOperands.lookup(MappedC)) {
+      // We have a placeholder that resolves to a constant that needs to
+      // be rewritten; note it so we can process it later.
+      ConstToPlaceholders[MappedC].push_back(ResolveConstant.first);
+      AddConstantToOrder(ResolveConstant.first);
+      ConstantMapping[ResolveConstant.first] = MappedC;
+    }
+  ResolveConstants.clear();
 
-        NewOps.push_back(cast<Constant>(NewOp));
-      }
+  // Perform the topological sort. This is Kahn's algorithm.
+  while (!RewriteWorklist.empty()) {
+    Constant *C = RewriteWorklist.pop_back_val();
+
+    // Add this constant, and placeholders that refer to it, to the sorted
+    // order. Add its uses to the worklist.
+    AddConstantToOrder(C);
+    for (Constant *Placeholder : ConstToPlaceholders.lookup(C))
+      AddConstantToOrder(Placeholder);
+  }
+
+  // Compute the rest of the ConstantMapping map. The topological sort ensures
+  // we can look up the rewritten operands in ConstantMapping.
+  for (Constant *C : RewrittenConstants) {
+    if (isa<ConstantPlaceHolder>(C)) {
+      assert(ConstantMapping.lookup(C) && "Should already be mapped");
+      continue;
+    }
+
+    // Remap the constant's operands.
+    for (auto &Operand : C->operands()) {
+      Constant *OperandVal = cast<Constant>(Operand);
+      Constant *NewOp = ConstantMapping.lookup(OperandVal);
+      NewOps.push_back(NewOp ? NewOp : OperandVal);
 
-      // Make the new constant.
-      Constant *NewC;
-      if (ConstantArray *UserCA = dyn_cast<ConstantArray>(UserC)) {
-        NewC = ConstantArray::get(UserCA->getType(), NewOps);
-      } else if (ConstantStruct *UserCS = dyn_cast<ConstantStruct>(UserC)) {
-        NewC = ConstantStruct::get(UserCS->getType(), NewOps);
-      } else if (isa<ConstantVector>(UserC)) {
-        NewC = ConstantVector::get(NewOps);
-      } else {
-        assert(isa<ConstantExpr>(UserC) && "Must be a ConstantExpr.");
-        NewC = cast<ConstantExpr>(UserC)->getWithOperands(NewOps);
-      }
-
-      UserC->replaceAllUsesWith(NewC);
-      UserC->destroyConstant();
-      NewOps.clear();
-    }
+    if (ConstantArray *CA = dyn_cast<ConstantArray>(C)) {
+      NewC = ConstantArray::get(CA->getType(), NewOps);
+    } else if (ConstantStruct *CS = dyn_cast<ConstantStruct>(C)) {
+      NewC = ConstantStruct::get(CS->getType(), NewOps);
+    } else if (isa<ConstantVector>(C)) {
+      assert(isa<ConstantExpr>(C) && "Must be a ConstantExpr.");
+      NewC = cast<ConstantExpr>(C)->getWithOperands(NewOps);
 
-    // Update all ValueHandles, they should be the only users at this point.
-    Placeholder->replaceAllUsesWith(RealVal);
-    delete cast<ConstantPlaceHolder>(Placeholder);
+    // Map this constant, and any placeholders that correspond to this
+    // constant, to the new value.
+    ConstantMapping[C] = NewC;
+    for (Constant *Placeholder : ConstToPlaceholders.lookup(C))
+      ConstantMapping[Placeholder] = NewC;
+  }
+
+  // Rewrite the uses of each constant.  The topological sort ensures we never
+  // RAUW a constant that's used by other constants.
+  for (Constant *C : reverse(RewrittenConstants)) {
+    assert(all_of(C->users(),
+                  [](User *U) -> bool {
+                    return !isa<Constant>(U) || isa<GlobalValue>(U);
+                  }) &&
+           "Rewriting constants in wrong order");
+    C->replaceAllUsesWith(ConstantMapping.lookup(C));
+    if (auto *Placeholder = dyn_cast<ConstantPlaceHolder>(C))
+      delete Placeholder;
+    else
+      C->destroyConstant();
   }
 }