[mlir][transforms] Add topological sort analysis

[springerm]

Summary

This change add a helper function for computing a topological sorting of a list of ops. E.g. this can be useful in transforms where a subset of ops should be cloned without dominance errors.

The analysis reuses the existing implementation in TopologicalSortUtils.cpp.

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

101–103

?

Otherwise, I'm not able to understand the comment.

104

I wouldn't use SmallVector if the number of ops being sorted is potentially large (is it)?

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

58

Is ops expected to be small? If so, can you use SetVector<Operation *, SmallVector<Operation *>, SmallPtrSet<Operation *>>. Otherwise, it's fine as it is.

80

Erasing from a set vector is linear time

88

It seems like instead of using a SetVector and maintaining a separate list of sorted ops, you can use a DenseSet (like before) and accept a MutableArrayRef<Operation *> and swap sorted operations to the start of the list like in the other implementation.

96

nit: can you swap the definition of this function to be before the other? To match the header file.

mlir/test/Transforms/test-toposort.mlir

15

Can you add a test where a subrange of ops are sorted,or even a subset?

[springerm]

address comments

mlir/lib/Transforms/Utils/TopologicalSortUtils.cpp

88

Switched to MutableArrayRef. The sort may no longer be stable though. (That's not a property that I need.) It's an array, not a linked list where I can move stuff around in a cheap way. Also mentioning that in the docs.

mlir/test/Transforms/test-toposort.mlir

15

I added a test with a subrange. Due to the swapping, the sort is no longer stable. Although if it's already sorted, nothing should change. (But it's not a property that is mentioned in the docs.) Should I add a test for that as well?

[rriddle]

mlir/test/lib/Transforms/TestTopologicalSort.cpp

10

This looks off.

[springerm]

address comments

[Mogball]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

101

Can you clarify that the above sorts are stable?

mlir/test/Transforms/test-toposort.mlir

17

Can you remove the underscores from these attribute names? They aren't conflicting with anything

[Mogball]

LGTM overall. If you don't need the sort to be stable, that's fine, but please make sure it's obvious that the other sorts are stable.

[springerm]

mlir/include/mlir/Transforms/TopologicalSortUtils.h

101

Added a comment to one of the above functions (the other one already had one).

[springerm]

address comments

mlir/include/mlir/Transforms/TopologicalSortUtils.h

 /// %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.
+/// terminator if it has one. This sort is stable.
 bool sortTopologically(
     Block *block,
     function_ref<bool(Value, Operation *)> isOperandReady = nullptr);
 
+/// Compute a topological ordering of the given ops. All ops must belong to the
+/// specified block.
+///
+/// This sort is not stable.
+///
+/// Note: If the specified ops contain incomplete/interrupted SSA use-def
+/// chains, the result may not actually be a topological sorting with respect to
+/// the entire program.
+bool computeTopologicalSorting(
+    Block *block, MutableArrayRef<Operation *> ops,
+    function_ref<bool(Value, Operation *)> isOperandReady = nullptr);
+
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_TOPOLOGICALSORTUTILS_H

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"
+#include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
-bool mlir::sortTopologically(
-    Block *block, llvm::iterator_range<Block::iterator> ops,
+/// Return `true` if the given operation is ready to be scheduled.
+static bool isOpReady(Block *block, Operation *op,
+                      DenseSet<Operation *> &unscheduledOps,
                       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))
+    if (isOperandReady && isOperandReady(value, op))
       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)
+    if (ancestor == op)
       return true;
     // - its ancestor in the block is scheduled.
     return !unscheduledOps.contains(ancestor);
   };
 
+  // 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();
+  });
+  return !readyToSchedule.wasInterrupted();
+}
+
+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();
+
   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())
+      if (!isOpReady(block, &op, unscheduledOps, isOperandReady))
         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);
 }
+
+bool mlir::computeTopologicalSorting(
+    Block *block, MutableArrayRef<Operation *> 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) {
+    assert(op->getBlock() == block && "op must belong to block");
+    unscheduledOps.insert(op);
+  }
+
+  unsigned nextScheduledOp = 0;
+
+  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 (swap it with the op at `nextScheduledOp`).
+    for (unsigned i = nextScheduledOp; i < ops.size(); ++i) {
+      if (!isOpReady(block, ops[i], unscheduledOps, isOperandReady))
+        continue;
+
+      // Schedule the operation by moving it to the start.
+      unscheduledOps.erase(ops[i]);
+      std::swap(ops[i], ops[nextScheduledOp]);
+      scheduledAtLeastOnce = true;
+      ++nextScheduledOp;
+    }
+
+    // If no operations were scheduled, just schedule the first op and continue.
+    if (!scheduledAtLeastOnce) {
+      allOpsScheduled = false;
+      unscheduledOps.erase(ops[nextScheduledOp++]);
+    }
+  }
+
+  return allOpsScheduled;
+}

mlir/test/Transforms/test-toposort.mlir

 // RUN: mlir-opt -topological-sort %s | FileCheck %s
