From c74923e7e73c26ab947101be965023f3d6d0901f Mon Sep 17 00:00:00 2001 From: Scott Manley Date: Mon, 9 Feb 2026 07:36:37 -0800 Subject: [PATCH 1/3] [RegionUtils] replace uses in nested regions when isolating from above When making a region IsolatedFromAbove, replace uses in any region within the parent region, not just the immediate parent region. --- mlir/lib/Transforms/Utils/RegionUtils.cpp | 10 +++++++- .../Transforms/make-isolated-from-above.mlir | 25 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp index 5406a51d2ab7f..87125beafa065 100644 --- a/mlir/lib/Transforms/Utils/RegionUtils.cpp +++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp @@ -156,9 +156,17 @@ SmallVector mlir::makeRegionIsolatedFromAbove( // Create a mapping between the captured values and the new arguments added. IRMapping map; + Operation *regionParentOp = region.getParentOp(); auto replaceIfFn = [&](OpOperand &use) { - return use.getOwner()->getBlock()->getParent() == ®ion; + Operation *parentOp = use.getOwner()->getParentOp(); + while (parentOp) { + if (parentOp == regionParentOp) + return true; + parentOp = parentOp->getParentOp(); + } + return false; }; + for (auto [arg, capturedVal] : llvm::zip(newEntryBlockArgs.take_back(finalCapturedValues.size()), finalCapturedValues)) { diff --git a/mlir/test/Transforms/make-isolated-from-above.mlir b/mlir/test/Transforms/make-isolated-from-above.mlir index a9d4325944fd9..3b0084d6e0007 100644 --- a/mlir/test/Transforms/make-isolated-from-above.mlir +++ b/mlir/test/Transforms/make-isolated-from-above.mlir @@ -113,3 +113,28 @@ func.func @make_isolated_from_above_multiple_blocks(%arg0 : index, %arg1 : index // CLONE2-NEXT: cf.br ^bb1 // CLONE2: ^bb1: // CLONE2: "foo.yield"(%[[C0]], %[[C1]], %[[D0]], %[[D1]], %[[B0]]) + + +// ----- + +// CHECK-LABEL: func @make_isolated_from_above_nested_region +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<8xindex> +// CHECK: %[[C1:.+]] = arith.constant 1 : index +// CHECK: %[[C8:.+]] = arith.constant 8 : index +// CHECK: test.isolated_one_region_op %[[C1]], %[[ARG0]], %[[C8]] +// CHECK: ^bb0(%[[B0:[a-zA-Z0-9]+]]: index, %[[B1:[a-zA-Z0-9]+]]: memref<8xindex>, %[[B2:[a-zA-Z0-9]+]]: index) +// CHECK: scf.for %arg4 = %[[B0]] to %[[B2]] step %[[B0]] +// CHECK: memref.store %[[B0]], %[[B1]][%arg4] : memref<8xindex> +// CHECK: "foo.yield"() : () -> () + +func.func @make_isolated_from_above_nested_region(%arg0 : memref<8xindex>) { + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + "test.one_region_with_operands_op"() ({ + scf.for %arg1 = %c1 to %c8 step %c1 { + memref.store %c1, %arg0[%arg1] : memref<8xindex> + } + "foo.yield"() : () -> () + }) : () -> () + return +} From 9046c36df60fea5b30347ad9660abda2169cf6b0 Mon Sep 17 00:00:00 2001 From: Scott Manley Date: Mon, 9 Feb 2026 09:45:06 -0800 Subject: [PATCH 2/3] check parent region instead of parent op --- mlir/lib/Transforms/Utils/RegionUtils.cpp | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp index 87125beafa065..6f0b977c08588 100644 --- a/mlir/lib/Transforms/Utils/RegionUtils.cpp +++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp @@ -156,13 +156,12 @@ SmallVector mlir::makeRegionIsolatedFromAbove( // Create a mapping between the captured values and the new arguments added. IRMapping map; - Operation *regionParentOp = region.getParentOp(); auto replaceIfFn = [&](OpOperand &use) { - Operation *parentOp = use.getOwner()->getParentOp(); - while (parentOp) { - if (parentOp == regionParentOp) + Region *parentRegion = use.getOwner()->getParentRegion(); + while (parentRegion) { + if (parentRegion == ®ion) return true; - parentOp = parentOp->getParentOp(); + parentRegion = parentRegion->getParentRegion(); } return false; }; From ccf5c7dfff8dfc27a4a2ff4b741a2a7ffc9afa00 Mon Sep 17 00:00:00 2001 From: Scott Manley Date: Mon, 9 Feb 2026 10:27:53 -0800 Subject: [PATCH 3/3] use isAncestor --- mlir/lib/Transforms/Utils/RegionUtils.cpp | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp index 6f0b977c08588..a87cb15625130 100644 --- a/mlir/lib/Transforms/Utils/RegionUtils.cpp +++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp @@ -157,13 +157,7 @@ SmallVector mlir::makeRegionIsolatedFromAbove( // Create a mapping between the captured values and the new arguments added. IRMapping map; auto replaceIfFn = [&](OpOperand &use) { - Region *parentRegion = use.getOwner()->getParentRegion(); - while (parentRegion) { - if (parentRegion == ®ion) - return true; - parentRegion = parentRegion->getParentRegion(); - } - return false; + return region.isAncestor(use.getOwner()->getParentRegion()); }; for (auto [arg, capturedVal] :