[mlir][spirv] Add resource limits into target environment

[antiagainst]

Summary

This commit adds two resource limits, max_compute_workgroup_size
and max_compute_workgroup_invocations as resource limits to
the target environment. They are not used at the current moment,
but they will affect the SPIR-V CodeGen. Adding for now to have
a proper target environment modelling.

[merge_guards_bot]

Unit tests: pass. 62417 tests passed, 0 failed and 845 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h

 
 /// Returns the attribute name for specifying entry point information.
 StringRef getEntryPointABIAttrName();
 
 /// Gets the EntryPointABIAttr given its fields.
 EntryPointABIAttr getEntryPointABIAttr(ArrayRef<int32_t> localSize,
                                        MLIRContext *context);
 
+/// Returns a default resource limits attribute that uses numbers from
+/// "Table 46. Required Limits" of the Vulkan spec.
+ResourceLimitsAttr getDefaultResourceLimits(MLIRContext *context);
+
 /// Returns the attribute name for specifying SPIR-V target environment.
 StringRef getTargetEnvAttrName();
 
 /// Returns the default target environment: SPIR-V 1.0 with Shader capability
 /// and no extra extensions.
 TargetEnvAttr getDefaultTargetEnv(MLIRContext *context);
 
 /// Queries the target environment from the given `op` or returns the default

mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td

 ]>;
 
 def SPV_ExtensionArrayAttr : TypedArrayAttrBase<
     SPV_ExtensionAttr, "SPIR-V extension array attribute">;
 
 def SPV_CapabilityArrayAttr : TypedArrayAttrBase<
     SPV_CapabilityAttr, "SPIR-V capability array attribute">;
 
+// This attribute specifies the limits for various resources on the target
+// architecture.
+//
+// See https://renderdoc.org/vkspec_chunked/chap36.html#limits for the complete
+// list of limits and their explanation for the Vulkan API. The following ones
+// are those affecting SPIR-V CodeGen.
+def SPV_ResourceLimitsAttr : StructAttr<"ResourceLimitsAttr", SPV_Dialect, [
+    StructFieldAttr<"max_compute_workgroup_invocations", I32Attr>,
+    StructFieldAttr<"max_compute_workgroup_size", I32ElementsAttr>
+]>;
+
 // For the generated SPIR-V module, this attribute specifies the target version,
-// allowed extensions and capabilities.
+// allowed extensions and capabilities, and resource limits.
 def SPV_TargetEnvAttr : StructAttr<"TargetEnvAttr", SPV_Dialect, [
     StructFieldAttr<"version", SPV_VersionAttr>,
     StructFieldAttr<"extensions", SPV_ExtensionArrayAttr>,
-    StructFieldAttr<"capabilities", SPV_CapabilityArrayAttr>
+    StructFieldAttr<"capabilities", SPV_CapabilityArrayAttr>,
+    StructFieldAttr<"limits", SPV_ResourceLimitsAttr>
 ]>;
 
 #endif // SPIRV_TARGET_AND_ABI

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

if (!attr.isa<spirv::EntryPointABIAttr>())
              << "' attribute must be a dictionary attribute containing one "
                 "32-bit integer elements attribute: 'local_size'";
   } else if (symbol == spirv::getTargetEnvAttrName()) {
     if (!attr.isa<spirv::TargetEnvAttr>())
       return op->emitError("'")
              << symbol
              << "' must be a dictionary attribute containing one 32-bit "
                 "integer attribute 'version', one string array attribute "
-                "'extensions', and one 32-bit integer array attribute "
-                "'capabilities'";
+                "'extensions', one 32-bit integer array attribute "
+                "'capabilities', and one dictionary attribute 'limits'";
   } else {
     return op->emitError("found unsupported '")
            << symbol << "' attribute on operation";
   }
 
   return success();
 }
 

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

spirv::getEntryPointABIAttr(ArrayRef<int32_t> localSize, MLIRContext *context) {
   assert(localSize.size() == 3);
   return spirv::EntryPointABIAttr::get(
       DenseElementsAttr::get<int32_t>(
           VectorType::get(3, IntegerType::get(32, context)), localSize)
           .cast<DenseIntElementsAttr>(),
       context);
 }
 
