[mlir][transforms] Add a topological sort utility and pass

[Mogball]

Summary

This patch adds a topological sort utility and pass. A topological sort reorders
the operations in a block without SSA dominance such that, as much as possible,
users of values come after their producers.

The utility function sorts topologically the operation range in a given block
with an optional user-provided callback that can be used to virtually break cycles.
The toposort pass itself recursively sorts graph regions under the target op.

[Mogball]

Context: this pass was sitting around internally and since it's not anything specific, thought I'd upstream it

[rriddle]

mlir/test/lib/Dialect/Test/TestDialect.cpp

725

Can you split this out into an NFC and submit separately?

[rriddle]

There is also the pre-existing mlir::topologicalSort which we should either rename or make private (given that it isn't great).

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

The missing SSA dominance is fairly significant of a limitation. I feel like a lot of people are going to want to use this, and then get burned immediately. We should either rename this, or add proper asserts before executing.

Also, this seems to completely disregard side effects?

mlir/lib/Transforms/TopologicalSort.cpp

16–26

Honestly I'd just inline this function into the single use.

[mehdi_amini]

mlir/include/mlir/Transforms/Passes.td

270

Please document that this is a stable sort (nothing changes if the block is already topo-sorted right?) and the behavior for cycles.

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

Should we say "Given a block in a Graph Region" instead of "without SSA dominance"?

Also, this seems to completely disregard side effects?

In a Graph-region there is no implicit ordering, am I missing something?

69

Can you document the callback?
Also function_ref can be "null", would it be an OK default and handling it in the code instead of calling it?

69

Also what is the use-case for the callback?

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

48

This isn't handling ops with region and implicit captures I believe.

[rriddle]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

You're right, I forgot that we explicitly codified that operations can be freely reordered.

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

Is there a difference in meaning to say "Block in a Graph Region" vs. "without SSA dominance"?

69

The sort will break cycles once all other ops have been scheduled and will break the cycle at the first op it encounters (top-down) in the cycle. This callback allows cycles to be broken early in special cases (e.g. NextIteration into Merge)

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

48

Can you elaborate on the particular case that isn't handled?

I added two tests below, one which is implicit capture of an SSA value from a predecessor block and another is implicit capture of values defined above the region.

[mehdi_amini]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

Better naming positively (what it is) than negatively (what it isn't) and Graph Region is the name used in the doc: https://mlir.llvm.org/docs/LangRef/#graph-regions

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

48

should be ok indeed

[Mogball]

split out TestOps change as NFC

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

What would be the purpose of running this on a block with SSA dominance?

[mehdi_amini]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

I had such need once though where I would perform a transform (I don't remember what) and it what more convenient to just append always to the block and toposort later.

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

16

Ah, interesting. There's no actual restriction that the block be in a graph region so I can drop that requirement from the documentation.

[Mogball]

review comments, and drop requirement that blocks need to not have SSA dominance

[Mogball]

fix cmake build

[rriddle]

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

16

nit: Can you drop the llvm:: here?

32

nit: Drop the const here.

mlir/include/mlir/Transforms/Passes.h

 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();
 
 /// Creates a pass which marks top-level symbol operations as `private` unless
 /// listed in `excludeSymbols`.
 std::unique_ptr<Pass>
 createSymbolPrivatizePass(ArrayRef<std::string> excludeSymbols = {});
 
+/// Creates a pass that recursively sorts nested regions without SSA dominance
+/// topologically such that, as much as possible, users of values appear after
+/// their producers.
+std::unique_ptr<Pass> createTopologicalSortPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/include/mlir/Transforms/Passes.td

let options = [
     Option<"printDataFlowEdges", "print-data-flow-edges", "bool",
            /*default=*/"true", "Print data flow edges">,
     Option<"printResultTypes", "print-result-types", "bool",
             /*default=*/"true", "Print result types of operations">
   ];
   let constructor = "mlir::createPrintOpGraphPass()";
 }
 
+def TopologicalSort : Pass<"topological-sort"> {
+  let summary = "Sort regions without SSA dominance in topological order";
+  let description = [{
+    Recursively sorts all nested regions without SSA dominance in topological
+    order. The main purpose is readability, as well as potentially processing of
+    certain transformations and analyses. The function sorts the operations in
+    all nested regions such that, as much as possible, all users appear after
+    their producers.
+  }];
+
+  let constructor = "mlir::createTopologicalSortPass()";
+}
+
 #endif // MLIR_TRANSFORMS_PASSES

mlir/include/mlir/Transforms/TopologicalSortUtils.h

+//===- TopologicalSortUtils.h - Topological sort utilities ------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_TOPOLOGICALSORTUTILS_H
+#define MLIR_TRANSFORMS_TOPOLOGICALSORTUTILS_H
+
+#include "mlir/IR/Block.h"
+
+namespace mlir {
+
+/// Given a block without SSA dominance, sort the range operations in
+/// topological order. The main purpose is readability, as well as potentially
+/// faster processing of certain transformations and analyses. The function
+/// sorts the given operations such that, as much as possible, all users appear
+/// after their producers.
+///
+/// For example:
+///
+/// ```mlir
+/// %0 = test.foo
+/// %1 = test.bar %0, %2
+/// %2 = test.baz
+/// ```
+///
+/// Will become:
+///
+/// ```mlir
+/// %0 = test.foo
+/// %1 = test.baz
+/// %2 = test.bar %0, %1
+/// ```
+///
+/// The sort also works on operations with regions and implicit captures. For
+/// example:
+///
+/// ```mlir
+/// %0 = test.foo {
+///   test.baz %1
+///   %1 = test.bar %2
+/// }
+/// %2 = test.foo
+/// ```
+///
+/// Will become:
+///
+/// ```mlir
+/// %0 = test.foo
+/// %1 = test.foo {
+///   test.baz %2
+///   %2 = test.bar %0
+/// }
+/// ```
+///
+/// Note that the sort is not recursive on nested regions.
+///
+/// If the sort is left with only operations that form a cycle, it breaks the
+/// cycle by marking the first encountered operation as ready and moving on.
+///
+/// The function optionally accepts a callback that can be provided by users to
+/// virtually break cycles early.
+bool sortTopologically(
+    Block *block, iterator_range<Block::iterator> ops,
+    function_ref<bool(Value, Operation *)> isOperandReady =
+        [](Value, Operation *) { return false; });
+
+/// Given a block without SSA dominance, sort its operations in order, excluding
+/// its terminator if it has one.
+bool sortTopologically(
+    Block *block, function_ref<bool(Value, Operation *)> isOperandReady =
+                      [](Value, Operation *) { return false; });
+
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_TOPOLOGICALSORTUTILS_H

mlir/lib/Transforms/TopologicalSort.cpp

+//===- TopologicalSort.cpp - Topological sort pass ------------------------===//
+//
+// 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 "PassDetail.h"
+#include "mlir/IR/RegionKindInterface.h"
+#include "mlir/Transforms/TopologicalSortUtils.h"
+
+using namespace mlir;
+
+/// Topologically sort the regions of the operation without SSA dominance.
+static void sortRegionsTopologically(Operation *op) {
+  auto regionInterface = dyn_cast<RegionKindInterface>(op);
+  if (!regionInterface)
+    return;
+  for (auto &it : llvm::enumerate(op->getRegions())) {
+    if (regionInterface.hasSSADominance(it.index()))
+      continue;
+    for (Block &block : it.value())
+      sortTopologically(&block);
+  }
+}
+
+namespace {
+struct TopologicalSortPass : public TopologicalSortBase<TopologicalSortPass> {
+  void runOnOperation() override {
+    getOperation()->walk(sortRegionsTopologically);
+  }
+};
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createTopologicalSortPass() {
+  return std::make_unique<TopologicalSortPass>();
+}

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

+//===- TopologicalSortUtils.h - Topological sort utilities ------*- C++ -*-===//
+//
+// 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/Transforms/TopologicalSortUtils.h"
+#include "mlir/IR/OpDefinition.h"
+
+using namespace mlir;
+
+bool mlir::sortTopologically(
+    Block *block, llvm::iterator_range<Block::iterator> ops,
+    function_ref<bool(Value, Operation *)> isOperandReady) {
+  if (ops.empty())
+    return true;
+
+  // The set of operations that have not yet been scheduled.
+  DenseSet<Operation *> unscheduledOps;
+  // Mark all operations as unscheduled.
+  for (Operation &op : ops)
+    unscheduledOps.insert(&op);
+
+  Block::iterator nextScheduledOp = ops.begin();
+  Block::iterator end = ops.end();
+
+  // An operation is ready to be scheduled if all its operands are ready. An
+  // operation is ready if:
+  const auto isReady = [&](Value value, Operation *top) {
+    // - the user-provided callback marks it as ready,
+    if (isOperandReady(value, top))
+      return true;
+    Operation *parent = value.getDefiningOp();
+    // - it is a block argument,
+    if (!parent)
+      return true;
+    Operation *ancestor = block->findAncestorOpInBlock(*parent);
+    // - it is an implicit capture,
+    if (!ancestor)
+      return true;
+    // - it is defined in a nested region, or
+    if (ancestor == top)
+      return true;
+    // - its ancestor in the block is scheduled.
+    return !unscheduledOps.contains(ancestor);
+  };
+
+  bool allOpsScheduled = true;
+  while (!unscheduledOps.empty()) {
+    bool scheduledAtLeastOnce = false;
+
+    // Loop over the ops that are not sorted yet, try to find the ones "ready",
+    // i.e. the ones for which there aren't any operand produced by an op in the
+    // set, and "schedule" it (move it before the `nextScheduledOp`).
+    for (Operation &op :
+         llvm::make_early_inc_range(llvm::make_range(nextScheduledOp, end))) {
+      // An operation is recursively ready to be scheduled of it and its nested
+      // operations are ready.
+      WalkResult readyToSchedule = op.walk([&](Operation *nestedOp) {
+        return llvm::all_of(
+                   nestedOp->getOperands(),
+                   [&](Value operand) { return isReady(operand, &op); })
+                   ? WalkResult::advance()
+                   : WalkResult::interrupt();
+      });
+      if (readyToSchedule.wasInterrupted())
+        continue;
+
+      // Schedule the operation by moving it to the start.
+      unscheduledOps.erase(&op);
+      op.moveBefore(block, nextScheduledOp);
+      scheduledAtLeastOnce = true;
+      // Move the iterator forward if we schedule the operation at the front.
+      if (&op == &*nextScheduledOp)
+        ++nextScheduledOp;
+    }
+    // If no operations were scheduled, give up and advance the iterator.
+    if (!scheduledAtLeastOnce) {
+      allOpsScheduled = false;
+      unscheduledOps.erase(&*nextScheduledOp);
+      ++nextScheduledOp;
+    }
+  }
+
+  return allOpsScheduled;
+}
+
+bool mlir::sortTopologically(
+    Block *block, function_ref<bool(Value, Operation *)> isOperandReady) {
+  if (block->empty())
+    return true;
+  if (block->back().hasTrait<OpTrait::IsTerminator>())
+    return sortTopologically(block, block->without_terminator(),
+                             isOperandReady);
+  return sortTopologically(block, *block, isOperandReady);
+}

mlir/test/Transforms/test-toposort.mlir

+// RUN: mlir-opt -topological-sort %s | FileCheck %s
+
+// Test producer is after user.
+// CHECK-LABEL: test.graph_region
+test.graph_region {
+  // CHECK-NEXT: test.foo
+  // CHECK-NEXT: test.baz
+  // CHECK-NEXT: test.bar
+  %0 = "test.foo"() : () -> i32
+  "test.bar"(%1, %0) : (i32, i32) -> ()
+  %1 = "test.baz"() : () -> i32
+}
+
+// Test cycles.
+// CHECK-LABEL: test.graph_region
+test.graph_region {
+  // CHECK-NEXT: test.d
+  // CHECK-NEXT: test.a
+  // CHECK-NEXT: test.c
+  // CHECK-NEXT: test.b
+  %2 = "test.c"(%1) : (i32) -> i32
+  %1 = "test.b"(%0, %2) : (i32, i32) -> i32
+  %0 = "test.a"(%3) : (i32) -> i32
+  %3 = "test.d"() : () -> i32
+}
+
+// Test block arguments.
+// CHECK-LABEL: test.graph_region
+test.graph_region {
+// CHECK-NEXT: (%{{.*}}:
+^entry(%arg0: i32):
+  // CHECK-NEXT: test.foo
+  // CHECK-NEXT: test.baz
+  // CHECK-NEXT: test.bar
+  %0 = "test.foo"(%arg0) : (i32) -> i32
+  "test.bar"(%1, %0) : (i32, i32) -> ()
+  %1 = "test.baz"(%arg0) : (i32) -> i32
+}
+
+// Test implicit block capture (and sort nested region).
+// CHECK-LABEL: test.graph_region
+func.func @test_graph_cfg() -> () {
+  %0 = "test.foo"() : () -> i32
+  cf.br ^next(%0 : i32)
+
+^next(%1: i32):
+  test.graph_region {
+    // CHECK-NEXT: test.foo
+    // CHECK-NEXT: test.baz
+    // CHECK-NEXT: test.bar
+    %2 = "test.foo"(%1) : (i32) -> i32
+    "test.bar"(%3, %2) : (i32, i32) -> ()
+    %3 = "test.baz"(%0) : (i32) -> i32
+  }
+  return
+}
+
+// Test region ops (and recursive sort).
+// CHECK-LABEL: test.graph_region
+test.graph_region {
+  // CHECK-NEXT: test.baz
+  // CHECK-NEXT: test.graph_region attributes {a} {
+  // CHECK-NEXT:   test.b
+  // CHECK-NEXT:   test.a
+  // CHECK-NEXT: }
+  // CHECK-NEXT: test.bar
+  // CHECK-NEXT: test.foo
+  %0 = "test.foo"(%1) : (i32) -> i32
+  test.graph_region attributes {a} {
+    %a = "test.a"(%b) : (i32) -> i32
+    %b = "test.b"(%2) : (i32) -> i32
+  }
+  %1 = "test.bar"(%2) : (i32) -> i32
+  %2 = "test.baz"() : () -> i32
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 
 RegionKind SSACFGRegionOp::getRegionKind(unsigned index) {
   return RegionKind::SSACFG;
 }
 
 //===----------------------------------------------------------------------===//
 // Test GraphRegionOp
 //===----------------------------------------------------------------------===//
 
-ParseResult GraphRegionOp::parse(OpAsmParser &parser, OperationState &result) {
-  // Parse the body region, and reuse the operand info as the argument info.
-  Region *body = result.addRegion();
-  return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
-}
-
-void GraphRegionOp::print(OpAsmPrinter &p) {
-  p << "test.graph_region ";
-  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
-}
-
 RegionKind GraphRegionOp::getRegionKind(unsigned index) {
   return RegionKind::Graph;
 }
 
 //===----------------------------------------------------------------------===//
 // Test AffineScopeOp
 //===----------------------------------------------------------------------===//
 

mlir/test/lib/Dialect/Test/TestOps.td

 def GraphRegionOp : TEST_Op<"graph_region",  [
     DeclareOpInterfaceMethods<RegionKindInterface>]> {
   let summary =  "operation with a graph region";
   let description = [{
     Test op that defines a graph region.
   }];
 
   let regions = (region AnyRegion:$region);
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = "attr-dict-with-keyword $region";
 }
 
 def AffineScopeOp : TEST_Op<"affine_scope", [AffineScope]> {
   let summary =  "affine scope operation";
   let description = [{
     Test op that defines a new affine scope.
   }];