[Mlir-commits] [mlir] [mlir][CFGToSCF] Fix crash when region contains unconvertible multi-successor op (PR #183935)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`transformCFGToSCF` would crash with a use-list assertion when it encountered an op like `spirv.BranchConditional` that implements `BranchOpInterface` (passing the existing precondition checks) but is not handled by `createStructuredBranchRegionOp`. The algorithm mutated the IR significantly before discovering the op was unsupported, leaving it in a corrupt half-transformed state that triggered the assertion on teardown.

Fix by adding `canConvertBranchOp` to `CFGToSCFInterface` (default: accept all ops) and calling it inside `checkTransformationPreconditions` for every block with more than one successor, before any IR modifications are made. `ControlFlowToSCFTransformation` overrides the method to accept only `cf.cond_br` and `cf.switch`.

Add a regression test covering `spirv.BranchConditional` inside a `spirv.mlir.loop` region nested in a `func.func`.

Fixes #<!-- -->173566

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


5 Files Affected:

- (modified) mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h (+4) 
- (modified) mlir/include/mlir/Transforms/CFGToSCF.h (+6) 
- (modified) mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp (+4) 
- (modified) mlir/lib/Transforms/Utils/CFGToSCF.cpp (+18-3) 
- (added) mlir/test/Conversion/ControlFlowToSCF/invalid.mlir (+18) 


``````````diff
diff --git a/mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h b/mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h
index dab8e7d51f24a..44e067f5031a0 100644
--- a/mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h
+++ b/mlir/include/mlir/Conversion/ControlFlowToSCF/ControlFlowToSCF.h
@@ -24,6 +24,10 @@ class Pass;
 /// operations to SCF Dialect operations.
 class ControlFlowToSCFTransformation : public CFGToSCFInterface {
 public:
+  /// Returns true only for `cf.cond_br` and `cf.switch`, the two multi-
+  /// successor ops this transformation knows how to convert.
+  bool canConvertBranchOp(Operation *op) override;
+
   /// Creates an `scf.if` op if `controlFlowCondOp` is a `cf.cond_br` op or
   /// an `scf.index_switch` if `controlFlowCondOp` is a `cf.switch`.
   /// Returns failure otherwise.
diff --git a/mlir/include/mlir/Transforms/CFGToSCF.h b/mlir/include/mlir/Transforms/CFGToSCF.h
index e3d4b6f3d7e7e..b18b1e610a8a1 100644
--- a/mlir/include/mlir/Transforms/CFGToSCF.h
+++ b/mlir/include/mlir/Transforms/CFGToSCF.h
@@ -96,6 +96,12 @@ class CFGToSCFInterface {
   /// control-flow paths where an SSA-value is undefined.
   virtual Value getUndefValue(Location loc, OpBuilder &builder, Type type) = 0;
 
+  /// Returns true if this operation (which has >1 successors) can be
+  /// converted to structured control flow by `createStructuredBranchRegionOp`.
+  /// Called during precondition checking, before any IR modifications.
+  /// Default implementation accepts all ops.
+  virtual bool canConvertBranchOp(Operation *op) { return true; }
+
   /// Creates a return-like terminator indicating unreachable.
   /// This is required when the transformation encounters a statically known
   /// infinite loop. Since structured control flow ops are not terminators,
diff --git a/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp b/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
index 5ac838cad6f0f..8b7bc1a221276 100644
--- a/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
+++ b/mlir/lib/Conversion/ControlFlowToSCF/ControlFlowToSCF.cpp
@@ -27,6 +27,10 @@ namespace mlir {
 
 using namespace mlir;
 
+bool ControlFlowToSCFTransformation::canConvertBranchOp(Operation *op) {
+  return isa<cf::CondBranchOp, cf::SwitchOp>(op);
+}
+
 FailureOr<Operation *>
 ControlFlowToSCFTransformation::createStructuredBranchRegionOp(
     OpBuilder &builder, Operation *controlFlowCondOp, TypeRange resultTypes,
diff --git a/mlir/lib/Transforms/Utils/CFGToSCF.cpp b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
index dbde75e17b3b8..e1816cfcd2dc4 100644
--- a/mlir/lib/Transforms/Utils/CFGToSCF.cpp
+++ b/mlir/lib/Transforms/Utils/CFGToSCF.cpp
@@ -1223,7 +1223,8 @@ static ReturnLikeExitCombiner createSingleExitBlocksForReturnLike(
 
 /// Checks all preconditions of the transformation prior to any transformations.
 /// Returns failure if any precondition is violated.
-static LogicalResult checkTransformationPreconditions(Region &region) {
+static LogicalResult
+checkTransformationPreconditions(Region &region, CFGToSCFInterface &interface) {
   for (Block &block : region.getBlocks())
     if (block.hasNoPredecessors() && !block.isEntryBlock())
       return block.front().emitOpError(
@@ -1267,7 +1268,21 @@ static LogicalResult checkTransformationPreconditions(Region &region) {
     }
     return WalkResult::advance();
   });
-  return failure(result.wasInterrupted());
+  if (result.wasInterrupted())
+    return failure();
+
+  // Verify all multi-successor terminators are convertible before touching IR.
+  for (Block &block : region.getBlocks()) {
+    if (block.getNumSuccessors() <= 1)
+      continue;
+    Operation *terminator = block.getTerminator();
+    if (!interface.canConvertBranchOp(terminator)) {
+      terminator->emitOpError(
+          "Cannot convert unknown control flow op to structured control flow");
+      return failure();
+    }
+  }
+  return success();
 }
 
 FailureOr<bool> mlir::transformCFGToSCF(Region &region,
@@ -1276,7 +1291,7 @@ FailureOr<bool> mlir::transformCFGToSCF(Region &region,
   if (region.empty() || region.hasOneBlock())
     return false;
 
-  if (failed(checkTransformationPreconditions(region)))
+  if (failed(checkTransformationPreconditions(region, interface)))
     return failure();
 
   DenseMap<Type, Value> typedUndefCache;
diff --git a/mlir/test/Conversion/ControlFlowToSCF/invalid.mlir b/mlir/test/Conversion/ControlFlowToSCF/invalid.mlir
new file mode 100644
index 0000000000000..ee2cb6cc5f762
--- /dev/null
+++ b/mlir/test/Conversion/ControlFlowToSCF/invalid.mlir
@@ -0,0 +1,18 @@
+// RUN: mlir-opt --lift-cf-to-scf -split-input-file --verify-diagnostics %s
+
+// Regression test: spirv.BranchConditional inside a spirv.mlir.loop region
+// must produce a clean diagnostic rather than an assertion crash.
+
+func.func @spirv_loop_no_crash(%cond: i1) {
+  spirv.mlir.loop {
+    spirv.Branch ^bb1
+  ^bb1:
+    // expected-error at below {{Cannot convert unknown control flow op to structured control flow}}
+    spirv.BranchConditional %cond, ^bb2, ^bb3
+  ^bb2:
+    spirv.Branch ^bb1
+  ^bb3:
+    spirv.mlir.merge
+  }
+  return
+}

``````````

</details>


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