Skip to content

Commit 30fe876

Browse files
authored
[mlir][cfg-to-scf] Fix invalid transformation when value is used in a subregion (#67544)
The current loop-reduce-form transformation incorrectly assumes that any value that is used in a block that isn't in the set of loop blocks is a block outside the loop. This is correct for a pure CFG but is incorrect if operations with subregions are present. In that case, a use may be in a subregion of an operation part of the loop and incorrectly deemed outside the loop. This would later lead to transformations with code that does not verify. This PR fixes that issue by checking the transitive parent block that is in the same region as the loop rather than the immediate parent block.
1 parent d338acf commit 30fe876

File tree

2 files changed

+53
-1
lines changed

2 files changed

+53
-1
lines changed

mlir/lib/Transforms/Utils/CFGToSCF.cpp

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -720,7 +720,13 @@ transformToReduceLoop(Block *loopHeader, Block *exitBlock,
720720
auto checkValue = [&](Value value) {
721721
Value blockArgument;
722722
for (OpOperand &use : llvm::make_early_inc_range(value.getUses())) {
723-
if (loopBlocks.contains(use.getOwner()->getBlock()))
723+
// Go through all the parent blocks and find the one part of the region
724+
// of the loop. If the block is part of the loop, then the value does
725+
// not escape the loop through this use.
726+
Block *currBlock = use.getOwner()->getBlock();
727+
while (currBlock && currBlock->getParent() != loopHeader->getParent())
728+
currBlock = currBlock->getParentOp()->getBlock();
729+
if (loopBlocks.contains(currBlock))
724730
continue;
725731

726732
// Block argument is only created the first time it is required.

mlir/test/Conversion/ControlFlowToSCF/test.mlir

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -710,3 +710,49 @@ func.func @nested_region() {
710710
// CHECK-NEXT: scf.yield
711711
// CHECK-NEXT: }
712712
// CHECK-NEXT: return
713+
714+
// -----
715+
716+
func.func @nested_region_inside_loop_use() {
717+
cf.br ^bb1
718+
719+
^bb1:
720+
%3 = "test.test1"() : () -> i32
721+
scf.execute_region {
722+
"test.foo"(%3) : (i32) -> ()
723+
scf.yield
724+
}
725+
cf.br ^bb1
726+
}
727+
728+
// CHECK-LABEL: func @nested_region_inside_loop_use
729+
// CHECK: scf.while
730+
// CHECK-NEXT: %[[DEF:.*]] = "test.test1"()
731+
// CHECK-NEXT: scf.execute_region
732+
// CHECK-NEXT: "test.foo"(%[[DEF]])
733+
734+
// -----
735+
736+
func.func @nested_region_outside_loop_use() {
737+
cf.br ^bb1
738+
739+
^bb1:
740+
%3 = "test.test1"() : () -> i32
741+
%cond = "test.test2"() : () -> i1
742+
cf.cond_br %cond, ^bb1, ^bb2
743+
744+
^bb2:
745+
scf.execute_region {
746+
"test.foo"(%3) : (i32) -> ()
747+
scf.yield
748+
}
749+
return
750+
}
751+
752+
// CHECK-LABEL: func @nested_region_outside_loop_use
753+
// CHECK: %[[RES:.*]] = scf.while
754+
// CHECK: %[[DEF:.*]] = "test.test1"()
755+
// CHECK: scf.condition(%{{.*}}) %[[DEF]]
756+
757+
// CHECK: scf.execute_region
758+
// CHECK-NEXT: "test.foo"(%[[RES]])

0 commit comments

Comments
 (0)