[MergeFuncs] Efficiently defer functions on merge

[• jrkoenig]

Summary

This patch introduces a side table in Merge Functions to
efficiently remove functions from the function set when functions
they refer to are merged. Previously these functions would need to
be compared lg(N) times to find the appropriate FunctionNode in the
tree to defer. With the recent determinism changes, this comparison
is more expensive. In addition, the removal function would not always
actually remove the function from the set (i.e. after remove(F),
there would sometimes still be a node in the tree which contains F).

With these changes, these functions are properly deferred, and so more
functions can be merged. In addition, when there are many merged
functions (and thus more deferred functions), there is a speedup:

chromium: 48678 merged -> 49380 merged; 6.58s -> 5.49s
libxul.so: 41004 merged -> 41030 merged; 8.02s -> 6.94s
mysqld: 1607 merged -> 1607 merged (same); 0.215s -> 0.212s (probably noise)

[jfb]

I'm not super comfortable with keeping iterators around: iterator invalidation is a great gotcha in C++. Is there a way to avoid keeping the iterators, without changing complexity (potentially by doing a replace on ->second)? e.g. set insertion doesn't invalidate, but set deletion does, and you have both so I can't naively convince myself that your code is valid.

lib/Transforms/IPO/MergeFunctions.cpp

1731

Move this assert up one line, after find.

[• jrkoenig]

Clarified invariant linking FnTree and FNodesInTree, use more descriptive
names, and add more comments.

[• jrkoenig]

According to this C++ standard draft/prerelease/whatever, Section 23.2.4, paragraph 9 claims that only the elements erased are invalidated. This is subtle but correct usage. Sounds like I need to add some comments.

The alternative is to iterate over the tree and remove any copy of F that is found, which is quite expensive. The current implementation (which uses the ordering to find the right node) sometimes missed functions, which most likely occurs due to functions changed before they are actually removed from the tree, changing the implicit invariant of the tree that it's elements remain in a stable order.

[jfb]

The current implementation (which uses the ordering to find the right node) sometimes missed functions, which most likely occurs due to functions changed before they are actually removed from the tree, changing the implicit invariant of the tree that it's elements remain in a stable order.

Can you clarify this? Is that a new thing? It at least seems like something that should be in a FIXME.

lib/Transforms/IPO/MergeFunctions.cpp

1734

It's a more common style to do the following instead of having a comment:

assert(foo && "I expected foo because bar");

[jfb]

This comment was removed by jfb.

[• jrkoenig]

Move comments into asserts.

[• jrkoenig]

In r245139, before hashing and the other recent changes, remove() was buggy as well. The problem manifests in all three LTO builds of chromium, mysql, and libxul.

[jfb]

lgtm

lib/Transforms/IPO/MergeFunctions.cpp

 /// by considering all pointer types to be equivalent. Once identified,
 /// MergeFunctions will fold them by replacing a call to one to a call to a
 /// bitcast of the other.
 ///
 class MergeFunctions : public ModulePass {
 public:
   static char ID;
   MergeFunctions()
-    : ModulePass(ID), FnTree(FunctionNodeCmp(&GlobalNumbers)),
+    : ModulePass(ID), FnTree(FunctionNodeCmp(&GlobalNumbers)), FNodesInTree(),
       HasGlobalAliases(false) {
     initializeMergeFunctionsPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnModule(Module &M) override;
 
 private:
   // The function comparison operator is provided here so that FunctionNodes do
private:
   /// Replace G with a simple tail call to bitcast(F). Also replace direct uses
   /// of G with bitcast(F). Deletes G.
   void writeThunk(Function *F, Function *G);
 
   /// Replace G with an alias to F. Deletes G.
   void writeAlias(Function *F, Function *G);
 
   /// Replace function F with function G in the function tree.
-  void replaceFunctionInTree(FnTreeType::iterator &IterToF, Function *G);
+  void replaceFunctionInTree(const FunctionNode &FN, Function *G);
 
   /// The set of all distinct functions. Use the insert() and remove() methods
-  /// to modify it.
+  /// to modify it. The map allows efficient lookup and deferring of Functions.
   FnTreeType FnTree;
+  // Map functions to the iterators of the FunctionNode which contains them
+  // in the FnTree. This must be updated carefully whenever the FnTree is
+  // modified, i.e. in insert(), remove(), and replaceFunctionInTree(), to avoid
+  // dangling iterators into FnTree. The invariant that preserves this is that
+  // there is exactly one mapping F -> FN for each FunctionNode FN in FnTree.
+  ValueMap<Function*, FnTreeType::iterator> FNodesInTree;
 
   /// Whether or not the target supports global aliases.
   bool HasGlobalAliases;
 };
 
 }  // end anonymous namespace
 
 char MergeFunctions::ID = 0;
