[Mlir-commits] [mlir] [mlir] Extend moveValueDefinitions/moveOperationDependencies with cross-region support (PR #176343)

[Jorn Tuyls]

================
@@ -1096,31 +1096,116 @@ LogicalResult mlir::simplifyRegions(RewriterBase &rewriter,
 // Move operation dependencies
 //===---------------------------------------------------------------------===//
 
+/// Check if a block argument would dominate at the insertion point.
+/// Returns true if there's a dominance issue.
+static bool blockArgHasDominanceIssue(BlockArgument arg, Block *insertionBlock,
+                                      DominanceInfo &dominance) {
+  Block *argBlock = arg.getOwner();
+  return argBlock != insertionBlock &&
+         !dominance.dominates(argBlock, insertionBlock);
+}
+
+/// Check if moving operations in the slice before `insertionPoint` would break
+/// dominance due to block argument operands. Returns true if there's an issue.
+/// If `failingOp` is provided, it will be set to the first problematic op.
+///
+/// For operands defined by ops: either the defining op is in the slice (so
+/// dominance preserved), or it already dominates insertionPoint (otherwise it
+/// would be in the slice). So we only need to check block argument operands,
+/// both as direct operands and as values captured inside regions.
+static bool hasBlockArgDominanceIssue(const llvm::SetVector<Operation *> &slice,
+                                      Operation *insertionPoint,
+                                      DominanceInfo &dominance,
+                                      Operation **failingOp = nullptr) {
+  Block *insertionBlock = insertionPoint->getBlock();
+  for (Operation *op : slice) {
+    // Check direct operands.
+    for (Value operand : op->getOperands()) {
+      if (operand.getDefiningOp())
+        continue;
+      auto arg = cast<BlockArgument>(operand);
+      if (blockArgHasDominanceIssue(arg, insertionBlock, dominance)) {
+        if (failingOp)
+          *failingOp = op;
+        return true;
+      }
+    }
+
+    // Check block arguments captured inside regions. Process one region at a
+    // time to enable early exit without collecting values from all regions.
+    for (Region &region : op->getRegions()) {
+      SetVector<Value> capturedValues;
+      getUsedValuesDefinedAbove(region, region, capturedValues);
+      for (Value val : capturedValues) {
+        auto arg = dyn_cast<BlockArgument>(val);
+        if (!arg)
+          continue;
+        if (blockArgHasDominanceIssue(arg, insertionBlock, dominance)) {
+          if (failingOp)
+            *failingOp = op;
+          return true;
+        }
+      }
+    }
+  }
+  return false;
+}
+
+/// Check if any region between an operation and an ancestor block is
+/// isolated from above. If so, moving the operation out would break
+/// the isolation semantics.
+static bool hasIsolatedRegionBetween(Operation *op, Block *ancestorBlock) {
+  Region *opRegion = op->getParentRegion();
+  Region *ancestorRegion = ancestorBlock->getParent();
+
+  // If both are in the same region (just different blocks), no isolation.
+  if (opRegion == ancestorRegion)
+    return false;
+
+  // Walk up from the op's region to find if there's an isolated region
+  // between the op and the ancestor.
+  Region *region = opRegion;
+  while (region && region != ancestorRegion) {
+    Operation *parentOp = region->getParentOp();
+    if (!parentOp)
+      break;
+
+    if (parentOp->hasTrait<OpTrait::IsIsolatedFromAbove>())
+      return true;
+
+    region = parentOp->getParentRegion();
+  }
+  return false;
+}
+
 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 case where operation and insertion point are not in "
-            "the same basic block");
-  }
-  // If `insertionPoint` does not dominate `op`, do nothing
+  Block *insertionBlock = insertionPoint->getBlock();
+  Block *opBlock = op->getBlock();
+
+  // If `insertionPoint` does not dominate `op`, do nothing.
   if (!dominance.properlyDominates(insertionPoint, op)) {
     return rewriter.notifyMatchFailure(op,
                                        "insertion point does not dominate op");
   }
 
+  // For cross-block moves, verify we're not crossing an isolated region.
+  if (opBlock != insertionBlock) {
----------------
jtuyls wrote:

Thanks, removed it!

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