[mlir] Make topologicalSort consider op regions

[ThomasRaoux]

Summary

When doing topological sort we need to make sure an op is scheduled before any of the ops within its regions.

[mehdi_amini]

Is this something that you can make visible with a test somehow?

[ThomasRaoux]

In D113423#3116686, @mehdi_amini wrote:

Is this something that you can make visible with a test somehow?

I have a following patch using this, I can combine it or do you think it should be tested independently of any other user? I would need to add some test pass for this and somehow test slicing/sorting individually.

[mehdi_amini]

Can you stack up the dependent change in Phabricator?

mlir/lib/Analysis/SliceAnalysis.cpp

173–174

Reading the code a bit more now, this seems like a recursive algorithm on use-def chain?
Seems like a ticking bomb for a stack overflow (we call this out in the list here: https://mlir.llvm.org/getting_started/DeveloperGuide/#style-guide )

This really should be iterative instead.

[ThomasRaoux]

Remove recursion and add a test for topological sort

[ThomasRaoux]

In D113423#3116686, @mehdi_amini wrote:

Is this something that you can make visible with a test somehow?

So I ended up adding a test pass just for this. Let me know what you think.

mlir/lib/Analysis/SliceAnalysis.cpp

173–174

Changed it to be iterative.

[rriddle]

mlir/lib/Analysis/SliceAnalysis.cpp

185–188

190–192

nit: drop trivial braces

[ThomasRaoux]

Address review comments

[mehdi_amini]

Thanks!

mlir/lib/Analysis/SliceAnalysis.cpp

struct DFSState {
   const SetVector<Operation *> &toSort;
   SmallVector<Operation *, 16> topologicalCounts;
   DenseSet<Operation *> seen;
 };
 } // namespace
 
 static void DFSPostorder(Operation *current, DFSState *state) {
   for (Value result : current->getResults()) {
     for (Operation *op : result.getUsers())
       DFSPostorder(op, state);
+    for (Region &region : current->getRegions()) {
+      for (Block &block : region.getBlocks()) {
+        for (Operation &op : block.getOperations())
+          DFSPostorder(&op, state);
+      }
+    }
   }
   bool inserted;
   using IterTy = decltype(state->seen.begin());
   IterTy iter;
   std::tie(iter, inserted) = state->seen.insert(current);
   if (inserted) {
     if (state->toSort.count(current) > 0) {
       state->topologicalCounts.push_back(current);
     }
   }
 }
 
 SetVector<Operation *>
 mlir::topologicalSort(const SetVector<Operation *> &toSort) {
   if (toSort.empty()) {
     return toSort;
   }
 
   // Run from each root with global count and `seen` set.
   DFSState state(toSort);