+// RUN: mlir-opt -test-topological-sort-analysis %s | FileCheck %s -check-prefix=CHECK-ANALYSIS
 
 // Test producer is after user.
 // CHECK-LABEL: test.graph_region
-test.graph_region {
+// CHECK-ANALYSIS-LABEL: test.graph_region
+test.graph_region attributes{"root"} {
   // 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
+
+  // CHECK-ANALYSIS-NEXT: test.foo{{.*}} {pos = 0
+  // CHECK-ANALYSIS-NEXT: test.bar{{.*}} {pos = 2
+  // CHECK-ANALYSIS-NEXT: test.baz{{.*}} {pos = 1
+  %0 = "test.foo"() {selected} : () -> i32
+  "test.bar"(%1, %0) {selected} : (i32, i32) -> ()
+  %1 = "test.baz"() {selected} : () -> i32
 }
 
 // Test cycles.
 // CHECK-LABEL: test.graph_region
-test.graph_region {
+// CHECK-ANALYSIS-LABEL: test.graph_region
+test.graph_region attributes{"root"} {
   // CHECK-NEXT: test.d
   // CHECK-NEXT: test.a
   // CHECK-NEXT: test.c
   // CHECK-NEXT: test.b
-  %2 = "test.c"(%1) : (i32) -> i32
+
+  // CHECK-ANALYSIS-NEXT: test.c{{.*}} {pos = 0
+  // CHECK-ANALYSIS-NEXT: test.b{{.*}} : (
+  // CHECK-ANALYSIS-NEXT: test.a{{.*}} {pos = 2
+  // CHECK-ANALYSIS-NEXT: test.d{{.*}} {pos = 1
+  %2 = "test.c"(%1) {selected} : (i32) -> i32
   %1 = "test.b"(%0, %2) : (i32, i32) -> i32
-  %0 = "test.a"(%3) : (i32) -> i32
-  %3 = "test.d"() : () -> i32
+  %0 = "test.a"(%3) {selected} : (i32) -> i32
+  %3 = "test.d"() {selected} : () -> i32
 }
 
 // Test block arguments.
 // CHECK-LABEL: test.graph_region
 test.graph_region {
 // CHECK-NEXT: (%{{.*}}:
 ^entry(%arg0: i32):
   // CHECK-NEXT: test.foo

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestCommutativityUtils.cpp
   TestConstantFold.cpp
   TestControlFlowSink.cpp
   TestInlining.cpp
   TestIntRangeInference.cpp
+  TestTopologicalSort.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
   MLIRAnalysis

mlir/test/lib/Transforms/TestTopologicalSort.cpp

+//===- TestTopologicalSort.cpp - Pass to test topological sort analysis ---===//
+//
+// 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/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/TopologicalSortUtils.h"
+
+using namespace mlir;
+
+namespace {
+struct TestTopologicalSortAnalysisPass
+    : public PassWrapper<TestTopologicalSortAnalysisPass,
+                         OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTopologicalSortAnalysisPass)
+
+  StringRef getArgument() const final {
+    return "test-topological-sort-analysis";
+  }
+  StringRef getDescription() const final {
+    return "Test topological sorting of ops";
+  }
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    OpBuilder builder(op->getContext());
+
+    op->walk([&](Operation *root) {
+      if (!root->hasAttr("root"))
+        return WalkResult::advance();
+
+      assert(root->getNumRegions() == 1 && root->getRegion(0).hasOneBlock() &&
+             "expected one block");
+      Block *block = &root->getRegion(0).front();
+      SmallVector<Operation *> selectedOps;
+      block->walk([&](Operation *op) {
+        if (op->hasAttr("selected"))
+          selectedOps.push_back(op);
+      });
+
+      computeTopologicalSorting(block, selectedOps);
+      for (const auto &it : llvm::enumerate(selectedOps))
+        it.value()->setAttr("pos", builder.getIndexAttr(it.index()));
+
+      return WalkResult::advance();
+    });
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestTopologicalSortAnalysisPass() {
+  PassRegistration<TestTopologicalSortAnalysisPass>();
+}
+} // namespace test
+} // namespace mlir

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

 void registerTestPDLByteCodePass();
 void registerTestPDLLPasses();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
 void registerTestSliceAnalysisPass();
 void registerTestTensorTransforms();
 void registerTestTilingInterface();
+void registerTestTopologicalSortAnalysisPass();
 void registerTestTransformDialectInterpreterPass();
 void registerTestVectorLowerings();
 void registerTestNvgpuLowerings();
 } // namespace test
 } // namespace mlir
 
 namespace test {
 void registerTestDialect(DialectRegistry &);
#endif
   mlir::test::registerTestPadFusion();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestPDLLPasses();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();
   mlir::test::registerTestSliceAnalysisPass();
   mlir::test::registerTestTensorTransforms();
   mlir::test::registerTestTilingInterface();
+  mlir::test::registerTestTopologicalSortAnalysisPass();
   mlir::test::registerTestTransformDialectInterpreterPass();
   mlir::test::registerTestVectorLowerings();
   mlir::test::registerTestNvgpuLowerings();
 }
 #endif
 
 int main(int argc, char **argv) {
   registerAllPasses();