[Mlir-commits] [mlir] 30fe876 - [mlir][cfg-to-scf] Fix invalid transformation when value is used in a subregion (#67544)

[llvmlistbot at llvm.org]

Author: Markus Böck
Date: 2023-09-27T14:02:54+02:00
New Revision: 30fe8762446ca91a34174ab9c5dde34bde4d4394

URL: https://github.com/llvm/llvm-project/commit/30fe8762446ca91a34174ab9c5dde34bde4d4394
DIFF: https://github.com/llvm/llvm-project/commit/30fe8762446ca91a34174ab9c5dde34bde4d4394.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/CFGToSCF.cpp
    mlir/test/Conversion/ControlFlowToSCF/test.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index e7bf6628ccbd7e0..def91a5593df410 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -720,7 +720,13 @@ transformToReduceLoop(Block *loopHeader, Block *exitBlock,
     auto checkValue = [&](Value value) {
       Value blockArgument;
       for (OpOperand &use : llvm::make_early_inc_range(value.getUses())) {
-        if (loopBlocks.contains(use.getOwner()->getBlock()))
+        // Go through all the parent blocks and find the one part of the region
+        // of the loop. If the block is part of the loop, then the value does
+        // not escape the loop through this use.
+        Block *currBlock = use.getOwner()->getBlock();
+        while (currBlock && currBlock->getParent() != loopHeader->getParent())
+          currBlock = currBlock->getParentOp()->getBlock();
+        if (loopBlocks.contains(currBlock))
           continue;
 
         // Block argument is only created the first time it is required.

diff  --git a/mlir/test/Conversion/ControlFlowToSCF/test.mlir b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
index ec35b4bc944213b..f6e7fb265790b78 100644
--- a/mlir/test/Conversion/ControlFlowToSCF/test.mlir
+++ b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
@@ -710,3 +710,49 @@ func.func @nested_region() {
 // CHECK-NEXT: scf.yield
 // CHECK-NEXT: }
 // CHECK-NEXT: return
+
+// -----
+
+func.func @nested_region_inside_loop_use() {
+  cf.br ^bb1
+
+^bb1:
+  %3 = "test.test1"() : () -> i32
+  scf.execute_region {
+    "test.foo"(%3) : (i32) -> ()
+    scf.yield
+  }
+  cf.br ^bb1
+}
+
+// CHECK-LABEL: func @nested_region_inside_loop_use
+// CHECK: scf.while
+// CHECK-NEXT: %[[DEF:.*]] = "test.test1"()
+// CHECK-NEXT: scf.execute_region
+// CHECK-NEXT: "test.foo"(%[[DEF]])
+
+// -----
+
+func.func @nested_region_outside_loop_use() {
+  cf.br ^bb1
+
+^bb1:
+  %3 = "test.test1"() : () -> i32
+  %cond = "test.test2"() : () -> i1
+  cf.cond_br %cond, ^bb1, ^bb2
+
+^bb2:
+  scf.execute_region {
+    "test.foo"(%3) : (i32) -> ()
+    scf.yield
+  }
+  return
+}
+
+// CHECK-LABEL: func @nested_region_outside_loop_use
+// CHECK: %[[RES:.*]] = scf.while
+// CHECK: %[[DEF:.*]] = "test.test1"()
+// CHECK: scf.condition(%{{.*}}) %[[DEF]]
+
+// CHECK: scf.execute_region
+// CHECK-NEXT: "test.foo"(%[[RES]])