[Mlir-commits] [mlir] a0055a9 - [mlir][spirv] Verify yielded types match result of `loop`/`selection` (#179205)

[llvmlistbot at llvm.org]

Author: Igor Wodiany
Date: 2026-02-02T14:56:59Z
New Revision: a0055a9b47bb91586f606d5bc59c0ec7df1328de

URL: https://github.com/llvm/llvm-project/commit/a0055a9b47bb91586f606d5bc59c0ec7df1328de
DIFF: https://github.com/llvm/llvm-project/commit/a0055a9b47bb91586f606d5bc59c0ec7df1328de.diff

LOG: [mlir][spirv] Verify yielded types match result of `loop`/`selection` (#179205)

This patch extends the verifyRegion functions to ensure that values
returned by `spirv.mlir.loop` and `spirv.mlir.selection` are consistent
with values yielded with `spirv.mlir.merge` inside the region.

Added: 
    

Modified: 
    mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
    mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
index 4d0aedca27d42..5c2d538a00b17 100644
--- a/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
@@ -360,6 +360,14 @@ static bool hasOtherMerge(Region &region) {
   });
 }
 
+/// Returns true if types yielded by `spirv.mlir.merge` in the region match
+/// those returned by the `op`.
+static bool returnTypesMatch(Region &region, Operation *op) {
+  auto mergeOps = region.getOps<spirv::MergeOp>();
+  Operation *mergeOp = llvm::getSingleElement(mergeOps);
+  return llvm::equal(mergeOp->getOperandTypes(), op->getResultTypes());
+}
+
 LogicalResult LoopOp::verifyRegions() {
   auto *op = getOperation();
 
@@ -445,6 +453,10 @@ LogicalResult LoopOp::verifyRegions() {
     }
   }
 
+  if (!returnTypesMatch(region, op))
+    return emitOpError(
+        "result types do not match types yielded with `spirv.mlir.merge`");
+
   return success();
 }
 
@@ -609,6 +621,10 @@ LogicalResult SelectionOp::verifyRegions() {
   if (region.hasOneBlock())
     return emitOpError("must have a selection header block");
 
+  if (!returnTypesMatch(region, op))
+    return emitOpError(
+        "result types do not match types yielded with `spirv.mlir.merge`");
+
   return success();
 }
 

diff  --git a/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir b/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir
index b70bb40dae97f..81dce9822db48 100644
--- a/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir
@@ -496,6 +496,62 @@ func.func @loop_yield(%count : i32) -> () {
 
 // -----
 
+func.func @loop_yield_result_type_mismatch(%count : i32) -> () {
+  %zero = spirv.Constant 0: i32
+  %one = spirv.Constant 1: i32
+
+  // expected-error at +1{{result types do not match types yielded with `spirv.mlir.merge`}}
+  %final_i = spirv.mlir.loop -> f32 {
+    spirv.Branch ^header(%zero: i32)
+
+  ^header(%i : i32):
+    %cmp = spirv.SLessThan %i, %count : i32
+    spirv.BranchConditional %cmp, ^body, ^merge
+
+  ^body:
+    spirv.Branch ^continue
+
+  ^continue:
+    %new_i = spirv.IAdd %i, %one : i32
+    spirv.Branch ^header(%new_i: i32)
+
+  ^merge:
+    spirv.mlir.merge %i : i32
+  }
+
+  return
+}
+
+// -----
+
+func.func @loop_yield_result_count_mismatch(%count : i32) -> () {
+  %zero = spirv.Constant 0: i32
+  %one = spirv.Constant 1: i32
+
+  // expected-error at +1{{result types do not match types yielded with `spirv.mlir.merge`}}
+  %final_i = spirv.mlir.loop -> i32 {
+    spirv.Branch ^header(%zero: i32)
+
+  ^header(%i : i32):
+    %cmp = spirv.SLessThan %i, %count : i32
+    spirv.BranchConditional %cmp, ^body, ^merge
+
+  ^body:
+    spirv.Branch ^continue
+
+  ^continue:
+    %new_i = spirv.IAdd %i, %one : i32
+    spirv.Branch ^header(%new_i: i32)
+
+  ^merge:
+    spirv.mlir.merge %i, %i : i32, i32
+  }
+
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spirv.mlir.merge
 //===----------------------------------------------------------------------===//
@@ -922,6 +978,57 @@ func.func @selection_yield(%cond: i1) -> () {
 
 // -----
 
+func.func @selection_yield_result_type_mismatch(%cond: i1) -> () {
+  %zero = spirv.Constant 0: i32
+
+  // expected-error at +1{{result types do not match types yielded with `spirv.mlir.merge`}}
+  %yield:2 = spirv.mlir.selection -> f32, f32 {
+    spirv.BranchConditional %cond, ^then, ^else
+
+  ^then:
+    %one = spirv.Constant 1: i32
+    %three = spirv.Constant 3: i32
+    spirv.Branch ^merge(%one, %three : i32, i32)
+
+  ^else:
+    %two = spirv.Constant 2: i32
+    %four = spirv.Constant 4 : i32
+    spirv.Branch ^merge(%two, %four : i32, i32)
+
+  ^merge(%merged_1_2: i32, %merged_3_4: i32):
+    spirv.mlir.merge %merged_1_2, %merged_3_4 : i32, i32
+  }
+
+  spirv.Return
+}
+
+// -----
+
+func.func @selection_yield_result_count_mismatch(%cond: i1) -> () {
+  %zero = spirv.Constant 0: i32
+
+  // expected-error at +1{{result types do not match types yielded with `spirv.mlir.merge`}}
+  %yield:2 = spirv.mlir.selection -> f32, f32 {
+    spirv.BranchConditional %cond, ^then, ^else
+
+  ^then:
+    %one = spirv.Constant 1: i32
+    spirv.Branch ^merge(%one :i32)
+
+  ^else:
+    %two = spirv.Constant 2: i32
+    spirv.Branch ^merge(%two : i32)
+
+  ^merge(%merged_1_2: i32):
+    spirv.mlir.merge %merged_1_2 : i32
+  }
+
+  spirv.Return
+}
+
+// -----
+
+
 //===----------------------------------------------------------------------===//
 // spirv.Unreachable
 //===----------------------------------------------------------------------===//