diff --git a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h --- a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h +++ b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.h @@ -43,6 +43,10 @@ EntryPointABIAttr getEntryPointABIAttr(ArrayRef 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(); diff --git a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td --- a/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td +++ b/mlir/include/mlir/Dialect/SPIRV/TargetAndABI.td @@ -48,12 +48,24 @@ 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 diff --git a/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp b/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp --- a/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp +++ b/mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp @@ -664,8 +664,8 @@ << 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"; diff --git a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp --- a/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp +++ b/mlir/lib/Dialect/SPIRV/TargetAndABI.cpp @@ -45,6 +45,17 @@ 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(v3i32Type, {128, 128, 64}), context); +} + StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; } spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) { @@ -54,7 +65,7 @@ builder.getI32ArrayAttr({}), builder.getI32ArrayAttr( {static_cast(spirv::Capability::Shader)}), - context); + spirv::getDefaultResourceLimits(context), context); } spirv::TargetEnvAttr spirv::lookupTargetEnvOrDefault(Operation *op) { diff --git a/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir b/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir --- a/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir +++ b/mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir @@ -18,7 +18,11 @@ 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> + } } } { @@ -79,7 +83,11 @@ 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>) { @@ -111,7 +119,11 @@ 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 { @@ -145,7 +157,11 @@ 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 { diff --git a/mlir/test/Dialect/SPIRV/target-and-abi.mlir b/mlir/test/Dialect/SPIRV/target-and-abi.mlir --- a/mlir/test/Dialect/SPIRV/target-and-abi.mlir +++ b/mlir/test/Dialect/SPIRV/target-and-abi.mlir @@ -106,28 +106,28 @@ // 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 } @@ -135,13 +135,21 @@ // ----- 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 } diff --git a/mlir/test/Dialect/SPIRV/target-env.mlir b/mlir/test/Dialect/SPIRV/target-env.mlir --- a/mlir/test/Dialect/SPIRV/target-env.mlir +++ b/mlir/test/Dialect/SPIRV/target-env.mlir @@ -42,7 +42,7 @@ // CHECK-LABEL: @cmp_exchange_weak_suitable_version_capabilities func @cmp_exchange_weak_suitable_version_capabilities(%ptr: !spv.ptr, %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, i32) -> (i32) @@ -51,7 +51,7 @@ // CHECK-LABEL: @cmp_exchange_weak_unsupported_version func @cmp_exchange_weak_unsupported_version(%ptr: !spv.ptr, %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, i32) -> (i32) @@ -64,7 +64,7 @@ // 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>) @@ -73,7 +73,7 @@ // 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>) @@ -86,7 +86,7 @@ // CHECK-LABEL: @cmp_exchange_weak_missing_capability_kernel func @cmp_exchange_weak_missing_capability_kernel(%ptr: !spv.ptr, %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, i32) -> (i32) @@ -95,7 +95,7 @@ // CHECK-LABEL: @cmp_exchange_weak_missing_capability_atomic_storage func @cmp_exchange_weak_missing_capability_atomic_storage(%ptr: !spv.ptr, %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, i32) -> (i32) @@ -104,7 +104,7 @@ // 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>) @@ -113,7 +113,7 @@ // 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) @@ -122,7 +122,7 @@ // 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) @@ -135,7 +135,7 @@ // 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>) @@ -144,7 +144,7 @@ // 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>) @@ -153,7 +153,7 @@ // 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"() : () ->() @@ -162,7 +162,7 @@ // 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"() : () -> () @@ -171,7 +171,7 @@ // 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"() : () -> () @@ -180,7 +180,7 @@ // 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"() : () -> () @@ -190,7 +190,7 @@ // 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"() : () -> ()