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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Igor Wodiany (IgWod-IMG)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp (+16) 
- (modified) mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir (+107) 


``````````diff
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
 //===----------------------------------------------------------------------===//

``````````

</details>


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