[Mlir-commits] [mlir] [mlir] Add a utility method to move operation dependencies. (PR #129975)

[Ian Wood]

================
@@ -1054,3 +1057,65 @@ LogicalResult mlir::simplifyRegions(RewriterBase &rewriter,
   return success(eliminatedBlocks || eliminatedOpsOrArgs ||
                  mergedIdenticalBlocks || droppedRedundantArguments);
 }
+
+//===---------------------------------------------------------------------===//
+// Move operation dependencies
+//===---------------------------------------------------------------------===//
+
+LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,
+                                              Operation *op,
+                                              Operation *insertionPoint,
+                                              DominanceInfo &dominance) {
+  // Currently unsupported case where the op and insertion point are
+  // in different basic blocks.
+  if (op->getBlock() != insertionPoint->getBlock()) {
+    return rewriter.notifyMatchFailure(
+        op, "unsupported caes where operation and insertion point are not in "
+            "the sme basic block");
+  }
+
+  // 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 = true;
+  options.filter = [&](Operation *sliceBoundaryOp) {
+    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()) {
----------------
IanWood1 wrote:

I think this needs to be iterative. A chain of >1 "uses from above" will cause the transform to produce invalid IR:

```mlir
func.func @move_region_dependencies() -> f32 {
  %0 = "before"() : () -> (f32)
  %1 = "moved_op"() ({
    "yield"(%0) : (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
  }
}
```


`options.omitUsesFromAbove=false` could also work but you might hit the assert [here](https://github.com/llvm/llvm-project/blob/a6ccda28f7569c1a03620d7520de7cfadc11f4a5/mlir/lib/Analysis/SliceAnalysis.cpp#L109-L110). 

Maybe this assert could just be removed and assume the user handles it with the filter function. The other overload of `getBackwardSlice` doesn't do this check on the initial `Value` see: https://github.com/llvm/llvm-project/blob/a6ccda28f7569c1a03620d7520de7cfadc11f4a5/mlir/lib/Analysis/SliceAnalysis.cpp#L156-L157

https://github.com/llvm/llvm-project/pull/129975