[mlir][SPIRV] Adding an attribute to capture configuration for cooperative matrix operations.

[mravishankar]

Summary

Each hardware that supports SPV_C_CooperativeMatrixNV has a list of
configurations that are supported natively. Add an attribute to
specify the configurations supported to the spv.target_env.

[ThomasRaoux]

LGTM

[antiagainst]

Awesome, thanks!

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

39

super nit: s/Cooperative/cooperative/

41

We can just use the man page link here: https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkCooperativeMatrixPropertiesNV.html.

It's more concentrated and simpler to load in browser. :)

[mravishankar]

Make the StructAttrField names snake_case.

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">;
 
+// Description of cooperative matrix operations supported on the
+// target. Represents `VkCooperativeMatrixPropertiesNV`. See
+// https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/VkCooperativeMatrixPropertiesNV.html
+def SPV_CooperativeMatrixPropertiesNVAttr :
+    StructAttr<"CooperativeMatrixPropertiesNVAttr", SPIRV_Dialect, [
+    StructFieldAttr<"m_size", I32Attr>,
+    StructFieldAttr<"n_size", I32Attr>,
+    StructFieldAttr<"k_size", I32Attr>,
+    StructFieldAttr<"a_type", TypeAttr>,
+    StructFieldAttr<"b_type", TypeAttr>,
+    StructFieldAttr<"c_type", TypeAttr>,
+    StructFieldAttr<"result_type", TypeAttr>,
+    StructFieldAttr<"scope", SPV_ScopeAttr>
+]>;
+
+def SPV_CooperativeMatrixPropertiesNVArrayAttr :
+    TypedArrayAttrBase<SPV_CooperativeMatrixPropertiesNVAttr,
+                       "CooperativeMatrixPropertiesNV array attribute">;
+
 // This attribute specifies the limits for various resources on the target
 // architecture.
 //
 // See https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#limits
 // for the complete list of limits and their explanation for the Vulkan API.
 // The following ones are those affecting SPIR-V CodeGen. Their default value
 // are the from Vulkan limit requirements:
 // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/html/vkspec.html#limits-minmax
 def SPV_ResourceLimitsAttr : StructAttr<"ResourceLimitsAttr", SPIRV_Dialect, [
     // The maximum total storage size, in bytes, available for variables
     // declared with the Workgroup storage class.
     StructFieldAttr<"max_compute_shared_memory_size",
                     DefaultValuedAttr<I32Attr, "16384">>,
 
     // The maximum total number of compute shader invocations in a single local
     // workgroup.
     StructFieldAttr<"max_compute_workgroup_invocations",
                     DefaultValuedAttr<I32Attr, "128">>,
     // The maximum size of a local compute workgroup, per dimension.
     StructFieldAttr<"max_compute_workgroup_size",
                     DefaultValuedAttr<I32ElementsAttr, "{128, 128, 64}">>,
 
     // The default number of invocations in each subgroup.
     // 0x7FFFFFFF means unknown.
-    StructFieldAttr<"subgroup_size", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>
+    StructFieldAttr<"subgroup_size", DefaultValuedAttr<I32Attr, "0x7FFFFFFF">>,
+
+    // The configurations of cooperative matrix operations
+    // supported. Default is an empty list.
+    StructFieldAttr<
+        "cooperative_matrix_properties_nv",
+        DefaultValuedAttr<SPV_CooperativeMatrixPropertiesNVArrayAttr, "{}">>
 ]>;
 
 #endif // SPIRV_TARGET_AND_ABI

mlir/lib/Dialect/SPIRV/TargetAndABI.cpp

 spirv::ResourceLimitsAttr
 spirv::getDefaultResourceLimits(MLIRContext *context) {
   // All the fields have default values. Here we just provide a nicer way to
   // construct a default resource limit attribute.
   return spirv::ResourceLimitsAttr ::get(
       /*max_compute_shared_memory_size=*/nullptr,
       /*max_compute_workgroup_invocations=*/nullptr,
       /*max_compute_workgroup_size=*/nullptr,
-      /*subgroup_size=*/nullptr, context);
+      /*subgroup_size=*/nullptr,
+      /*cooperative_matrix_properties_nv=*/nullptr, context);
 }
 
 StringRef spirv::getTargetEnvAttrName() { return "spv.target_env"; }
 
 spirv::TargetEnvAttr spirv::getDefaultTargetEnv(MLIRContext *context) {
   auto triple = spirv::VerCapExtAttr::get(spirv::Version::V_1_0,
                                           {spirv::Capability::Shader},
                                           ArrayRef<Extension>(), context);

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

#spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
       max_compute_workgroup_size = dense<[128, 64, 64]> : vector<3xi32>
     },
     more_stuff
   >
 } { return }
 
 // -----
 
+func @target_env_cooperative_matrix() attributes{
+  // CHECK:      spv.target_env = #spv.target_env<
+  // CHECK-SAME:   SPV_NV_cooperative_matrix
+  // CHECK-SAME:   cooperative_matrix_properties_nv = [
+  // CHECK-SAME:     {a_type = i8, b_type = i8, c_type = i32,
+  // CHECK-SAME:      k_size = 32 : i32, m_size = 8 : i32, n_size = 8 : i32
+  // CHECK-SAME:      result_type = i32, scope = 3 : i32}
+  // CHECK-SAME:     {a_type = f16, b_type = f16, c_type = f16,
+  // CHECK-SAME:      k_size = 16 : i32, m_size = 8 : i32, n_size = 8 : i32
+  // CHECK-SAME:      result_type = f16, scope = 3 : i32}
+  spv.target_env = #spv.target_env<
+  #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class,
+                            SPV_NV_cooperative_matrix]>,
+  {
+    cooperative_matrix_properties_nv = [{
+      m_size = 8: i32,
+      n_size = 8: i32,
+      k_size = 32: i32,
+      a_type = i8,
+      b_type = i8,
+      c_type = i32,
+      result_type = i32,
+      scope = 3: i32
+    }, {
+      m_size = 8: i32,
+      n_size = 8: i32,
+      k_size = 16: i32,
+      a_type = f16,
+      b_type = f16,
+      c_type = f16,
+      result_type = f16,
+      scope = 3: i32
+    }]
+  }>
+} { return }
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spv.vce
 //===----------------------------------------------------------------------===//
 
 func @vce_wrong_type() attributes {
   // expected-error @+1 {{expected valid keyword}}
   vce = #spv.vce<64>
 } { return }