[Mlir-commits] [mlir] [mlir][CFGToSCF] Fix crash when encountering unknown control flow ops (PR #184103)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/184103

>From 9c4269dc3ada8085ae3a372878ec683d64068516 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 2 Mar 2026 03:04:56 -0800
Subject: [PATCH] [mlir][CFGToSCF] Fix crash when encountering unknown control
 flow ops

When transformToStructuredCFBranches encountered a control flow op not
handled by the CFGToSCFInterface (e.g., spirv.BranchConditional with
--lift-cf-to-scf), it correctly emitted an error and returned failure.
However, blocks had already been moved from the parent region into
temporary local Region objects before the failure was detected.

When those temporary Region objects went out of scope, their destructor
tried to destroy the contained blocks. But those blocks still had live
predecessor references from the parent region (the regionEntry's
terminator still pointed to them), causing an assertion failure:
  use_empty() && "Cannot destroy a value that still has uses\!"

Fix: on failure from createStructuredBranchRegionOp, move the blocks
from the temporary conditionalRegions back into the parent region
before returning failure. This restores IR consistency and allows the
Region destructor to run safely.

Fixes #120883
---
 mlir/lib/Transforms/Utils/CFGToSCF.cpp        | 14 +++++++-
 .../ControlFlowToSCF/unknown-cf-op.mlir       | 36 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Conversion/ControlFlowToSCF/unknown-cf-op.mlir

diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index dbde75e17b3b8..9e15d8555135a 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -1156,8 +1156,20 @@ static FailureOr<SmallVector<Block *>> transformToStructuredCFBranches(
     FailureOr<Operation *> result = interface.createStructuredBranchRegionOp(
         opBuilder, regionEntry->getTerminator(),
         continuation->getArgumentTypes(), conditionalRegions);
-    if (failed(result))
+    if (failed(result)) {
+      // Blocks were moved from the parent region into conditionalRegions before
+      // calling createStructuredBranchRegionOp. On failure, move them back to
+      // avoid use-after-free crashes: the moved blocks may still be referenced
+      // as successors by blocks remaining in the parent region, so destroying
+      // conditionalRegions with live uses would trigger an assertion.
+      // This patching does not undo the change, it barely makes it so that the
+      // pass can gracefully fail instead of crashing.
+      Region *parentRegion = regionEntry->getParent();
+      for (Region &conditionalRegion : conditionalRegions)
+        parentRegion->getBlocks().splice(parentRegion->getBlocks().end(),
+                                         conditionalRegion.getBlocks());
       return failure();
+    }
     structuredCondOp = *result;
     regionEntry->getTerminator()->erase();
   }
diff --git a/mlir/test/Conversion/ControlFlowToSCF/unknown-cf-op.mlir b/mlir/test/Conversion/ControlFlowToSCF/unknown-cf-op.mlir
new file mode 100644
index 0000000000000..957cc06761e41
--- /dev/null
+++ b/mlir/test/Conversion/ControlFlowToSCF/unknown-cf-op.mlir
@@ -0,0 +1,36 @@
+// RUN: mlir-opt --lift-cf-to-scf %s -verify-diagnostics -split-input-file
+
+// Verify that --lift-cf-to-scf does not crash when it encounters an
+// unknown control-flow op (not cf.cond_br or cf.switch) in a multi-block
+// region. Instead it should emit a clean error.
+// See: https://github.com/llvm/llvm-project/issues/120883
+
+// The spirv.BranchConditional op implements BranchOpInterface but is not
+// handled by ControlFlowToSCFTransformation::createStructuredBranchRegionOp.
+// Previously, blocks were moved into temporary Region objects before the
+// failure was detected, and the Region destructor would assert because the
+// moved blocks still had live predecessor references.
+
+func.func private @unknown_cf_in_loop(%arg0: f32) -> (i32, i32) {
+  %cst0_i32 = spirv.Constant 0 : i32
+  %cst-1_i32 = spirv.Constant -1 : i32
+  %0 = spirv.Variable : !spirv.ptr<i32, Function>
+  %1 = spirv.Variable : !spirv.ptr<i32, Function>
+  spirv.mlir.loop {
+    spirv.Branch ^bb1(%cst0_i32, %cst-1_i32 : i32, i32)
+  ^bb1(%3: i32, %4: i32):  // 2 preds: ^bb0, ^bb2
+    %5 = spirv.SLessThan %3, %cst-1_i32 : i32
+    // expected-error at below {{'spirv.BranchConditional' op Cannot convert unknown control flow op to structured control flow}}
+    spirv.BranchConditional %5, ^bb2, ^bb3
+  ^bb2:  // pred: ^bb1
+    %6 = spirv.IAdd %3, %4 : i32
+    spirv.Store "Function" %0, %6 : i32
+    %7 = spirv.IAdd %3, %cst0_i32 : i32
+    spirv.Branch ^bb1(%7, %6 : i32, i32)
+  ^bb3:  // pred: ^bb1
+    spirv.mlir.merge
+  }
+  %2 = spirv.Load "Function" %1 : i32
+  %3 = spirv.Load "Function" %0 : i32
+  return %3, %2 : i32, i32
+}