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

[Markus Böck]

https://github.com/zero9178 updated https://github.com/llvm/llvm-project/pull/67544

>From 3356715a9910defef4f2fb915ce9dadcbf6b0d4e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Markus=20B=C3=B6ck?= <markus.bock+llvm at nextsilicon.com>
Date: Wed, 27 Sep 2023 13:47:56 +0200
Subject: [PATCH 1/2] [mlir][cf] Fix invalid transformation when value is used
 in a subregion

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.
---
 mlir/lib/Transforms/Utils/CFGToSCF.cpp        |  8 +++-
 .../Conversion/ControlFlowToSCF/test.mlir     | 47 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index e7bf6628ccbd7e0..dce1e58c74926c8 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 it 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..ecf3e8a2119aa32 100644
--- a/mlir/test/Conversion/ControlFlowToSCF/test.mlir
+++ b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
@@ -710,3 +710,50 @@ 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]])
+

>From 2c9e4f392bd7b046f61091e7125de38ec66a5e12 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Markus=20B=C3=B6ck?= <markus.bock+llvm at nextsilicon.com>
Date: Wed, 27 Sep 2023 13:50:53 +0200
Subject: [PATCH 2/2] super minor comment and whitespace change

---
 mlir/lib/Transforms/Utils/CFGToSCF.cpp          | 4 ++--
 mlir/test/Conversion/ControlFlowToSCF/test.mlir | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index dce1e58c74926c8..def91a5593df410 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -721,8 +721,8 @@ transformToReduceLoop(Block *loopHeader, Block *exitBlock,
       Value blockArgument;
       for (OpOperand &use : llvm::make_early_inc_range(value.getUses())) {
         // 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 it does not
-        // escape the loop through this use.
+        // 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();
diff --git a/mlir/test/Conversion/ControlFlowToSCF/test.mlir b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
index ecf3e8a2119aa32..f6e7fb265790b78 100644
--- a/mlir/test/Conversion/ControlFlowToSCF/test.mlir
+++ b/mlir/test/Conversion/ControlFlowToSCF/test.mlir
@@ -756,4 +756,3 @@ func.func @nested_region_outside_loop_use() {
 
 // CHECK: scf.execute_region
 // CHECK-NEXT: "test.foo"(%[[RES]])
-