[Mlir-commits] [mlir] [mlir] Fix edge case in move operation dependencies utility (PR #131081)

[Ian Wood]

https://github.com/IanWood1 created https://github.com/llvm/llvm-project/pull/131081

This change simplifies `moveOperationDependencies` and handles the edge case where if the given op directly uses `insertionPoint`, the op's dependencies can still be moved.





>From 10a44686d4da24c5f34890767c346b439d1545a3 Mon Sep 17 00:00:00 2001
From: Ian Wood <ianwood2024 at u.northwestern.edu>
Date: Thu, 13 Mar 2025 10:27:41 -0700
Subject: [PATCH] Simplify and handle edge case

Signed-off-by: Ian Wood <ianwood2024 at u.northwestern.edu>
---
 mlir/lib/Transforms/Utils/RegionUtils.cpp     | 33 +++----------------
 mlir/test/Transforms/move-operation-deps.mlir | 27 +++++++++++++++
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index 18e079d153161..84a6182ae2168 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -1079,34 +1079,11 @@ LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
                                        "insertion point does not dominate op");
   }
 
-  // Find the backward slice of operation for each `Value` the operation
-  // depends on. Prune the slice to only include operations not already
-  // dominated by the `insertionPoint`
-  BackwardSliceOptions options;
-  options.inclusive = false;
-  options.omitUsesFromAbove = false;
-  // Since current support is to only move within a same basic block,
-  // the slices dont need to look past block arguments.
-  options.omitBlockArguments = true;
-  options.filter = [&](Operation *sliceBoundaryOp) {
-    return !dominance.properlyDominates(sliceBoundaryOp, insertionPoint);
-  };
-  llvm::SetVector<Operation *> slice;
-  getBackwardSlice(op, &slice, options);
-
-  // If the slice contains `insertionPoint` cannot move the dependencies.
-  if (slice.contains(insertionPoint)) {
-    return rewriter.notifyMatchFailure(
-        op,
-        "cannot move dependencies before operation in backward slice of op");
-  }
-
-  // We should move the slice in topological order, but `getBackwardSlice`
-  // already does that. So no need to sort again.
-  for (Operation *op : slice) {
-    rewriter.moveOpBefore(op, insertionPoint);
-  }
-  return success();
+  auto operandRange = op->getOperands();
+  SetVector<Value> defs(operandRange.begin(), operandRange.end());
+  getUsedValuesDefinedAbove(op->getRegions(), defs);
+  return moveValueDefinitions(rewriter, defs.takeVector(), insertionPoint,
+                              dominance);
 }
 
 LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
diff --git a/mlir/test/Transforms/move-operation-deps.mlir b/mlir/test/Transforms/move-operation-deps.mlir
index aa7b5dc2a240a..68f287ea2d910 100644
--- a/mlir/test/Transforms/move-operation-deps.mlir
+++ b/mlir/test/Transforms/move-operation-deps.mlir
@@ -27,6 +27,33 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// Check that `moved_op` gets moved even if `foo` is defined by `before`.
+func.func @move_with_direct_use() -> f32 {
+  %0 = "before"() : () -> (f32)
+  %1 = "moved_op"() : () -> (f32)
+  %2 = "foo"(%0, %1) : (f32, f32) -> (f32)
+  return %2 : f32
+}
+// CHECK-LABEL: func @move_with_direct_use()
+//       CHECK:   %[[MOVED:.+]] = "moved_op"
+//       CHECK:   %[[BEFORE:.+]] = "before"
+//       CHECK:   %[[FOO:.+]] = "foo"(%[[BEFORE]], %[[MOVED]])
+//       CHECK:   return %[[FOO]]
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0 : !transform.any_op {transform.readonly}) {
+    %op1 = transform.structured.match ops{["foo"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+    %op2 = transform.structured.match ops{["before"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+    transform.test.move_operand_deps %op1 before %op2
+        : !transform.any_op, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
 // Move operands that are implicitly captured by the op
 func.func @move_region_dependencies() -> f32 {
   %0 = "before"() : () -> (f32)