+spirv::ResourceLimitsAttr
+spirv::getDefaultResourceLimits(MLIRContext *context) {
+  auto i32Type = IntegerType::get(32, context);
+  auto v3i32Type = VectorType::get(3, i32Type);
+
+  // These numbers are from "Table 46. Required Limits" of the Vulkan spec.
+  return spirv::ResourceLimitsAttr ::get(
+      IntegerAttr::get(i32Type, 128),
+      DenseIntElementsAttr::get<int32_t>(v3i32Type, {128, 128, 64}), context);
+}
+
 StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; }
 
 spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
   Builder builder(context);
   return spirv::TargetEnvAttr::get(
       builder.getI32IntegerAttr(static_cast<uint32_t>(spirv::Version::V_1_0)),
       builder.getI32ArrayAttr({}),
       builder.getI32ArrayAttr(
           {static_cast<uint32_t>(spirv::Capability::Shader)}),
-      context);
+      spirv::getDefaultResourceLimits(context), context);
 }
 
 spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) {
   if (auto attr = op->getAttrOfType<spirv::TargetEnvAttr>(
           spirv::getTargetEnvAttrName()))
     return attr;
   return getDefaultTargetEnv(op->getContext());
 }

mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir

indexing_maps = [
     affine_map<(i) -> (0)>
   ]
 }
 
 module attributes {
   spv.target_env = {
     version = 3 : i32,
     extensions = [],
-    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic
+    limits = {
+      max_compute_workgroup_invocations = 128 : i32,
+      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>
+    }
   }
 } {
 
 // CHECK:      spv.globalVariable
 // CHECK-SAME: built_in("LocalInvocationId")
 
 // CHECK:      func @single_workgroup_reduction
 // CHECK-SAME: (%[[INPUT:.+]]: !spv.ptr{{.+}}, %[[OUTPUT:.+]]: !spv.ptr{{.+}})
indexing_maps = [
     affine_map<(i) -> (0)>
   ]
 }
 
 module attributes {
   spv.target_env = {
     version = 3 : i32,
     extensions = [],
-    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic
+    limits = {
+      max_compute_workgroup_invocations = 128 : i32,
+      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>
+    }
   }
 } {
 func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
   linalg.generic #single_workgroup_reduction_trait %input, %output {
     ^bb(%in: i32, %out: i32):
       %sum = addi %in, %out : i32
       linalg.yield %sum : i32
indexing_maps = [
     affine_map<(i) -> (0)>
   ]
 }
 
 module attributes {
   spv.target_env = {
     version = 3 : i32,
     extensions = [],
-    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic
+    limits = {
+      max_compute_workgroup_invocations = 128 : i32,
+      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>
+    }
   }
 } {
 func @single_workgroup_reduction(%input: memref<16xi32>, %output: memref<1xi32>) attributes {
   spv.entry_point_abi = {local_size = dense<[32, 1, 1]>: vector<3xi32>}
 } {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
   linalg.generic #single_workgroup_reduction_trait %input, %output {
     ^bb(%in: i32, %out: i32):
indexing_maps = [
     affine_map<(i, j) -> (i)>
   ]
 }
 
 module attributes {
   spv.target_env = {
     version = 3 : i32,
     extensions = [],
-    capabilities = [1: i32, 63: i32] // Shader, GroupNonUniformArithmetic
+    capabilities = [1: i32, 63: i32], // Shader, GroupNonUniformArithmetic
+    limits = {
+      max_compute_workgroup_invocations = 128 : i32,
+      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>
+    }
   }
 } {
 func @single_workgroup_reduction(%input: memref<16x8xi32>, %output: memref<16xi32>) attributes {
   spv.entry_point_abi = {local_size = dense<[16, 8, 1]>: vector<3xi32>}
 } {
   // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
   linalg.generic #single_workgroup_reduction_trait %input, %output {
     ^bb(%in: i32, %out: i32):
       %sum = addi %in, %out : i32
       linalg.yield %sum : i32
   } : memref<16x8xi32>, memref<16xi32>
   spv.Return
 }
 }

