[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 ↗

(On Diff #427506)

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

17

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

17–27

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

17

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?

70

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?

70

Also what is the use-case for the callback?

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

49

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

[rriddle]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

17

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

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

17

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

70

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

49

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

17

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

49

should be ok indeed

[Mogball]

split out TestOps change as NFC

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

17

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

[mehdi_amini]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

17

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

17

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

15

nit: Can you drop the llvm:: here?

31

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.
+
+    This sort is stable. If the block is already topologically sorted, the IR
+    is not changed. Operations that form a cycle are moved to the end of the
+    regions in a stable order.
+  }];
+
+  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, sort a range operations in said block in topological order.
+/// The main purpose is readability of graph regions, potentially faster
+/// processing of certain transformations and analyses, or fixing the SSA
+/// dominance of blocks that require it after transformations. 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. This sort is stable;
+/// if the operations are already topologically sorted, nothing changes.
+///
+/// Operations that form cycles are moved to the end of the block in order. 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. It is called on top-level operations in the
+/// block with value uses at or below those operations. The function should
+/// return true to mark that value as ready to be scheduled.
+///
+/// For example, if `isOperandReady` is set to always mark edges from `foo.A` to
+/// `foo.B` as ready, these operations:
+///
+/// ```mlir
+/// %0 = foo.B(%1)
+/// %1 = foo.C(%2)
+/// %2 = foo.A(%0)
+/// ```
+///
+/// Are sorted as:
+///
+/// ```mlir
+/// %0 = foo.A(%2)
+/// %1 = foo.C(%0)
+/// %2 = foo.B(%1)
+/// ```
+bool sortTopologically(
+    Block *block, iterator_range<Block::iterator> ops,
+    function_ref<bool(Value, Operation *)> isOperandReady = nullptr);
+
+/// Given a block, sort its operations in topological order, excluding its
+/// terminator if it has one.
+bool sortTopologically(
+    Block *block,
+    function_ref<bool(Value, Operation *)> isOperandReady = nullptr);
+
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_TOPOLOGICALSORTUTILS_H

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
   ControlFlowSink.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   OpStats.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   SymbolPrivatize.cpp
+  TopologicalSort.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRTransformsPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRCopyOpInterface
   MLIRLoopLikeInterface
   MLIRPass
   MLIRSupport
   MLIRTransformUtils
   )

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;
+
+namespace {
+struct TopologicalSortPass : public TopologicalSortBase<TopologicalSortPass> {
+  void runOnOperation() override {
+    // Topologically sort the regions of the operation without SSA dominance.
+    getOperation()->walk([](RegionKindInterface op) {
+      for (auto &it : llvm::enumerate(op->getRegions())) {
+        if (op.hasSSADominance(it.index()))
+          continue;
+        for (Block &block : it.value())
+          sortTopologically(&block);
+      }
+    });
+  }
+};
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createTopologicalSortPass() {
+  return std::make_unique<TopologicalSortPass>();
+}

mlir/lib/Transforms/Utils/CMakeLists.txt

 add_mlir_library(MLIRTransformUtils
   ControlFlowSinkUtils.cpp
   DialectConversion.cpp
   FoldUtils.cpp
   GreedyPatternRewriteDriver.cpp
   InliningUtils.cpp
   LoopInvariantCodeMotionUtils.cpp
   RegionUtils.cpp
   SideEffectUtils.cpp
+  TopologicalSortUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRLoopLikeInterface
   MLIRRewrite
   )

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 && 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
+}