Use SingleBlockImplicitTerminator trait for spv.module

This trait provides the ensureTerminator() utility function and
the checks to make sure a spv.module is indeed terminated with
spv._module_end.

PiperOrigin-RevId: 261664153

diff --git a/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td b/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
index 054da98..1d52860 100644 (file)
--- a/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/SPIRVStructureOps.td
@@ -30,7 +30,8 @@
 include "mlir/SPIRV/SPIRVBase.td"
 #endif // SPIRV_BASE
 
-def SPV_ModuleOp : SPV_Op<"module", []> {
+def SPV_ModuleOp : SPV_Op<"module",
+                          [SingleBlockImplicitTerminator<"ModuleEndOp">]> {
   let summary = "The top-level op that defines a SPIR-V module";
 
   let description = [{

diff --git a/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
index 2e59809..9366acc 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
@@ -683,12 +683,8 @@ static LogicalResult verify(spirv::LoadOp loadOp) {
 // spv.module
 //===----------------------------------------------------------------------===//
 
-static void ensureModuleEnd(Region *region, Builder builder, Location loc) {
-  impl::ensureRegionTerminator<spirv::ModuleEndOp>(*region, builder, loc);
-}
-
 void spirv::ModuleOp::build(Builder *builder, OperationState *state) {
-  ensureModuleEnd(state->addRegion(), *builder, state->location);
+  ensureTerminator(*state->addRegion(), *builder, state->location);
 }
 
 void spirv::ModuleOp::build(Builder *builder, OperationState *state,
@@ -704,7 +700,7 @@ void spirv::ModuleOp::build(Builder *builder, OperationState *state,
     state->addAttribute("extensions", extensions);
   if (extended_instruction_sets)
     state->addAttribute("extended_instruction_sets", extended_instruction_sets);
-  ensureModuleEnd(state->addRegion(), *builder, state->location);
+  ensureTerminator(*state->addRegion(), *builder, state->location);
 }
 
 static ParseResult parseModuleOp(OpAsmParser *parser, OperationState *state) {
@@ -726,8 +722,8 @@ static ParseResult parseModuleOp(OpAsmParser *parser, OperationState *state) {
       return failure();
   }
 
-  ensureModuleEnd(body, parser->getBuilder(), state->location);
-
+  spirv::ModuleOp::ensureTerminator(*body, parser->getBuilder(),
+                                    state->location);
   return success();
 }
 

diff --git a/mlir/test/Dialect/SPIRV/structure-ops.mlir b/mlir/test/Dialect/SPIRV/structure-ops.mlir
index e032ca6..db51b17 100644 (file)
--- a/mlir/test/Dialect/SPIRV/structure-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/structure-ops.mlir
@@ -107,8 +107,8 @@ spv.module "Logical" "Bla" { }
 
 // -----
 
-// Module_with_multiple_blocks
-// expected-error @+1 {{failed to verify constraint: region with 1 blocks}}
+// Module with multiple blocks
+// expected-error @+1 {{expects region #0 to have 0 or 1 blocks}}
 spv.module "Logical" "VulkanKHR" {
 ^first:
   spv.Return
@@ -118,6 +118,15 @@ spv.module "Logical" "VulkanKHR" {
 
 // -----
 
+// Module with wrong terminator
+// expected-error@+2 {{expects regions to end with 'spv._module_end'}}
+// expected-note@+1 {{in custom textual format, the absence of terminator implies 'spv._module_end'}}
+"spv.module"() ({
+  %0 = spv.constant true
+}) {addressing_model = 0 : i32, memory_model = 1 : i32} : () -> ()
+
+// -----
+
 // Use non SPIR-V op inside.module
 spv.module "Logical" "VulkanKHR" {
   // expected-error @+1 {{'spv.module' can only contain func and spv.* ops}}