mlir/test/Dialect/SPIRV/target-and-abi.mlir

 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv.target_env
 //===----------------------------------------------------------------------===//
 
-// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', and one 32-bit integer array attribute 'capabilities'}}
+// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', one 32-bit integer array attribute 'capabilities', and one dictionary attribute 'limits'}}
 func @target_env_wrong_type() attributes {
   spv.target_env = 64
 } { return }
 
 // -----
 
-// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', and one 32-bit integer array attribute 'capabilities'}}
+// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', one 32-bit integer array attribute 'capabilities', and one dictionary attribute 'limits'}}
 func @target_env_missing_fields() attributes {
   spv.target_env = {version = 0: i32}
 } { return }
 
 // -----
 
-// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', and one 32-bit integer array attribute 'capabilities'}}
+// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', one 32-bit integer array attribute 'capabilities', and one dictionary attribute 'limits'}}
 func @target_env_wrong_extension_type() attributes {
   spv.target_env = {version = 0: i32, extensions = [32: i32], capabilities = [1: i32]}
 } { return }
 
 // -----
 
-// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', and one 32-bit integer array attribute 'capabilities'}}
+// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', one 32-bit integer array attribute 'capabilities', and one dictionary attribute 'limits'}}
 func @target_env_wrong_extension() attributes {
   spv.target_env = {version = 0: i32, extensions = ["SPV_Something"], capabilities = [1: i32]}
 } { return }
 
 // -----
 
 func @target_env() attributes {
-  // CHECK: spv.target_env = {capabilities = [1 : i32], extensions = ["SPV_KHR_storage_buffer_storage_class"], version = 0 : i32}
-  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_storage_buffer_storage_class"], capabilities = [1: i32]}
+  // CHECK: spv.target_env = {capabilities = [1 : i32], extensions = ["SPV_KHR_storage_buffer_storage_class"], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 64, 64]> : vector<3xi32>}, version = 0 : i32}
+  spv.target_env = {
+    version = 0: i32,
+    extensions = ["SPV_KHR_storage_buffer_storage_class"],
+    capabilities = [1: i32],
+    limits = {
+      max_compute_workgroup_invocations = 128 : i32,
+      max_compute_workgroup_size = dense<[128, 64, 64]> : vector<3xi32>
+    }
+  }
 } { return }
 
 // -----
 
-// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', and one 32-bit integer array attribute 'capabilities'}}
+// expected-error @+1 {{'spv.target_env' must be a dictionary attribute containing one 32-bit integer attribute 'version', one string array attribute 'extensions', one 32-bit integer array attribute 'capabilities', and one dictionary attribute 'limits'}}
 func @target_env_extra_fields() attributes {
   spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_storage_buffer_storage_class"], capabilities = [1: i32], extra = 32}
 } { return }

