[Mlir-commits] [mlir] [MLIR] Prevent CF to SCF transform from erasing blocks still in use (PR #206310)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Martin P (pechy)

<details>
<summary>Changes</summary>

While merging blocks, `transformToStructuredCFBranches` may try to erase blocks that still have uses. This may happen when the block has self-loop or when it has 1 successor, but the successor has multiple predecessors. This probably happens only with unreachable blocks in the function.

Note that running a DCE pass before CF to SCF transform would remove the triggering assert in the lit test.

Fixes #<!-- -->206086.

---
Full diff: https://github.com/llvm/llvm-project/pull/206310.diff


2 Files Affected:

- (modified) mlir/lib/Transforms/Utils/CFGToSCF.cpp (+7-1) 
- (added) mlir/test/Conversion/ControlFlowToSCF/dead_code.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index 0cc6bb43ddd38..3a9b979ca940e 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -950,8 +950,14 @@ static FailureOr<SmallVector<Block *>> transformToStructuredCFBranches(
     return SmallVector<Block *>{};
 
   if (regionEntry->getNumSuccessors() == 1) {
-    // Single successor we can just splice together.
+    // Single successor we can splice together, if they are not self-loops and
+    // the successor's only predecessor is `regionEntry`.
     Block *successor = regionEntry->getSuccessor(0);
+    if (successor == regionEntry ||
+        successor->getUniquePredecessor() != regionEntry) {
+      return SmallVector<Block *>{};
+    }
+
     for (auto &&[oldValue, newValue] : llvm::zip(
              successor->getArguments(), getSuccessorOperands(regionEntry, 0)))
       oldValue.replaceAllUsesWith(newValue);
diff --git a/mlir/test/Conversion/ControlFlowToSCF/dead_code.mlir b/mlir/test/Conversion/ControlFlowToSCF/dead_code.mlir
new file mode 100644
index 0000000000000..a6f13e92e1a3d
--- /dev/null
+++ b/mlir/test/Conversion/ControlFlowToSCF/dead_code.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt --lift-cf-to-scf -split-input-file %s
+
+// Regression test: regions that have 1 successor, but that successor has more
+// than 1 predecessor cannot be be simply merged and the successor erased. No
+// assertion crash must happen.
+
+module {
+  func.func @single_return_statement(%arg0: i32) -> i32 {
+    return %arg0 : i32
+  ^bb1:  // pred: ^bb1
+    %true = arith.constant true
+    cf.cond_br %true, ^bb2, ^bb1
+  ^bb2:  // pred: ^bb1
+    return %arg0 : i32
+  }
+}
+

``````````

</details>


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