[MLIR] Implemented iterative topological sorting

[georgemitenkov]

Summary

Implemented an iterative preorder DFS topological sort
instead of the previous recursive version.

[rriddle]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

363

nit: You could just use llvm::ReversePostOrderTraversal to get the topological sort.

364

nit: Drop the llvm:: here, SmallVector is re-exported into the mlir namespace.

368

Why the move here? Also, you can you pop_back_val instead.

[rriddle]

Thanks for working on this!

[kiranchandramohan]

Thanks for working on this.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

363

Can this be inlined into topologicalSort?

[georgemitenkov]

Addressed comments and reused llvm::ReversePostOrderTraversal for topological sorting.

[kiranchandramohan]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

389

Nit: MLIR adopts camelBack convention for names. Switch from RPOT-> rpot?
https://mlir.llvm.org/getting_started/DeveloperGuide/

[rriddle]

LGTM after resolving the remaining comments. You might want to clarify the commit title and description a bit more, it is very broad for what you are fixing here.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

390

nit: Use the range based insert instead.

Also in the future, please cache the end iterator in loops when it is unchanged within the loop.

[georgemitenkov]

Rebased and addressed comments.

[kiranchandramohan]

LGTM.

[dcaballe]

Hey! Thanks for the patch! I wondered if it would make sense to move this utility to a public utility header. I found myself needing it locally and I wouldn't like to just copy-paste it. There is another topological sort for ops under 'include/mlir/Analysis/SliceAnalysis.h'. Maybe we could move them both to 'include/mlir/IR/Utils.h'? WDYT?

[georgemitenkov]

In D88544#2309019, @dcaballe wrote:

Hey! Thanks for the patch! I wondered if it would make sense to move this utility to a public utility header. I found myself needing it locally and I wouldn't like to just copy-paste it. There is another topological sort for ops under 'include/mlir/Analysis/SliceAnalysis.h'. Maybe we could move them both to 'include/mlir/IR/Utils.h'? WDYT?

I think that if this functionality is needed elsewhere, then it should make sense to publicly expose it. Having 'include/mlir/IR/Utils.h' sounds reasonable - maybe we can have something similar to LLVM's ADT with MLIR specific utilities/algorithms built on top (I guess this is not only the case of topological sorting)? @rriddle, @mehdi_amini, @ftynse WDYT?

[mehdi_amini]

We use https://github.com/llvm/llvm-project/tree/master/mlir/include/mlir/Support in general for our helpers.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (auto &numberedPhiNode : llvm::enumerate(phis)) {
         phiNode.addIncoming(valueMapping.lookup(getPHISourceValue(
                                 bb, pred, numArguments, index)),
                             blockMapping.lookup(pred));
       }
     }
   }
 }
 
-// TODO: implement an iterative version
-static void topologicalSortImpl(llvm::SetVector<Block *> &blocks, Block *b) {
+static void topologicalSortImpl(llvm::SetVector<Block *> &blocks, Block *root) {
+  llvm::SmallVector<Block *, 8> stack;
+  stack.push_back(root);
+
+  while (!stack.empty()) {
+    Block *b = std::move(stack.back());
+    stack.pop_back();
     blocks.insert(b);
-  for (Block *bb : b->getSuccessors()) {
-    if (blocks.count(bb) == 0)
-      topologicalSortImpl(blocks, bb);
+
+    for (Block *succ : b->getSuccessors()) {
+      if (blocks.count(succ) == 0)
+        stack.push_back(succ);
+    }
   }
 }
 
 /// Sort function blocks topologically.
 template <typename T>
 static llvm::SetVector<Block *> topologicalSort(T &f) {
-  // For each blocks that has not been visited yet (i.e. that has no
-  // predecessors), add it to the list and traverse its successors in DFS
-  // preorder.
+  // For each block that has not been visited yet (i.e. that has no
+  // predecessors), add it to the list and traverse its successors iteratively
+  // in DFS preorder.
   llvm::SetVector<Block *> blocks;
   for (Block &b : f) {
     if (blocks.count(&b) == 0)
       topologicalSortImpl(blocks, &b);
   }
   assert(blocks.size() == f.getBlocks().size() && "some blocks are not sorted");
 
   return blocks;
 }
 
 /// Convert the OpenMP parallel Operation to LLVM IR.
 LogicalResult
 ModuleTranslation::convertOmpParallel(Operation &opInst,
                                       llvm::IRBuilder<> &builder) {