[llvm] [Reassociate] Compute value ranks iteratively (PR #204952)

[Hao Ren via llvm-commits]

================
@@ -212,24 +212,78 @@ unsigned ReassociatePass::getRank(Value *V) {
   if (unsigned Rank = ValueRankMap[I])
     return Rank;    // Rank already known?
 
-  // If this is an expression, return the 1+MAX(rank(LHS), rank(RHS)) so that
-  // we can reassociate expressions for code motion!  Since we do not recurse
-  // for PHI nodes, we cannot have infinite recursion here, because there
-  // cannot be loops in the value graph that do not go through PHI nodes.
-  unsigned Rank = 0, MaxRank = RankMap[I->getParent()];
-  for (unsigned i = 0, e = I->getNumOperands(); i != e && Rank != MaxRank; ++i)
-    Rank = std::max(Rank, getRank(I->getOperand(i)));
-
-  // If this is a 'not' or 'neg' instruction, do not count it for rank. This
-  // assures us that X and ~X will have the same rank.
-  if (!match(I, m_Not(m_Value())) && !match(I, m_Neg(m_Value())) &&
-      !match(I, m_FNeg(m_Value())))
-    ++Rank;
-
-  LLVM_DEBUG(dbgs() << "Calculated Rank[" << V->getName() << "] = " << Rank
-                    << "\n");
-
-  return ValueRankMap[I] = Rank;
+  // Return 1+MAX(rank(LHS), rank(RHS)) for expressions so we can reassociate
+  // expressions for code motion. Use an explicit worklist rather than native
+  // recursion so long acyclic use-def chains do not overflow the stack.
+  struct RankWorkItem {
+    Instruction *I;
+    unsigned OpNo;
+    unsigned Rank;
+    unsigned MaxRank;
+  };
+
+  auto GetLeafRank = [this](Value *V) {
+    return isa<Argument>(V) ? ValueRankMap[V] : 0;
+  };
+
+  SmallVector<RankWorkItem, 16> Worklist;
+  Worklist.push_back(RankWorkItem{I, 0, 0, RankMap[I->getParent()]});
+
+  auto CompleteRank = [&](unsigned Rank) {
+    // A rank has been computed for the current work item. Pop it and, if a
+    // parent is waiting, fold that rank into the parent before resuming there.
+    Worklist.pop_back();
+    if (Worklist.empty())
+      return true;
+
+    RankWorkItem &Parent = Worklist.back();
+    Parent.Rank = std::max(Parent.Rank, Rank);
+    ++Parent.OpNo;
+    return false;
+  };
+
+  do {
----------------
nvidia-moomoo wrote:

changed

https://github.com/llvm/llvm-project/pull/204952