if (F->mayBeOverridden()) {
     ++NumDoubleWeak;
   } else {
     writeThunkOrAlias(F, G);
   }
 
   ++NumFunctionsMerged;
 }
 
-/// Replace function F for function G in the map.
-void MergeFunctions::replaceFunctionInTree(FnTreeType::iterator &IterToF,
+/// Replace function F by function G.
+void MergeFunctions::replaceFunctionInTree(const FunctionNode &FN,
                                            Function *G) {
-  Function *F = IterToF->getFunc();
-
-  // A total order is already guaranteed otherwise because we process strong
-  // functions before weak functions.
-  assert(((F->mayBeOverridden() && G->mayBeOverridden()) ||
-          (!F->mayBeOverridden() && !G->mayBeOverridden())) &&
-         "Only change functions if both are strong or both are weak");
-  (void)F;
+  Function *F = FN.getFunc();
   assert(FunctionComparator(F, G, &GlobalNumbers).compare() == 0 &&
          "The two functions must be equal");
   
-  IterToF->replaceBy(G);
+  auto I = FNodesInTree.find(F);
+  // We should have that FNodesInTree contains a mapping for F
+  assert(I != FNodesInTree.end());
+  // There shouldn't be a mapping for G
+  assert(FNodesInTree.count(G) == 0);
+  
+  FnTreeType::iterator IterToFNInFnTree = I->second;
+  // Make sure F mapped to FN.
+  assert(&(*IterToFNInFnTree) == &FN);
+  // Remove F -> FN and insert G -> FN
+  FNodesInTree.erase(I);
+  FNodesInTree.insert({G, IterToFNInFnTree});
+  // Replace F with G in FN, which is stored inside the FnTree.
+  FN.replaceBy(G);
 }
 
 // Insert a ComparableFunction into the FnTree, or merge it away if equal to one
 // that was already inserted.
 bool MergeFunctions::insert(Function *NewFunction) {
   std::pair<FnTreeType::iterator, bool> Result =
       FnTree.insert(FunctionNode(NewFunction));
 
   if (Result.second) {
+    assert(FNodesInTree.count(NewFunction) == 0);
+    FNodesInTree.insert({NewFunction, Result.first});
     DEBUG(dbgs() << "Inserting as unique: " << NewFunction->getName() << '\n');
     return false;
   }
 
   const FunctionNode &OldF = *Result.first;
 
   // Don't merge tiny functions, since it can just end up making the function
   // larger.
bool MergeFunctions::insert(Function *NewFunction) {
   //
   // When one function is weak and the other is strong there is an order imposed
   // already. We process strong functions before weak functions.
   if ((OldF.getFunc()->mayBeOverridden() && NewFunction->mayBeOverridden()) ||
       (!OldF.getFunc()->mayBeOverridden() && !NewFunction->mayBeOverridden()))
     if (OldF.getFunc()->getName() > NewFunction->getName()) {
       // Swap the two functions.
       Function *F = OldF.getFunc();
-      replaceFunctionInTree(Result.first, NewFunction);
+      replaceFunctionInTree(*Result.first, NewFunction);
       NewFunction = F;
       assert(OldF.getFunc() != F && "Must have swapped the functions.");
     }
 
   // Never thunk a strong function to a weak function.
   assert(!OldF.getFunc()->mayBeOverridden() || NewFunction->mayBeOverridden());
 
   DEBUG(dbgs() << "  " << OldF.getFunc()->getName()
                << " == " << NewFunction->getName() << '\n');
 
   Function *DeleteF = NewFunction;
   mergeTwoFunctions(OldF.getFunc(), DeleteF);
   return true;
 }
 
 // Remove a function from FnTree. If it was already in FnTree, add
 // it to Deferred so that we'll look at it in the next round.
 void MergeFunctions::remove(Function *F) {
-  // We need to make sure we remove F, not a function "equal" to F per the
-  // function equality comparator.
-  FnTreeType::iterator found = FnTree.find(FunctionNode(F));
-  size_t Erased = 0;
-  if (found != FnTree.end() && found->getFunc() == F) {
-    Erased = 1;
-    FnTree.erase(found);
-  }
-
-  if (Erased) {
-    DEBUG(dbgs() << "Removed " << F->getName()
-                 << " from set and deferred it.\n");
+  auto I = FNodesInTree.find(F);
+  if (I != FNodesInTree.end()) {
+    DEBUG(dbgs() << "Deferred " << F->getName()<< ".\n");
+    FnTree.erase(I->second);
+    // I->second has been invalidated, remove it from the FNodesInTree map to
+    // preserve the invariant.
+    FNodesInTree.erase(I);
     Deferred.emplace_back(F);
   }
 }
 
 // For each instruction used by the value, remove() the function that contains
 // the instruction. This should happen right before a call to RAUW.
 void MergeFunctions::removeUsers(Value *V) {
   std::vector<Value *> Worklist;