commit 1878259c3960daeee91e2c95d487514d9e2984a4
author Longsheng Mou <longshengmou@gmail.com> Sat Mar 29 15:35:00 2025 +0800
committer GitHub <noreply@github.com> Sat Mar 29 15:35:00 2025 +0800
tree 08e56d02901f46beae13c15343827aee78503bce
parent 5b3e152e16950e551530b640335bb17dcc70a2fa

[mlir][spirv] Update verifier for `spirv.mlir.merge` (#133427)

- Moves the verification logic to the `verifyRegions` method of the
parent operation.
- Fixes a crash during verification when the last block lacks a
terminator.

Fixes #132850.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td
index ade20f9..cb7d27e 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td
@@ -351,7 +351,8 @@
 
 // -----
 
-def SPIRV_MergeOp : SPIRV_Op<"mlir.merge", [Pure, Terminator]> {
+def SPIRV_MergeOp : SPIRV_Op<"mlir.merge", [
+    Pure, Terminator, ParentOneOf<["SelectionOp", "LoopOp"]>]> {
   let summary = "A special terminator for merging a structured selection/loop.";
 
   let description = [{
@@ -371,6 +372,8 @@
   let hasOpcode = 0;
 
   let autogenSerialization = 0;
+
+  let hasVerifier = 0;
 }
 
 // -----

mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp

diff --git a/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
index 191bb33..bcfd7eb 100644
--- a/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/ControlFlowOps.cpp
@@ -259,6 +259,13 @@
          isa<spirv::MergeOp>(block.front());
 }
 
+/// Returns true if a `spirv.mlir.merge` op outside the merge block.
+static bool hasOtherMerge(Region &region) {
+  return !region.empty() && llvm::any_of(region.getOps(), [&](Operation &op) {
+    return isa<spirv::MergeOp>(op) && op.getBlock() != &region.back();
+  });
+}
+
 LogicalResult LoopOp::verifyRegions() {
   auto *op = getOperation();
 
@@ -298,6 +305,9 @@
   if (!isMergeBlock(merge))
     return emitOpError("last block must be the merge block with only one "
                        "'spirv.mlir.merge' op");
+  if (hasOtherMerge(region))
+    return emitOpError(
+        "should not have 'spirv.mlir.merge' op outside the merge block");
 
   if (std::next(region.begin()) == region.end())
     return emitOpError(
@@ -378,24 +388,6 @@
 }
 
 //===----------------------------------------------------------------------===//
-// spirv.mlir.merge
-//===----------------------------------------------------------------------===//
-
-LogicalResult MergeOp::verify() {
-  auto *parentOp = (*this)->getParentOp();
-  if (!parentOp || !isa<spirv::SelectionOp, spirv::LoopOp>(parentOp))
-    return emitOpError(
-        "expected parent op to be 'spirv.mlir.selection' or 'spirv.mlir.loop'");
-
-  // TODO: This check should be done in `verifyRegions` of parent op.
-  Block &parentLastBlock = (*this)->getParentRegion()->back();
-  if (getOperation() != parentLastBlock.getTerminator())
-    return emitOpError("can only be used in the last block of "
-                       "'spirv.mlir.selection' or 'spirv.mlir.loop'");
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // spirv.Return
 //===----------------------------------------------------------------------===//
 
@@ -507,6 +499,9 @@
   if (!isMergeBlock(region.back()))
     return emitOpError("last block must be the merge block with only one "
                        "'spirv.mlir.merge' op");
+  if (hasOtherMerge(region))
+    return emitOpError(
+        "should not have 'spirv.mlir.merge' op outside the merge block");
 
   if (std::next(region.begin()) == region.end())
     return emitOpError("must have a selection header block");

mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir

diff --git a/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir b/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir
index 1d1e284..188a55d 100644
--- a/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/control-flow-ops.mlir
@@ -431,29 +431,36 @@
 //===----------------------------------------------------------------------===//
 
 func.func @merge() -> () {
-  // expected-error @+1 {{expected parent op to be 'spirv.mlir.selection' or 'spirv.mlir.loop'}}
+  // expected-error @+1 {{expects parent op to be one of 'spirv.mlir.selection, spirv.mlir.loop'}}
   spirv.mlir.merge
 }
 
 // -----
 
 func.func @only_allowed_in_last_block(%cond : i1) -> () {
-  %zero = spirv.Constant 0: i32
-  %one = spirv.Constant 1: i32
-  %var = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
-
+  // expected-error @+1 {{'spirv.mlir.selection' op should not have 'spirv.mlir.merge' op outside the merge block}}
   spirv.mlir.selection {
     spirv.BranchConditional %cond, ^then, ^merge
-
   ^then:
-    spirv.Store "Function" %var, %one : i32
-    // expected-error @+1 {{can only be used in the last block of 'spirv.mlir.selection' or 'spirv.mlir.loop'}}
     spirv.mlir.merge
-
   ^merge:
     spirv.mlir.merge
   }
+  spirv.Return
+}
 
+// -----
+
+// Ensure this case not crash
+
+func.func @last_block_no_terminator(%cond : i1) -> () {
+  // expected-error @+1 {{empty block: expect at least a terminator}}
+  spirv.mlir.selection {
+    spirv.BranchConditional %cond, ^then, ^merge
+  ^then:
+    spirv.mlir.merge
+  ^merge:
+  }
   spirv.Return
 }
 
@@ -461,12 +468,12 @@
 
 func.func @only_allowed_in_last_block() -> () {
   %true = spirv.Constant true
+  // expected-error @+1 {{'spirv.mlir.loop' op should not have 'spirv.mlir.merge' op outside the merge block}}
   spirv.mlir.loop {
     spirv.Branch ^header
   ^header:
     spirv.BranchConditional %true, ^body, ^merge
   ^body:
-    // expected-error @+1 {{can only be used in the last block of 'spirv.mlir.selection' or 'spirv.mlir.loop'}}
     spirv.mlir.merge
   ^continue:
     spirv.Branch ^header