Simplify slice computation.

Signed-off-by: MaheshRavishankar <[email protected]>

Author: MaheshRavishankar

mlir/lib/Transforms/Utils/RegionUtils.cpp

@@ -1083,7 +1083,7 @@ LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
   // depends on. Prune the slice to only include operations not already
   // dominated by the `insertionPoint`
   BackwardSliceOptions options;
-  options.inclusive = true;
+  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.
@@ -1092,20 +1092,7 @@ LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
     return !dominance.properlyDominates(sliceBoundaryOp, insertionPoint);
   };
   llvm::SetVector<Operation *> slice;
-
-  // Get the defined slice for operands.
-  for (Value operand : op->getOperands()) {
-    getBackwardSlice(operand, &slice, options);
-  }
-  auto regions = op->getRegions();
-  if (!regions.empty()) {
-    // If op has region, get the defined slice for all captured values.
-    llvm::SetVector<Value> capturedVals;
-    mlir::getUsedValuesDefinedAbove(regions, capturedVals);
-    for (Value value : capturedVals) {
-      getBackwardSlice(value, &slice, options);
-    }
-  }
+  getBackwardSlice(op, &slice, options);
 
   // If the slice contains `insertionPoint` cannot move the dependencies.
   if (slice.contains(insertionPoint)) {
@@ -1114,8 +1101,8 @@ LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
         "cannot move dependencies before operation in backward slice of op");
   }
 
-  // Sort the slice topologically, and move in topological order.
-  mlir::topologicalSort(slice);
+  // 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);
   }

mlir/test/Transforms/move-operation-deps.mlir

@@ -125,6 +125,42 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// Current implementation omits following basic block argument when
+// computing slices. Verify that this gives expected result.
+func.func @ignore_basic_block_arguments() -> f32 {
+^bb0():
+  %8 = "test"() : () -> (f32)
+  return %8: f32
+^bb1(%bbArg : f32):
+  %0 = "before"() : () -> (f32)
+  %1 = "moved_op"() ({
+    "yield"(%bbArg) : (f32) -> ()
+  }) : () -> (f32)
+  %2 = "foo"() ({
+    "yield"(%1) : (f32) -> ()
+  }) : () -> (f32)
+  return %2 : f32
+}
+
+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
+  }
+}
+// CHECK-LABEL: func @ignore_basic_block_arguments()
+//       CHECK:   %[[MOVED_1:.+]] = "moved_op"
+//       CHECK:   "before"
+//       CHECK:   %[[FOO:.+]] = "foo"
+//       CHECK:   return %[[FOO]]
+
+// -----
+
 // Fail when the "before" operation is part of the operation slice.
 func.func @do_not_move_slice() -> f32 {
   %0 = "before"() : () -> (f32)