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

[Martin P]

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

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.

>From a3bc583e9daa994713a662dedb7aaf3f6502f71c Mon Sep 17 00:00:00 2001
From: Martin Petracek <3456763+pechy at users.noreply.github.com>
Date: Sun, 28 Jun 2026 07:50:34 +0200
Subject: [PATCH] [MLIR] Prevent CF to SCF transform from erasing blocks still
 in use

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.
---
 mlir/lib/Transforms/Utils/CFGToSCF.cpp          |  8 +++++++-
 .../Conversion/ControlFlowToSCF/dead_code.mlir  | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Conversion/ControlFlowToSCF/dead_code.mlir

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
+  }
+}
+