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

[Ian Wood]

================
@@ -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()) {
----------------
IanWood1 wrote:

You can use this overload to not have to loop over each region: https://github.com/llvm/llvm-project/blob/8f1427d26929bdec364f46a9043303e6b1ed4200/mlir/include/mlir/Transforms/RegionUtils.h#L51-L54

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