[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.

3 files changed, 34 insertions, 29 deletions

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_LoopOp : SPIRV_Op<"mlir.loop", [InFunctionScope]> {
 
 // -----
 
-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 @@ def SPIRV_MergeOp : SPIRV_Op<"mlir.merge", [Pure, Terminator]> {
   let hasOpcode = 0;
 
   let autogenSerialization = 0;
+
+  let hasVerifier = 0;
 }
 
 // -----
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 @@ static bool isMergeBlock(Block &block) {
          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 @@ LogicalResult LoopOp::verifyRegions() {
   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 @@ void LoopOp::addEntryAndMergeBlock(OpBuilder &builder) {
 }
 
 //===----------------------------------------------------------------------===//
-// 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 @@ LogicalResult SelectionOp::verifyRegions() {
   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");
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 @only_entry_and_continue_branch_to_header() -> () {
 //===----------------------------------------------------------------------===//
 
 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(%cond : i1) -> () {
 
 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