[SelectionDAG] Get rid of recursion in CalcNodeSethiUllmanNumber

[mkazantsev]

Summary

The recursive implementation of CalcNodeSethiUllmanNumber may
overflow stack on extremely long pred chains. This patch replaces it
with an equivalent iterative implementation.

[mkazantsev]

Ping

[reames]

lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

1864 ↗

(On Diff #101004)

I think you can use a much simpler bit of code here. The key bit is that we can't compute the current node until all predecessors are done, but that (by assumption from the existing code) the graph is a tree. (Add an assert that enforces that please!) Also, Wikipedia helps here. :)

If you do something along the lines of:
while (!Worklist.empty())

Cur = Worklist.pop_back();
if (not all preds done) {
   push(Cur)
   push(each pred)
  continue
}
compute answer using preds

}

I think you get the same result right? This will eliminate the need for the intermediate state.

[mkazantsev]

Rewritten in a more clear manner.

lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

1864 ↗

(On Diff #101004)

We still need a workstate to keep the current pred in order to avoid pushing an element more than once into the stack, but you are right, it can be done in a much more straightforward way. :)

[reames]

lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

1864 ↗

(On Diff #101004)

Why do you need the predecessors in order or to track which ones have been previously processed?

This is a reduction over a set of child nodes. The order shouldn't matter and the "recursive" walk is responsible for filling in the value. It shouldn't even matter if the code is DFS or BFS if I'm reading the original right.

[mkazantsev]

lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

1864 ↗

(On Diff #101004)

It's not for tracking order, it's for preventing us from pushing one node more than once. Imagine the case:

preds(a) = (x, y, z, b)
preds(b) = (x, y, z, c)
preds(c) = (x, y, z, d)

I need the iterator to not push x, y, z into stack 3 times.

[mkazantsev]

lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

1864 ↗

(On Diff #101004)

Please note that Sethi-Ullman's algorithm described in Wikipedia works on trees, and here we are dealing with SUnits which are nodes in Scheduling DAG. So it is possible that one SUnit is a predecessor of many others. In this case we need to preserve the calculation order of recursive DFS unless we want one element to be pushed many times as in the example above. I just double-checked that it actually happens (and we fail the assertion if it does). So please don't be confused by Wilipedia's description of the algorithm. This modification is a bit different.

[reames]

Argument accepted, minor comment below.

lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

1864 ↗

(On Diff #101997)

Reduce nesting by inverting conditional or using impl helper pattern.

[reames]

LGTM subject to previously stated minor comment being fixed.

llvm/trunk/lib/CodeGen/SelectionDAG/ScheduleDAGRRList.cpp

if (LSchedLow != RSchedLow)
     return LSchedLow < RSchedLow ? 1 : -1;
   return 0;
 }
 
 /// CalcNodeSethiUllmanNumber - Compute Sethi Ullman number.
 /// Smaller number is the higher priority.
 static unsigned
 CalcNodeSethiUllmanNumber(const SUnit *SU, std::vector<unsigned> &SUNumbers) {
-  unsigned &SethiUllmanNumber = SUNumbers[SU->NodeNum];
-  if (SethiUllmanNumber != 0)
-    return SethiUllmanNumber;
+  if (SUNumbers[SU->NodeNum] != 0)
+    return SUNumbers[SU->NodeNum];
 
+  // Use WorkList to avoid stack overflow on excessively large IRs.
+  struct WorkState {
+    WorkState(const SUnit *SU) : SU(SU) {}
+    const SUnit *SU;
+    unsigned PredsProcessed = 0;
+  };
+
+  SmallVector<WorkState, 16> WorkList;
+  WorkList.push_back(SU);
+  while (!WorkList.empty()) {
+    auto &Temp = WorkList.back();
+    auto *TempSU = Temp.SU;
+    bool AllPredsKnown = true;
+    // Try to find a non-evaluated pred and push it into the processing stack.
+    for (unsigned P = Temp.PredsProcessed; P < TempSU->Preds.size(); ++P) {
+      auto &Pred = TempSU->Preds[P];
+      if (Pred.isCtrl()) continue;  // ignore chain preds
+      SUnit *PredSU = Pred.getSUnit();
+      if (SUNumbers[PredSU->NodeNum] == 0) {
+#ifndef NDEBUG
+        // In debug mode, check that we don't have such element in the stack.
+        for (auto It : WorkList)
+          assert(It.SU != PredSU && "Trying to push an element twice?");
+#endif
+        WorkList.push_back(PredSU);
+        AllPredsKnown = false;
+        // Next time start processing this one starting from the next pred.
+        Temp.PredsProcessed = P + 1;
+        break;
+      }
+    }
+
+    if (!AllPredsKnown)
+      continue;
+
+    // Once all preds are known, we can calculate the answer for this one.
+    unsigned SethiUllmanNumber = 0;
     unsigned Extra = 0;
-  for (const SDep &Pred : SU->Preds) {
+    for (const SDep &Pred : TempSU->Preds) {
       if (Pred.isCtrl()) continue;  // ignore chain preds
       SUnit *PredSU = Pred.getSUnit();
-    unsigned PredSethiUllman = CalcNodeSethiUllmanNumber(PredSU, SUNumbers);
+      unsigned PredSethiUllman = SUNumbers[PredSU->NodeNum];
+      assert(PredSethiUllman > 0 && "We should have evaluated this pred!");
       if (PredSethiUllman > SethiUllmanNumber) {
         SethiUllmanNumber = PredSethiUllman;
         Extra = 0;
       } else if (PredSethiUllman == SethiUllmanNumber)
         ++Extra;
     }
 
     SethiUllmanNumber += Extra;
-
     if (SethiUllmanNumber == 0)
       SethiUllmanNumber = 1;
+    SUNumbers[TempSU->NodeNum] = SethiUllmanNumber;
+    WorkList.pop_back();
+  }
 
-  return SethiUllmanNumber;
+  assert(SUNumbers[SU->NodeNum] > 0 && "SethiUllman should never be zero!");
+  return SUNumbers[SU->NodeNum];
 }
 
 /// CalculateSethiUllmanNumbers - Calculate Sethi-Ullman numbers of all
 /// scheduling units.
 void RegReductionPQBase::CalculateSethiUllmanNumbers() {
   SethiUllmanNumbers.assign(SUnits->size(), 0);
 
   for (const SUnit &SU : *SUnits)