[mlir][spirv] Change translation to use spirv.module

Update the SPIRV `mlir-translate` translations to translate to/from
`spirv.module` instead of `builtin.module`. This simplifies the
translation since the code no longer needs to walk the module looking
for a SPIRV module, however it requires passing `-no-implicit-module` to
all the tests.

Reviewed By: antiagainst

Differential Revision: https://reviews.llvm.org/D135819

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

@@ -51,17 +51,7 @@ deserializeModule(const llvm::MemoryBuffer *input, MLIRContext *context) {
 
   auto binary = llvm::makeArrayRef(reinterpret_cast<const uint32_t *>(start),
                                    size / sizeof(uint32_t));
-
-  OwningOpRef<spirv::ModuleOp> spirvModule =
-      spirv::deserialize(binary, context);
-  if (!spirvModule)
-    return {};
-
-  OwningOpRef<ModuleOp> module(ModuleOp::create(FileLineColLoc::get(
-      context, input->getBufferIdentifier(), /*line=*/0, /*column=*/0)));
-  module->getBody()->push_front(spirvModule.release());
-
-  return std::move(module);
+  return spirv::deserialize(binary, context);
 }
 
 namespace mlir {
@@ -80,22 +70,10 @@ void registerFromSPIRVTranslation() {
 // Serialization registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult serializeModule(ModuleOp module, raw_ostream &output) {
-  if (!module)
-    return failure();
-
+static LogicalResult serializeModule(spirv::ModuleOp module,
+                                     raw_ostream &output) {
   SmallVector<uint32_t, 0> binary;
-
-  SmallVector<spirv::ModuleOp, 1> spirvModules;
-  module.walk([&](spirv::ModuleOp op) { spirvModules.push_back(op); });
-
-  if (spirvModules.empty())
-    return module.emitError("found no 'spirv.module' op");
-
-  if (spirvModules.size() != 1)
-    return module.emitError("found more than one 'spirv.module' op");
-
-  if (failed(spirv::serialize(spirvModules[0], binary)))
+  if (failed(spirv::serialize(module, binary)))
     return failure();
 
   output.write(reinterpret_cast<char *>(binary.data()),
@@ -108,7 +86,7 @@ namespace mlir {
 void registerToSPIRVTranslation() {
   TranslateFromMLIRRegistration toBinary(
       "serialize-spirv", "serialize SPIR-V dialect",
-      [](ModuleOp module, raw_ostream &output) {
+      [](spirv::ModuleOp module, raw_ostream &output) {
         return serializeModule(module, output);
       },
       [](DialectRegistry &registry) {
@@ -121,21 +99,14 @@ void registerToSPIRVTranslation() {
 // Round-trip registration
 //===----------------------------------------------------------------------===//
 
-static LogicalResult roundTripModule(ModuleOp srcModule, bool emitDebugInfo,
+static LogicalResult roundTripModule(spirv::ModuleOp module, bool emitDebugInfo,
                                      raw_ostream &output) {
   SmallVector<uint32_t, 0> binary;
-  MLIRContext *context = srcModule.getContext();
-  auto spirvModules = srcModule.getOps<spirv::ModuleOp>();
-
-  if (spirvModules.begin() == spirvModules.end())
-    return srcModule.emitError("found no 'spirv.module' op");
-
-  if (std::next(spirvModules.begin()) != spirvModules.end())
-    return srcModule.emitError("found more than one 'spirv.module' op");
+  MLIRContext *context = module->getContext();
 
   spirv::SerializationOptions options;
   options.emitDebugInfo = emitDebugInfo;
-  if (failed(spirv::serialize(*spirvModules.begin(), binary, options)))
+  if (failed(spirv::serialize(module, binary, options)))
     return failure();
 
   MLIRContext deserializationContext(context->getDialectRegistry());
@@ -146,15 +117,7 @@ static LogicalResult roundTripModule(ModuleOp srcModule, bool emitDebugInfo,
       spirv::deserialize(binary, &deserializationContext);
   if (!spirvModule)
     return failure();
-
-  // Wrap around in a new MLIR module.
-  OwningOpRef<ModuleOp> dstModule(ModuleOp::create(
-      FileLineColLoc::get(&deserializationContext,
-                          /*filename=*/"", /*line=*/0, /*column=*/0)));
-  dstModule->getBody()->push_front(spirvModule.release());
-  if (failed(verify(*dstModule)))
-    return failure();
-  dstModule->print(output);
+  spirvModule->print(output);
 
   return mlir::success();
 }
@@ -163,7 +126,7 @@ namespace mlir {
 void registerTestRoundtripSPIRV() {
   TranslateFromMLIRRegistration roundtrip(
       "test-spirv-roundtrip", "test roundtrip in SPIR-V dialect",
-      [](ModuleOp module, raw_ostream &output) {
+      [](spirv::ModuleOp module, raw_ostream &output) {
         return roundTripModule(module, /*emitDebugInfo=*/false, output);
       },
       [](DialectRegistry &registry) {
@@ -174,7 +137,7 @@ void registerTestRoundtripSPIRV() {
 void registerTestRoundtripDebugSPIRV() {
   TranslateFromMLIRRegistration roundtrip(
       "test-spirv-roundtrip-debug", "test roundtrip debug in SPIR-V",
-      [](ModuleOp module, raw_ostream &output) {
+      [](spirv::ModuleOp module, raw_ostream &output) {
         return roundTripModule(module, /*emitDebugInfo=*/true, output);
       },
       [](DialectRegistry &registry) {

mlir/test/Target/SPIRV/arithmetic-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @fmul(%arg0 : f32, %arg1 : f32) "None" {

mlir/test/Target/SPIRV/array.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @array_stride(%arg0 : !spirv.ptr<!spirv.array<4x!spirv.array<4xf32, stride=4>, stride=128>, StorageBuffer>, %arg1 : i32, %arg2 : i32) "None" {

mlir/test/Target/SPIRV/atomic-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @test_int_atomics

mlir/test/Target/SPIRV/barrier-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @memory_barrier_0() -> () "None" {

mlir/test/Target/SPIRV/bit-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @bitcount(%arg: i32) -> i32 "None" {

mlir/test/Target/SPIRV/cast-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @bit_cast(%arg0 : f32) "None" {

mlir/test/Target/SPIRV/composite-op.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @composite_insert(%arg0 : !spirv.struct<(f32, !spirv.struct<(!spirv.array<4xf32>, f32)>)>, %arg1: !spirv.array<4xf32>) -> !spirv.struct<(f32, !spirv.struct<(!spirv.array<4xf32>, f32)>)> "None" {

mlir/test/Target/SPIRV/constant.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @bool_const

mlir/test/Target/SPIRV/cooperative-matrix-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [CooperativeMatrixNV], [SPV_NV_cooperative_matrix]> {
   // CHECK-LABEL: @cooperative_matrix_load

mlir/test/Target/SPIRV/debug.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip-debug -mlir-print-debuginfo -mlir-print-local-scope %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip-debug -mlir-print-debuginfo -mlir-print-local-scope %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: loc({{".*debug.mlir"}}:5:3)

mlir/test/Target/SPIRV/decorations.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: location = 0 : i32

mlir/test/Target/SPIRV/entry-point.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @noop() -> () "None" {

mlir/test/Target/SPIRV/execution-mode.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @foo() -> () "None" {

mlir/test/Target/SPIRV/function-call.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.GlobalVariable @var1 : !spirv.ptr<!spirv.array<4xf32>, Input>

mlir/test/Target/SPIRV/gl-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @math(%arg0 : f32, %arg1 : f32, %arg2 : i32) "None" {

mlir/test/Target/SPIRV/global-variable.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 // CHECK:      spirv.GlobalVariable @var0 bind(1, 0) : !spirv.ptr<f32, Input>
 // CHECK-NEXT: spirv.GlobalVariable @var1 bind(0, 1) : !spirv.ptr<f32, Output>

mlir/test/Target/SPIRV/group-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @subgroup_ballot

mlir/test/Target/SPIRV/image-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @image(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32) "None" {

mlir/test/Target/SPIRV/image.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: !spirv.ptr<!spirv.image<f32, Dim1D, NoDepth, NonArrayed, SingleSampled, SamplerUnknown, Unknown>, UniformConstant>

mlir/test/Target/SPIRV/invalid-module.mlir

@@ -1,4 +1,4 @@
 // RUN: mlir-translate %s -serialize-spirv -no-implicit-module -verify-diagnostics
 
-// expected-error@below {{expected a 'builtin.module' op, got 'spirv.module'}}
-spirv.module Logical Simple {}
+// expected-error@below {{expected a 'spirv.module' op, got 'builtin.module'}}
+module {}

mlir/test/Target/SPIRV/joint-matrix-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [JointMatrixINTEL], [SPV_INTEL_joint_matrix]> {
   // CHECK-LABEL: @joint_matrix_load

mlir/test/Target/SPIRV/logical-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @iequal_scalar(%arg0: i32, %arg1: i32)  "None" {

mlir/test/Target/SPIRV/loop.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 // Single loop
 

mlir/test/Target/SPIRV/matrix.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @matrix_access_chain

mlir/test/Target/SPIRV/memory-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {

mlir/test/Target/SPIRV/module.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 // CHECK:      spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
 // CHECK-NEXT:   spirv.func @foo() "Inline" {

mlir/test/Target/SPIRV/non-uniform-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @group_non_uniform_ballot

mlir/test/Target/SPIRV/ocl-ops.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Physical64 OpenCL requires #spirv.vce<v1.0, [Kernel, Addresses], []> {
   spirv.func @float_insts(%arg0 : f32) "None" {

mlir/test/Target/SPIRV/phi.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 // Test branch with one block argument
 

mlir/test/Target/SPIRV/sampled-image.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: !spirv.ptr<!spirv.sampled_image<!spirv.image<f32, Dim1D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, UniformConstant>

mlir/test/Target/SPIRV/selection.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 // Selection with both then and else branches
 

mlir/test/Target/SPIRV/spec-constant.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: spirv.SpecConstant @sc_true = true

mlir/test/Target/SPIRV/struct.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK: !spirv.ptr<!spirv.struct<(!spirv.array<128 x f32, stride=4> [0])>, Input>

mlir/test/Target/SPIRV/terminator.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @ret

mlir/test/Target/SPIRV/undef.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   spirv.func @foo() -> () "None" {