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

[Jakub Kuderski]

================
@@ -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()) {
+    // If op has region, get the defined slice for all captured values.
+    llvm::SetVector<Value> capturedVals;
+    mlir::getUsedValuesDefinedAbove(regions, capturedVals);
+    for (auto value : capturedVals) {
+      getBackwardSlice(value, &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");
+  }
+
+  // Sort the slice topologically, ad move in topological order.
----------------
kuhar wrote:

```suggestion
  // Sort the slice topologically, and move in topological order.
```

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