[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

170–186

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

170–186

Changed it to be iterative.

[rriddle]

mlir/lib/Analysis/SliceAnalysis.cpp

182–185

190–192

nit: drop trivial braces

[ThomasRaoux]

Address review comments

[mehdi_amini]

Thanks!

mlir/lib/Analysis/SliceAnalysis.cpp

 struct DFSState {
   DFSState(const SetVector<Operation *> &set)
       : toSort(set), topologicalCounts(), seen() {}
   const SetVector<Operation *> &toSort;
   SmallVector<Operation *, 16> topologicalCounts;
   DenseSet<Operation *> seen;
 };
 } // namespace
 
-static void DFSPostorder(Operation *current, DFSState *state) {
+static void DFSPostorder(Operation *root, DFSState *state) {
+  SmallVector<Operation *> queue(1, root);
+  std::vector<Operation *> ops;
+  while (!queue.empty()) {
+    Operation *current = queue.pop_back_val();
+    ops.push_back(current);
     for (Value result : current->getResults()) {
       for (Operation *op : result.getUsers())
-      DFSPostorder(op, state);
+        queue.push_back(op);
     }
-  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);
+    for (Region &region : current->getRegions()) {
+      for (Block &block : region.getBlocks()) {
+        for (Operation &op : block.getOperations())
+          queue.push_back(&op);
+      }
+    }
+  }
+
+  for (Operation *op : llvm::reverse(ops)) {
+    if (state->seen.insert(op).second && state->toSort.count(op) > 0) {
+      state->topologicalCounts.push_back(op);
     }
   }
 }
 
 SetVector<Operation *>
 mlir::topologicalSort(const SetVector<Operation *> &toSort) {
   if (toSort.empty()) {
     return toSort;
   }

mlir/test/Analysis/test-topoligical-sort.mlir

+// RUN: mlir-opt %s -test-print-topological-sort 2>&1 | FileCheck %s
+
+// CHECK-LABEL: Testing : region
+//       CHECK: arith.addi {{.*}} : index
+//  CHECK-NEXT: scf.for
+//       CHECK: } {__test_sort_original_idx__ = 2 : i64}
+//  CHECK-NEXT: arith.addi {{.*}} : i32
+//  CHECK-NEXT: arith.subi {{.*}} : i32
+func @region(
+  %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index,
+  %arg4 : i32, %arg5 : i32, %arg6 : i32,
+  %buffer : memref<i32>) {
+  %0 = arith.addi %arg4, %arg5 {__test_sort_original_idx__ = 0} : i32
+  %idx = arith.addi %arg0, %arg1 {__test_sort_original_idx__ = 3} : index
+  scf.for %arg7 = %idx to %arg2 step %arg3  {
+    %2 = arith.addi %0, %arg5 : i32
+    %3 = arith.subi %2, %arg6 {__test_sort_original_idx__ = 1} : i32
+    memref.store %3, %buffer[] : memref<i32>
+  } {__test_sort_original_idx__ = 2}
+  return
+}

mlir/test/lib/Analysis/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestAnalysis
   TestAliasAnalysis.cpp
   TestCallGraph.cpp
   TestLiveness.cpp
   TestMatchReduction.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestNumberOfExecutions.cpp
+  TestSlice.cpp
 
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAnalysis
   MLIRMemRef

mlir/test/lib/Analysis/TestSlice.cpp

+//===------------- TestSlice.cpp - Test slice related analisis ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+static const StringLiteral kOrderMarker = "__test_sort_original_idx__";
+
+namespace {
+
+struct TestTopologicalSortPass
+    : public PassWrapper<TestTopologicalSortPass, FunctionPass> {
+  StringRef getArgument() const final { return "test-print-topological-sort"; }
+  StringRef getDescription() const final {
+    return "Print operations in topological order";
+  }
+  void runOnFunction() override {
+    std::map<int, Operation *> ops;
+    getFunction().walk([&ops](Operation *op) {
+      if (auto originalOrderAttr = op->getAttrOfType<IntegerAttr>(kOrderMarker))
+        ops[originalOrderAttr.getInt()] = op;
+    });
+    SetVector<Operation *> sortedOp;
+    for (auto op : ops)
+      sortedOp.insert(op.second);
+    sortedOp = topologicalSort(sortedOp);
+    llvm::errs() << "Testing : " << getFunction().getName() << "\n";
+    for (Operation *op : sortedOp) {
+      op->print(llvm::errs());
+      llvm::errs() << "\n";
+    }
+  }
+};
+
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestSliceAnalysisPass() {
+  PassRegistration<TestTopologicalSortPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMemRefStrideCalculation();
 void registerTestNumberOfBlockExecutionsPass();
 void registerTestNumberOfOperationExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPDLByteCodePass();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
+void registerTestSliceAnalysisPass();
 void registerTestVectorConversions();
 } // namespace test
 } // namespace mlir
 
 namespace test {
 void registerTestDialect(DialectRegistry &);
 } // namespace test
 
#endif
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestNumberOfBlockExecutionsPass();
   mlir::test::registerTestNumberOfOperationExecutionsPass();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
+  mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestVectorConversions();
 }
 #endif
 
 int main(int argc, char **argv) {
   registerAllPasses();
 #ifdef MLIR_INCLUDE_TESTS
   registerTestPasses();