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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Transforms/Utils/CFGToSCF.cpp (+11-1) 
- (added) mlir/test/Conversion/ControlFlowToSCF/unknown-cf-op.mlir (+36) 


``````````diff
diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index dbde75e17b3b8..d7f341539ad0f 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -1156,8 +1156,18 @@ 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.
+      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
+}

``````````

</details>


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