[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;
----------------
jtuyls wrote:

Yes, indeed, I liked the clarity, but removed it now.

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