mlir/test/Dialect/SPIRV/target-env.mlir

 //             PhysicalStorageBufferAddresses (5347)
 
 //===----------------------------------------------------------------------===//
 // MaxVersion
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @cmp_exchange_weak_suitable_version_capabilities
 func @cmp_exchange_weak_suitable_version_capabilities(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 attributes {
-  spv.target_env = {version = 1: i32, extensions = [], capabilities = [6: i32, 21: i32]}
+  spv.target_env = {version = 1: i32, extensions = [], capabilities = [6: i32, 21: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.AtomicCompareExchangeWeak "Workgroup" "AcquireRelease|AtomicCounterMemory" "Acquire"
   %0 = "test.convert_to_atomic_compare_exchange_weak_op"(%ptr, %value, %comparator): (!spv.ptr<i32, Workgroup>, i32, i32) -> (i32)
   return %0: i32
 }
 
 // CHECK-LABEL: @cmp_exchange_weak_unsupported_version
 func @cmp_exchange_weak_unsupported_version(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 attributes {
-  spv.target_env = {version = 4: i32, extensions = [], capabilities = [6: i32, 21: i32]}
+  spv.target_env = {version = 4: i32, extensions = [], capabilities = [6: i32, 21: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_atomic_compare_exchange_weak_op
   %0 = "test.convert_to_atomic_compare_exchange_weak_op"(%ptr, %value, %comparator): (!spv.ptr<i32, Workgroup>, i32, i32) -> (i32)
   return %0: i32
 }
 
 //===----------------------------------------------------------------------===//
 // MinVersion
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @group_non_uniform_ballot_suitable_version
 func @group_non_uniform_ballot_suitable_version(%predicate: i1) -> vector<4xi32> attributes {
-  spv.target_env = {version = 4: i32, extensions = [], capabilities = [64: i32]}
+  spv.target_env = {version = 4: i32, extensions = [], capabilities = [64: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.GroupNonUniformBallot "Workgroup"
   %0 = "test.convert_to_group_non_uniform_ballot_op"(%predicate): (i1) -> (vector<4xi32>)
   return %0: vector<4xi32>
 }
 
 // CHECK-LABEL: @group_non_uniform_ballot_unsupported_version
 func @group_non_uniform_ballot_unsupported_version(%predicate: i1) -> vector<4xi32> attributes {
-  spv.target_env = {version = 1: i32, extensions = [], capabilities = [64: i32]}
+  spv.target_env = {version = 1: i32, extensions = [], capabilities = [64: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_group_non_uniform_ballot_op
   %0 = "test.convert_to_group_non_uniform_ballot_op"(%predicate): (i1) -> (vector<4xi32>)
   return %0: vector<4xi32>
 }
 
 //===----------------------------------------------------------------------===//
 // Capability
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @cmp_exchange_weak_missing_capability_kernel
 func @cmp_exchange_weak_missing_capability_kernel(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 attributes {
-  spv.target_env = {version = 3: i32, extensions = [], capabilities = [21: i32]}
+  spv.target_env = {version = 3: i32, extensions = [], capabilities = [21: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_atomic_compare_exchange_weak_op
   %0 = "test.convert_to_atomic_compare_exchange_weak_op"(%ptr, %value, %comparator): (!spv.ptr<i32, Workgroup>, i32, i32) -> (i32)
   return %0: i32
 }
 
 // CHECK-LABEL: @cmp_exchange_weak_missing_capability_atomic_storage
 func @cmp_exchange_weak_missing_capability_atomic_storage(%ptr: !spv.ptr<i32, Workgroup>, %value: i32, %comparator: i32) -> i32 attributes {
-  spv.target_env = {version = 3: i32, extensions = [], capabilities = [6: i32]}
+  spv.target_env = {version = 3: i32, extensions = [], capabilities = [6: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_atomic_compare_exchange_weak_op
   %0 = "test.convert_to_atomic_compare_exchange_weak_op"(%ptr, %value, %comparator): (!spv.ptr<i32, Workgroup>, i32, i32) -> (i32)
   return %0: i32
 }
 
 // CHECK-LABEL: @subgroup_ballot_missing_capability
 func @subgroup_ballot_missing_capability(%predicate: i1) -> vector<4xi32> attributes {
-  spv.target_env = {version = 4: i32, extensions = ["SPV_KHR_shader_ballot"], capabilities = []}
+  spv.target_env = {version = 4: i32, extensions = ["SPV_KHR_shader_ballot"], capabilities = [], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_subgroup_ballot_op
   %0 = "test.convert_to_subgroup_ballot_op"(%predicate): (i1) -> (vector<4xi32>)
   return %0: vector<4xi32>
 }
 
 // CHECK-LABEL: @bit_reverse_directly_implied_capability
 func @bit_reverse_directly_implied_capability(%operand: i32) -> i32 attributes {
-  spv.target_env = {version = 0: i32, extensions = [], capabilities = [2: i32]}
+  spv.target_env = {version = 0: i32, extensions = [], capabilities = [2: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.BitReverse
   %0 = "test.convert_to_bit_reverse_op"(%operand): (i32) -> (i32)
   return %0: i32
 }
 
 // CHECK-LABEL: @bit_reverse_recursively_implied_capability
 func @bit_reverse_recursively_implied_capability(%operand: i32) -> i32 attributes {
-  spv.target_env = {version = 0: i32, extensions = [], capabilities = [24: i32]}
+  spv.target_env = {version = 0: i32, extensions = [], capabilities = [24: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.BitReverse
   %0 = "test.convert_to_bit_reverse_op"(%operand): (i32) -> (i32)
   return %0: i32
 }
 
 //===----------------------------------------------------------------------===//
 // Extension
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @subgroup_ballot_suitable_extension
 func @subgroup_ballot_suitable_extension(%predicate: i1) -> vector<4xi32> attributes {
-  spv.target_env = {version = 4: i32, extensions = ["SPV_KHR_shader_ballot"], capabilities = [4423: i32]}
+  spv.target_env = {version = 4: i32, extensions = ["SPV_KHR_shader_ballot"], capabilities = [4423: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.SubgroupBallotKHR
   %0 = "test.convert_to_subgroup_ballot_op"(%predicate): (i1) -> (vector<4xi32>)
   return %0: vector<4xi32>
 }
 
 // CHECK-LABEL: @subgroup_ballot_missing_extension
 func @subgroup_ballot_missing_extension(%predicate: i1) -> vector<4xi32> attributes {
-  spv.target_env = {version = 4: i32, extensions = [], capabilities = [4423: i32]}
+  spv.target_env = {version = 4: i32, extensions = [], capabilities = [4423: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_subgroup_ballot_op
   %0 = "test.convert_to_subgroup_ballot_op"(%predicate): (i1) -> (vector<4xi32>)
   return %0: vector<4xi32>
 }
 
 // CHECK-LABEL: @module_suitable_extension1
 func @module_suitable_extension1() attributes {
-  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_vulkan_memory_model", "SPV_EXT_physical_storage_buffer"], capabilities = [5345: i32, 5347: i32]}
+  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_vulkan_memory_model", "SPV_EXT_physical_storage_buffer"], capabilities = [5345: i32, 5347: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.module "PhysicalStorageBuffer64" "Vulkan"
   "test.convert_to_module_op"() : () ->()
   return
 }
 
 // CHECK-LABEL: @module_suitable_extension2
 func @module_suitable_extension2() attributes {
-  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_vulkan_memory_model", "SPV_KHR_physical_storage_buffer"], capabilities = [5345: i32, 5347: i32]}
+  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_vulkan_memory_model", "SPV_KHR_physical_storage_buffer"], capabilities = [5345: i32, 5347: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.module "PhysicalStorageBuffer64" "Vulkan"
   "test.convert_to_module_op"() : () -> ()
   return
 }
 
 // CHECK-LABEL: @module_missing_extension_mm
 func @module_missing_extension_mm() attributes {
-  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_physical_storage_buffer"], capabilities = [5345: i32, 5347: i32]}
+  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_physical_storage_buffer"], capabilities = [5345: i32, 5347: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_module_op
   "test.convert_to_module_op"() : () -> ()
   return
 }
 
 // CHECK-LABEL: @module_missing_extension_am
 func @module_missing_extension_am() attributes {
-  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_vulkan_memory_model"], capabilities = [5345: i32, 5347: i32]}
+  spv.target_env = {version = 0: i32, extensions = ["SPV_KHR_vulkan_memory_model"], capabilities = [5345: i32, 5347: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: test.convert_to_module_op
   "test.convert_to_module_op"() : () -> ()
   return
 }
 
 // CHECK-LABEL: @module_implied_extension
 func @module_implied_extension() attributes {
   // Version 1.5 implies SPV_KHR_vulkan_memory_model and SPV_KHR_physical_storage_buffer.
-  spv.target_env = {version = 5: i32, extensions = [], capabilities = [5345: i32, 5347: i32]}
+  spv.target_env = {version = 5: i32, extensions = [], capabilities = [5345: i32, 5347: i32], limits = {max_compute_workgroup_invocations = 128 : i32, max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}}
 } {
   // CHECK: spv.module "PhysicalStorageBuffer64" "Vulkan"
   "test.convert_to_module_op"() : () -> ()
   return
 }