[mlirv][spirv] Add KHR Cooperative Matrix type and extension

[kuhar]

Summary

Start plumbing through support for the SPV_KHR_cooperative_matrix
extension: https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_cooperative_matrix.html.

Register the extension, add new coop matrix type, and add
spirv.KHR.CooperativeMatrixLength op to exercise it.

Make sure that mixing of the KHR and NV coop matrix extensions is not
allowed. Make cast verification more robust.

[kuhar]

Clean up

[kuhar]

Drop some unnecessary auto

[qedawkins]

Might want to wait for @antiagainst to take a look as well.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td

33

nit: attribute

mlir/test/Dialect/SPIRV/IR/types.mlir

453

The extension documentation indicates that only 0, 1, and 2 are valid uses. Should we verify that as a part of the type?

[kuhar]

mlir/test/Dialect/SPIRV/IR/types.mlir

453

Good point, although we should probably name these (MatrixA, MatrixB, MatrixAcc) instead of using plain integers, just like we do with the scope part.

[kuhar]

Add a new enum attribute to represent coop matrix use.
Fix a typo.

[antiagainst]

Nice work, thanks Jakub!

mlir/test/Dialect/SPIRV/IR/types.mlir

453

Yes, the symbolic ones certainly are better.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

 def SPV_KHR_subgroup_uniform_control_flow    : I32EnumAttrCase<"SPV_KHR_subgroup_uniform_control_flow", 23>;
 def SPV_KHR_linkonce_odr                     : I32EnumAttrCase<"SPV_KHR_linkonce_odr", 24>;
 def SPV_KHR_fragment_shader_barycentric      : I32EnumAttrCase<"SPV_KHR_fragment_shader_barycentric", 25>;
 def SPV_KHR_ray_cull_mask                    : I32EnumAttrCase<"SPV_KHR_ray_cull_mask", 26>;
 def SPV_KHR_uniform_group_instructions       : I32EnumAttrCase<"SPV_KHR_uniform_group_instructions", 27>;
 def SPV_KHR_subgroup_rotate                  : I32EnumAttrCase<"SPV_KHR_subgroup_rotate", 28>;
 def SPV_KHR_non_semantic_info                : I32EnumAttrCase<"SPV_KHR_non_semantic_info", 29>;
 def SPV_KHR_terminate_invocation             : I32EnumAttrCase<"SPV_KHR_terminate_invocation", 30>;
+def SPV_KHR_cooperative_matrix               : I32EnumAttrCase<"SPV_KHR_cooperative_matrix", 31>;
 
 def SPV_EXT_demote_to_helper_invocation  : I32EnumAttrCase<"SPV_EXT_demote_to_helper_invocation", 1000>;
 def SPV_EXT_descriptor_indexing          : I32EnumAttrCase<"SPV_EXT_descriptor_indexing", 1001>;
 def SPV_EXT_fragment_fully_covered       : I32EnumAttrCase<"SPV_EXT_fragment_fully_covered", 1002>;
 def SPV_EXT_fragment_invocation_density  : I32EnumAttrCase<"SPV_EXT_fragment_invocation_density", 1003>;
 def SPV_EXT_fragment_shader_interlock    : I32EnumAttrCase<"SPV_EXT_fragment_shader_interlock", 1004>;
 def SPV_EXT_physical_storage_buffer      : I32EnumAttrCase<"SPV_EXT_physical_storage_buffer", 1005>;
 def SPV_EXT_shader_stencil_export        : I32EnumAttrCase<"SPV_EXT_shader_stencil_export", 1006>;
SPIRV_I32EnumAttr<"Extension", "supported SPIR-V extensions", "ext", [
       SPV_KHR_storage_buffer_storage_class, SPV_KHR_subgroup_vote,
       SPV_KHR_variable_pointers, SPV_KHR_vulkan_memory_model, SPV_KHR_expect_assume,
       SPV_KHR_integer_dot_product, SPV_KHR_bit_instructions, SPV_KHR_fragment_shading_rate,
       SPV_KHR_workgroup_memory_explicit_layout, SPV_KHR_ray_query,
       SPV_KHR_ray_tracing, SPV_KHR_subgroup_uniform_control_flow, SPV_KHR_linkonce_odr,
       SPV_KHR_fragment_shader_barycentric, SPV_KHR_ray_cull_mask,
       SPV_KHR_uniform_group_instructions, SPV_KHR_subgroup_rotate,
       SPV_KHR_non_semantic_info, SPV_KHR_terminate_invocation,
+      SPV_KHR_cooperative_matrix,
       SPV_EXT_demote_to_helper_invocation, SPV_EXT_descriptor_indexing,
       SPV_EXT_fragment_fully_covered, SPV_EXT_fragment_invocation_density,
       SPV_EXT_fragment_shader_interlock, SPV_EXT_physical_storage_buffer,
       SPV_EXT_shader_stencil_export, SPV_EXT_shader_viewport_index_layer,
       SPV_EXT_shader_atomic_float_add, SPV_EXT_shader_atomic_float_min_max,
       SPV_EXT_shader_image_int64, SPV_EXT_shader_atomic_float16_add,
       SPV_AMD_gpu_shader_half_float_fetch, SPV_AMD_shader_ballot,
       SPV_AMD_shader_explicit_vertex_parameter, SPV_AMD_shader_fragment_mask,
list<Availability> availability = [
     MinVersion<SPIRV_V_1_6>
   ];
 }
 def SPIRV_C_RayCullMaskKHR                              : I32EnumAttrCase<"RayCullMaskKHR", 6020> {
   list<Availability> availability = [
     Extension<[SPV_KHR_ray_cull_mask]>
   ];
 }
+def SPIRV_C_CooperativeMatrixKHR                        : I32EnumAttrCase<"CooperativeMatrixKHR", 6022> {
+  list<Availability> availability = [
+    Extension<[SPV_KHR_cooperative_matrix]>,
+    MinVersion<SPIRV_V_1_6>
+  ];
+}
 def SPIRV_C_BitInstructions                             : I32EnumAttrCase<"BitInstructions", 6025> {
   list<Availability> availability = [
     Extension<[SPV_KHR_bit_instructions]>
   ];
 }
 def SPIRV_C_AtomicFloat32AddEXT                         : I32EnumAttrCase<"AtomicFloat32AddEXT", 6033> {
   list<Availability> availability = [
     Extension<[SPV_EXT_shader_atomic_float_add]>
SPIRV_I32EnumAttr<"Capability", "valid SPIR-V Capability", "capability", [
       SPIRV_C_UnstructuredLoopControlsINTEL, SPIRV_C_FPGALoopControlsINTEL,
       SPIRV_C_KernelAttributesINTEL, SPIRV_C_FPGAKernelAttributesINTEL,
       SPIRV_C_FPGAMemoryAccessesINTEL, SPIRV_C_FPGAClusterAttributesINTEL,
       SPIRV_C_LoopFuseINTEL, SPIRV_C_MemoryAccessAliasingINTEL,
       SPIRV_C_FPGABufferLocationINTEL, SPIRV_C_ArbitraryPrecisionFixedPointINTEL,
       SPIRV_C_USMStorageClassesINTEL, SPIRV_C_IOPipesINTEL, SPIRV_C_BlockingPipesINTEL,
       SPIRV_C_FPGARegINTEL, SPIRV_C_DotProductInputAll,
       SPIRV_C_DotProductInput4x8BitPacked, SPIRV_C_DotProduct, SPIRV_C_RayCullMaskKHR,
+      SPIRV_C_CooperativeMatrixKHR,
       SPIRV_C_BitInstructions, SPIRV_C_AtomicFloat32AddEXT, SPIRV_C_AtomicFloat64AddEXT,
       SPIRV_C_LongConstantCompositeINTEL, SPIRV_C_OptNoneINTEL,
       SPIRV_C_AtomicFloat16AddEXT, SPIRV_C_DebugInfoModuleINTEL, SPIRV_C_SplitBarrierINTEL,
       SPIRV_C_GroupUniformArithmeticKHR, SPIRV_C_Shader, SPIRV_C_Vector16,
       SPIRV_C_Float16Buffer, SPIRV_C_Int64Atomics, SPIRV_C_ImageBasic, SPIRV_C_Pipes,
       SPIRV_C_DeviceEnqueue, SPIRV_C_LiteralSampler, SPIRV_C_GenericPointer, SPIRV_C_Image1D,
       SPIRV_C_ImageBuffer, SPIRV_C_NamedBarrier, SPIRV_C_GroupNonUniformVote,
       SPIRV_C_GroupNonUniformArithmetic, SPIRV_C_GroupNonUniformBallot,

 
 def SPIRV_ArrayedAttr : SPIRV_I32EnumAttr<
   "ImageArrayedInfo", "valid SPIR-V Image Arrayed specification",
   "image_arrayed_info", [SPIRV_IAI_NonArrayed, SPIRV_IAI_Arrayed]>;
 
 def SPIRV_ISI_SingleSampled : I32EnumAttrCase<"SingleSampled", 0>;
 def SPIRV_ISI_MultiSampled  : I32EnumAttrCase<"MultiSampled", 1>;
 
-def SPIRV_SamplingAttr: SPIRV_I32EnumAttr<
+def SPIRV_SamplingAttr : SPIRV_I32EnumAttr<
   "ImageSamplingInfo", "valid SPIR-V Image Sampling specification",
   "image_sampling_info", [SPIRV_ISI_SingleSampled, SPIRV_ISI_MultiSampled]>;
 
 def SPIRV_ISUI_SamplerUnknown : I32EnumAttrCase<"SamplerUnknown", 0>;
 def SPIRV_ISUI_NeedSampler    : I32EnumAttrCase<"NeedSampler", 1>;
 def SPIRV_ISUI_NoSampler      : I32EnumAttrCase<"NoSampler", 2>;
 
 def SPIRV_SamplerUseAttr: SPIRV_I32EnumAttr<
   "ImageSamplerUseInfo", "valid SPIR-V Sampler Use specification",
   "image_sampler_use_info",
   [SPIRV_ISUI_SamplerUnknown, SPIRV_ISUI_NeedSampler, SPIRV_ISUI_NoSampler]>;
 
 def SPIRV_ML_ColumnMajor : I32EnumAttrCase<"ColumnMajor", 0>;
 def SPIRV_ML_RowMajor    : I32EnumAttrCase<"RowMajor", 1>;
 def SPIRV_ML_PackedA     : I32EnumAttrCase<"PackedA", 2>;
 def SPIRV_ML_PackedB     : I32EnumAttrCase<"PackedB", 3>;
 
-def SPIRV_MatrixLayoutAttr  :
+def SPIRV_MatrixLayoutAttr :
     SPIRV_I32EnumAttr<"MatrixLayout", "valid SPIR-V MatrixLayout", "matrixLayout", [
       SPIRV_ML_ColumnMajor, SPIRV_ML_RowMajor, SPIRV_ML_PackedA, SPIRV_ML_PackedB
     ]>;
 
+// Cooperative Matrix Use for the SPV_KHR_cooperative_matrix extension.
+def SPIRV_KHR_CMU_MatrixA   : I32EnumAttrCase<"MatrixA", 0>;
+def SPIRV_KHR_CMU_MatrixB   : I32EnumAttrCase<"MatrixB", 1>;
+def SPIRV_KHR_CMU_MatrixAcc : I32EnumAttrCase<"MatrixAcc", 2>;
+
+def SPIRV_KHR_CooperativeMatrixUseAttr :
+    SPIRV_I32EnumAttr<"CooperativeMatrixUseKHR",
+                      "valid SPIR-V Cooperative Matrix Use (KHR)",
+                      "coop_matrix_use_khr", [
+      SPIRV_KHR_CMU_MatrixA, SPIRV_KHR_CMU_MatrixB, SPIRV_KHR_CMU_MatrixAcc
+    ]>;
+
 //===----------------------------------------------------------------------===//
 // SPIR-V attribute definitions
 //===----------------------------------------------------------------------===//
 
 def SPIRV_VerCapExtAttr : DialectAttr<
     SPIRV_Dialect,
     CPred<"::llvm::isa<::mlir::spirv::VerCapExtAttr>($_self)">,
     "version-capability-extension attribute"> {
class IOrUI<int width>
   int bitwidth = width;
 }
 
 class SignlessOrUnsignedIntOfWidths<list<int> widths> :
     AnyTypeOf<!foreach(w, widths, IOrUI<w>),
               !interleave(widths, "/") # "-bit signless/unsigned integer">;
 
 def SPIRV_IsArrayType : CPred<"::llvm::isa<::mlir::spirv::ArrayType>($_self)">;
+def SPIRV_IsCooperativeMatrixType :
+  CPred<"::llvm::isa<::mlir::spirv::CooperativeMatrixType>($_self)">;
 def SPIRV_IsCooperativeMatrixNVType :
   CPred<"::llvm::isa<::mlir::spirv::CooperativeMatrixNVType>($_self)">;
 def SPIRV_IsImageType : CPred<"::llvm::isa<::mlir::spirv::ImageType>($_self)">;
 def SPIRV_IsJointMatrixType :
   CPred<"::llvm::isa<::mlir::spirv::JointMatrixINTELType>($_self)">;
 def SPIRV_IsMatrixType : CPred<"::llvm::isa<::mlir::spirv::MatrixType>($_self)">;
 def SPIRV_IsPtrType : CPred<"::llvm::isa<::mlir::spirv::PointerType>($_self)">;
 def SPIRV_IsRTArrayType : CPred<"::llvm::isa<::mlir::spirv::RuntimeArrayType>($_self)">;

 def SPIRV_Vector : VectorOfLengthAndType<[2, 3, 4, 8, 16],
                                        [SPIRV_Bool, SPIRV_Integer, SPIRV_Float]>;
 // Component type check is done in the type parser for the following SPIR-V
 // dialect-specific types so we use "Any" here.
 def SPIRV_AnyPtr : DialectType<SPIRV_Dialect, SPIRV_IsPtrType,
                              "any SPIR-V pointer type">;
 def SPIRV_AnyArray : DialectType<SPIRV_Dialect, SPIRV_IsArrayType,
                                "any SPIR-V array type">;
+def SPIRV_AnyCooperativeMatrix : DialectType<SPIRV_Dialect,
+                                   SPIRV_IsCooperativeMatrixType,
+                                  "any SPIR-V cooperative matrix type">;
 def SPIRV_AnyCooperativeMatrixNV : DialectType<SPIRV_Dialect,
                                      SPIRV_IsCooperativeMatrixNVType,
                                      "any SPIR-V NV cooperative matrix type">;
 def SPIRV_AnyImage : DialectType<SPIRV_Dialect, SPIRV_IsImageType,
                                 "any SPIR-V image type">;
 def SPIRV_AnyJointMatrix : DialectType<SPIRV_Dialect, SPIRV_IsJointMatrixType,
                                 "any SPIR-V joint matrix type">;
 def SPIRV_AnyMatrix : DialectType<SPIRV_Dialect, SPIRV_IsMatrixType,
                                 "any SPIR-V matrix type">;
 def SPIRV_AnyRTArray : DialectType<SPIRV_Dialect, SPIRV_IsRTArrayType,
                                  "any SPIR-V runtime array type">;
 def SPIRV_AnyStruct : DialectType<SPIRV_Dialect, SPIRV_IsStructType,
                                 "any SPIR-V struct type">;
 def SPIRV_AnySampledImage : DialectType<SPIRV_Dialect, SPIRV_IsSampledImageType,
                                 "any SPIR-V sampled image type">;
 
 def SPIRV_Numerical : AnyTypeOf<[SPIRV_Integer, SPIRV_Float]>;
 def SPIRV_Scalar : AnyTypeOf<[SPIRV_Numerical, SPIRV_Bool]>;
 def SPIRV_Aggregate : AnyTypeOf<[SPIRV_AnyArray, SPIRV_AnyRTArray, SPIRV_AnyStruct]>;
 def SPIRV_Composite :
     AnyTypeOf<[SPIRV_Vector, SPIRV_AnyArray, SPIRV_AnyRTArray, SPIRV_AnyStruct,
-               SPIRV_AnyCooperativeMatrixNV, SPIRV_AnyJointMatrix, SPIRV_AnyMatrix]>;
+               SPIRV_AnyCooperativeMatrix, SPIRV_AnyCooperativeMatrixNV,
+               SPIRV_AnyJointMatrix, SPIRV_AnyMatrix]>;
 def SPIRV_Type : AnyTypeOf<[
     SPIRV_Void, SPIRV_Bool, SPIRV_Integer, SPIRV_Float, SPIRV_Vector,
     SPIRV_AnyPtr, SPIRV_AnyArray, SPIRV_AnyRTArray, SPIRV_AnyStruct,
-    SPIRV_AnyCooperativeMatrixNV, SPIRV_AnyJointMatrix, SPIRV_AnyMatrix,
-    SPIRV_AnySampledImage
+    SPIRV_AnyCooperativeMatrix, SPIRV_AnyCooperativeMatrixNV,
+    SPIRV_AnyJointMatrix, SPIRV_AnyMatrix, SPIRV_AnySampledImage
   ]>;
 
 def SPIRV_SignedInt : SignedIntOfWidths<[8, 16, 32, 64]>;
 def SPIRV_SignlessOrUnsignedInt : SignlessOrUnsignedIntOfWidths<[8, 16, 32, 64]>;
 
+class SPIRV_CoopMatrixOfType<list<Type> allowedTypes> :
+  ContainerType<AnyTypeOf<allowedTypes>, SPIRV_IsCooperativeMatrixType,
+    "::llvm::cast<::mlir::spirv::CooperativeMatrixType>($_self).getElementType()",
+    "Cooperative Matrix">;
+
 class SPIRV_CoopMatrixNVOfType<list<Type> allowedTypes> :
   ContainerType<AnyTypeOf<allowedTypes>, SPIRV_IsCooperativeMatrixNVType,
     "::llvm::cast<::mlir::spirv::CooperativeMatrixNVType>($_self).getElementType()",
     "Cooperative Matrix NV">;
 
 class SPIRV_JointMatrixOfType<list<Type> allowedTypes> :
   ContainerType<AnyTypeOf<allowedTypes>, SPIRV_IsJointMatrixType,
     "::llvm::cast<::mlir::spirv::JointMatrixINTELType>($_self).getElementType()",
     "Joint Matrix">;
 
 class SPIRV_ScalarOrVectorOf<Type type> :
     AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4, 8, 16], [type]>]>;
 
 class SPIRV_ScalarOrVectorOrCoopMatrixOf<Type type> :
     AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4, 8, 16], [type]>,
-               SPIRV_CoopMatrixNVOfType<[type]>]>;
+               SPIRV_CoopMatrixOfType<[type]>, SPIRV_CoopMatrixNVOfType<[type]>]>;
 
 class SPIRV_MatrixOrCoopMatrixOf<Type type> :
-    AnyTypeOf<[SPIRV_AnyMatrix, SPIRV_CoopMatrixNVOfType<[type]>]>;
+    AnyTypeOf<[SPIRV_AnyMatrix,
+               SPIRV_CoopMatrixOfType<[type]>,
+               SPIRV_CoopMatrixNVOfType<[type]>]>;
 
 def SPIRV_ScalarOrVector : AnyTypeOf<[SPIRV_Scalar, SPIRV_Vector]>;
 def SPIRV_ScalarOrVectorOrPtr : AnyTypeOf<[SPIRV_ScalarOrVector, SPIRV_AnyPtr]>;
 
 class SPIRV_Vec4<Type type> : VectorOfLengthAndType<[4], [type]>;
 def SPIRV_IntVec4 : SPIRV_Vec4<SPIRV_Integer>;
 def SPIRV_IOrUIVec4 : SPIRV_Vec4<SPIRV_SignlessOrUnsignedInt>;
 def SPIRV_Int32Vec4 : SPIRV_Vec4<AnyI32>;

 def SPIRV_OC_OpGroupNonUniformFMax        : I32EnumAttrCase<"OpGroupNonUniformFMax", 358>;
 def SPIRV_OC_OpSubgroupBallotKHR          : I32EnumAttrCase<"OpSubgroupBallotKHR", 4421>;
 def SPIRV_OC_OpSDot                       : I32EnumAttrCase<"OpSDot", 4450>;
 def SPIRV_OC_OpUDot                       : I32EnumAttrCase<"OpUDot", 4451>;
 def SPIRV_OC_OpSUDot                      : I32EnumAttrCase<"OpSUDot", 4452>;
 def SPIRV_OC_OpSDotAccSat                 : I32EnumAttrCase<"OpSDotAccSat", 4453>;
 def SPIRV_OC_OpUDotAccSat                 : I32EnumAttrCase<"OpUDotAccSat", 4454>;
 def SPIRV_OC_OpSUDotAccSat                : I32EnumAttrCase<"OpSUDotAccSat", 4455>;
+def SPIRV_OC_OpTypeCooperativeMatrixKHR   : I32EnumAttrCase<"OpTypeCooperativeMatrixKHR", 4456>;
+def SPIRV_OC_OpCooperativeMatrixLengthKHR : I32EnumAttrCase<"OpCooperativeMatrixLengthKHR", 4460>;
 def SPIRV_OC_OpTypeCooperativeMatrixNV    : I32EnumAttrCase<"OpTypeCooperativeMatrixNV", 5358>;
 def SPIRV_OC_OpCooperativeMatrixLoadNV    : I32EnumAttrCase<"OpCooperativeMatrixLoadNV", 5359>;
 def SPIRV_OC_OpCooperativeMatrixStoreNV   : I32EnumAttrCase<"OpCooperativeMatrixStoreNV", 5360>;
 def SPIRV_OC_OpCooperativeMatrixMulAddNV  : I32EnumAttrCase<"OpCooperativeMatrixMulAddNV", 5361>;
 def SPIRV_OC_OpCooperativeMatrixLengthNV  : I32EnumAttrCase<"OpCooperativeMatrixLengthNV", 5362>;
 def SPIRV_OC_OpSubgroupBlockReadINTEL     : I32EnumAttrCase<"OpSubgroupBlockReadINTEL", 5575>;
 def SPIRV_OC_OpSubgroupBlockWriteINTEL    : I32EnumAttrCase<"OpSubgroupBlockWriteINTEL", 5576>;
 def SPIRV_OC_OpAssumeTrueKHR              : I32EnumAttrCase<"OpAssumeTrueKHR", 5630>;
SPIRV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", "opcode", [
       SPIRV_OC_OpGroupNonUniformShuffleUp, SPIRV_OC_OpGroupNonUniformShuffleDown,
       SPIRV_OC_OpGroupNonUniformIAdd, SPIRV_OC_OpGroupNonUniformFAdd,
       SPIRV_OC_OpGroupNonUniformIMul, SPIRV_OC_OpGroupNonUniformFMul,
       SPIRV_OC_OpGroupNonUniformSMin, SPIRV_OC_OpGroupNonUniformUMin,
       SPIRV_OC_OpGroupNonUniformFMin, SPIRV_OC_OpGroupNonUniformSMax,
       SPIRV_OC_OpGroupNonUniformUMax, SPIRV_OC_OpGroupNonUniformFMax,
       SPIRV_OC_OpSubgroupBallotKHR, SPIRV_OC_OpSDot, SPIRV_OC_OpUDot,
       SPIRV_OC_OpSUDot, SPIRV_OC_OpSDotAccSat, SPIRV_OC_OpUDotAccSat,
-      SPIRV_OC_OpSUDotAccSat, SPIRV_OC_OpTypeCooperativeMatrixNV,
+      SPIRV_OC_OpSUDotAccSat, SPIRV_OC_OpCooperativeMatrixLengthKHR,
+      SPIRV_OC_OpTypeCooperativeMatrixNV,
       SPIRV_OC_OpCooperativeMatrixLoadNV, SPIRV_OC_OpCooperativeMatrixStoreNV,
       SPIRV_OC_OpCooperativeMatrixMulAddNV, SPIRV_OC_OpCooperativeMatrixLengthNV,
       SPIRV_OC_OpSubgroupBlockReadINTEL, SPIRV_OC_OpSubgroupBlockWriteINTEL,
       SPIRV_OC_OpAssumeTrueKHR, SPIRV_OC_OpAtomicFAddEXT, SPIRV_OC_OpGroupIMulKHR,
       SPIRV_OC_OpGroupFMulKHR,
 
       SPIRV_OC_OpTypeJointMatrixINTEL, SPIRV_OC_OpJointMatrixLoadINTEL,
       SPIRV_OC_OpJointMatrixStoreINTEL, SPIRV_OC_OpJointMatrixMadINTEL,

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td

 //===- SPIRVCooperativeMatrixOps.td - cooperative matmul ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is the op definition spec of cooperative matrix multiply extension ops.
+// We support both cooperative matrix extensions:
+//  - SPV_NV_cooperative_matrix
+//  - SPV_KHR_cooperative_matrix
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
 #define MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
 
 //===----------------------------------------------------------------------===//
+// SPV_KHR_cooperative_matrix extension ops.
+//===----------------------------------------------------------------------===//
+
+// -----
+
+def SPIRV_KHRCooperativeMatrixLengthOp :
+      SPIRV_KhrVendorOp<"CooperativeMatrixLength", [Pure]> {
+  let summary = "Queries the number of cooperative matrix components";
+
+  let description = [{
+    Number of components of a cooperative matrix type accessible to each
+    invocation when treated as a composite.
+
+    The type attribute must be a cooperative matrix type.
+
+    ``` {.ebnf}
+    cooperative-matrix-length-op ::= ssa-id `=` `spirv.KHR.CooperativeMatrixLength
+                                    ` : ` cooperative-matrix-type
+    ```
+
+    #### Example:
+
+    ```
+    %0 = spirv.KHR.CooperativeMatrixLength :
+           !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+    ```
+  }];
+
+  let assemblyFormat = "attr-dict `:` $cooperative_matrix_type";
+
+  let availability = [
+    MinVersion<SPIRV_V_1_6>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[SPV_KHR_cooperative_matrix]>,
+    Capability<[SPIRV_C_CooperativeMatrixKHR]>
+  ];
+
+  let arguments = (ins
+    TypeAttr:$cooperative_matrix_type
+  );
+
+  let results = (outs
+    SPIRV_Int32:$result
+  );
+}
+
+//===----------------------------------------------------------------------===//
 // SPV_NV_cooperative_matrix extension ops.
 //===----------------------------------------------------------------------===//
 
 // -----
 
 def SPIRV_NVCooperativeMatrixLengthOp : SPIRV_NvVendorOp<"CooperativeMatrixLength",
   [Pure]> {
   let summary = "See extension SPV_NV_cooperative_matrix";
def SPIRV_NVCooperativeMatrixLengthOp : SPIRV_NvVendorOp<"CooperativeMatrixLength",
 
   let arguments = (ins
     TypeAttr:$cooperative_matrix_type
   );
 
   let results = (outs
     SPIRV_Int32:$result
   );
-  let hasVerifier = 0;
 }
 
 // -----
 
 def SPIRV_NVCooperativeMatrixLoadOp : SPIRV_NvVendorOp<"CooperativeMatrixLoad", []> {
   let summary = "See extension SPV_NV_cooperative_matrix";
 
   let description = [{

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTypes.h

 
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 
+#include <cstdint>
 #include <tuple>
 
 namespace mlir {
 namespace spirv {
 
 namespace detail {
 struct ArrayTypeStorage;
+struct CooperativeMatrixTypeStorage;
 struct CooperativeMatrixNVTypeStorage;
 struct ImageTypeStorage;
 struct JointMatrixTypeStorage;
 struct MatrixTypeStorage;
 struct PointerTypeStorage;
 struct RuntimeArrayTypeStorage;
 struct SampledImageTypeStorage;
 struct StructTypeStorage;
void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      std::optional<StorageClass> storage = std::nullopt);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        std::optional<StorageClass> storage = std::nullopt);
 };
 
 llvm::hash_code
 hash_value(const StructType::MemberDecorationInfo &memberDecorationInfo);
 
+// SPIR-V KHR cooperative matrix type
+class CooperativeMatrixType
+    : public Type::TypeBase<CooperativeMatrixType, CompositeType,
+                            detail::CooperativeMatrixTypeStorage> {
+public:
+  using Base::Base;
+
+  static CooperativeMatrixType get(Type elementType, uint32_t rows,
+                                   uint32_t columns, Scope scope,
+                                   CooperativeMatrixUseKHR use);
+  Type getElementType() const;
+
+  /// Returns the scope of the matrix.
+  Scope getScope() const;
+  /// Returns the number of rows of the matrix.
+  uint32_t getRows() const;
+  /// Returns the number of columns of the matrix.
+  uint32_t getColumns() const;
+  /// Returns the use parameter of the cooperative matrix.
+  CooperativeMatrixUseKHR getUse() const;
+
+  void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
+                     std::optional<StorageClass> storage = std::nullopt);
+  void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
+                       std::optional<StorageClass> storage = std::nullopt);
+};
+
 // SPIR-V NV cooperative matrix type
 class CooperativeMatrixNVType
     : public Type::TypeBase<CooperativeMatrixNVType, CompositeType,
                             detail::CooperativeMatrixNVTypeStorage> {
 public:
   using Base::Base;
 
   static CooperativeMatrixNVType get(Type elementType, Scope scope,

mlir/lib/Dialect/SPIRV/IR/SPIRVDialect.cpp

if (failed(parseOptionalArrayStride(dialect, parser, stride)))
     return Type();
 
   if (parser.parseGreater())
     return Type();
   return ArrayType::get(elementType, count, stride);
 }
 
 // cooperative-matrix-type ::=
-//   `!spirv.NV.coopmatrix` `<` rows `x` columns `x` element-type ',' scope>
+//   `!spirv.coopmatrix` `<` rows `x` columns `x` element-type `,`
+//                           scope `,` use `>`
 static Type parseCooperativeMatrixType(SPIRVDialect const &dialect,
                                        DialectAsmParser &parser) {
   if (parser.parseLess())
+    return {};
+
+  SmallVector<int64_t, 2> dims;
+  SMLoc countLoc = parser.getCurrentLocation();
+  if (parser.parseDimensionList(dims, /*allowDynamic=*/false))
+    return {};
+
+  if (dims.size() != 2) {
+    parser.emitError(countLoc, "expected row and column count");
+    return {};
+  }
+
+  auto elementTy = parseAndVerifyType(dialect, parser);
+  if (!elementTy)
+    return {};
+
+  Scope scope;
+  if (parser.parseComma() || parseEnumKeywordAttr(scope, parser, "scope <id>"))
+    return {};
+
+  CooperativeMatrixUseKHR use;
+  if (parser.parseComma() || parseEnumKeywordAttr(use, parser, "use <id>"))
+    return {};
+
+  if (parser.parseGreater())
+    return {};
+
+  return CooperativeMatrixType::get(elementTy, dims[0], dims[1], scope, use);
+}
+
+// nv-cooperative-matrix-type ::=
+//   `!spirv.NV.coopmatrix` `<` rows `x` columns `x` element-type `,` scope `>`
+static Type parseCooperativeMatrixNVType(SPIRVDialect const &dialect,
+                                         DialectAsmParser &parser) {
+  if (parser.parseLess())
     return Type();
 
   SmallVector<int64_t, 2> dims;
   SMLoc countLoc = parser.getCurrentLocation();
   if (parser.parseDimensionList(dims, /*allowDynamic=*/false))
     return Type();
 
   if (dims.size() != 2) {

 //              | struct-type
 Type SPIRVDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
 
   if (keyword == "array")
     return parseArrayType(*this, parser);
-  if (keyword == "NV.coopmatrix")
+  if (keyword == "coopmatrix")
     return parseCooperativeMatrixType(*this, parser);
+  if (keyword == "NV.coopmatrix")
+    return parseCooperativeMatrixNVType(*this, parser);
   if (keyword == "jointmatrix")
     return parseJointMatrixType(*this, parser);
   if (keyword == "image")
     return parseImageType(*this, parser);
   if (keyword == "ptr")
     return parsePointerType(*this, parser);
   if (keyword == "rtarray")
     return parseRuntimeArrayType(*this, parser);
static void print(StructType type, DialectAsmPrinter &os) {
   llvm::interleaveComma(llvm::seq<unsigned>(0, type.getNumElements()), os,
                         printMember);
   os << ")>";
 
   if (type.isIdentified())
     structContext.remove(type.getIdentifier());
 }
 
+static void print(CooperativeMatrixType type, DialectAsmPrinter &os) {
+  os << "coopmatrix<" << type.getRows() << "x" << type.getColumns() << "x"
+     << type.getElementType() << ", " << type.getScope() << ", "
+     << type.getUse() << ">";
+}
+
 static void print(CooperativeMatrixNVType type, DialectAsmPrinter &os) {
   os << "NV.coopmatrix<" << type.getRows() << "x" << type.getColumns() << "x";
   os << type.getElementType() << ", " << stringifyScope(type.getScope());
   os << ">";
 }
 
 static void print(JointMatrixINTELType type, DialectAsmPrinter &os) {
   os << "jointmatrix<" << type.getRows() << "x" << type.getColumns() << "x";
   os << type.getElementType() << ", "
      << stringifyMatrixLayout(type.getMatrixLayout());
   os << ", " << stringifyScope(type.getScope()) << ">";
 }
 
 static void print(MatrixType type, DialectAsmPrinter &os) {
   os << "matrix<" << type.getNumColumns() << " x " << type.getColumnType();
   os << ">";
 }
 
 void SPIRVDialect::printType(Type type, DialectAsmPrinter &os) const {
   TypeSwitch<Type>(type)
-      .Case<ArrayType, CooperativeMatrixNVType, JointMatrixINTELType,
-            PointerType, RuntimeArrayType, ImageType, SampledImageType,
-            StructType, MatrixType>([&](auto type) { print(type, os); })
+      .Case<ArrayType, CooperativeMatrixType, CooperativeMatrixNVType,
+            JointMatrixINTELType, PointerType, RuntimeArrayType, ImageType,
+            SampledImageType, StructType, MatrixType>(
+          [&](auto type) { print(type, os); })
       .Default([](Type) { llvm_unreachable("unhandled SPIR-V type"); });
 }
 
 //===----------------------------------------------------------------------===//
 // Constant
 //===----------------------------------------------------------------------===//
 
 Operation *SPIRVDialect::materializeConstant(OpBuilder &builder,

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 #include <cassert>
 #include <numeric>
 
 using namespace mlir;
 
 // TODO: generate these strings using ODS.
 constexpr char kAlignmentAttrName[] = "alignment";
static LogicalResult verifyCastOp(Operation *op,
                                   bool skipBitWidthCheck = false) {
   // Some CastOps have no limit on bit widths for result and operand type.
   if (skipBitWidthCheck)
     return success();
 
   Type operandType = op->getOperand(0).getType();
   Type resultType = op->getResult(0).getType();
 
-  // ODS checks that result type and operand type have the same shape.
-  if (auto vectorType = llvm::dyn_cast<VectorType>(operandType)) {
-    operandType = vectorType.getElementType();
-    resultType = llvm::cast<VectorType>(resultType).getElementType();
-  }
+  // ODS checks that result type and operand type have the same shape. Check
+  // that composite types match and extract the element types, if any.
+  using TypePair = std::pair<Type, Type>;
+  auto [operandElemTy, resultElemTy] =
+      TypeSwitch<Type, TypePair>(operandType)
+          .Case<VectorType, spirv::CooperativeMatrixType,
+                spirv::CooperativeMatrixNVType, spirv::JointMatrixINTELType>(
+              [resultType](auto concreteOperandTy) -> TypePair {
+                if (auto concreteResultTy =
+                        dyn_cast<decltype(concreteOperandTy)>(resultType)) {
+                  return {concreteOperandTy.getElementType(),
+                          concreteResultTy.getElementType()};
+                }
+                return {};
+              })
+          .Default([resultType](Type operandType) -> TypePair {
+            return {operandType, resultType};
+          });
 
-  if (auto coopMatrixType =
-          llvm::dyn_cast<spirv::CooperativeMatrixNVType>(operandType)) {
-    operandType = coopMatrixType.getElementType();
-    resultType =
-        llvm::cast<spirv::CooperativeMatrixNVType>(resultType).getElementType();
-  }
+  if (!operandElemTy || !resultElemTy)
+    return op->emitOpError("incompatible operand and result types");
 
-  if (auto jointMatrixType =
-          llvm::dyn_cast<spirv::JointMatrixINTELType>(operandType)) {
-    operandType = jointMatrixType.getElementType();
-    resultType =
-        llvm::cast<spirv::JointMatrixINTELType>(resultType).getElementType();
-  }
-
-  auto operandTypeBitWidth = operandType.getIntOrFloatBitWidth();
-  auto resultTypeBitWidth = resultType.getIntOrFloatBitWidth();
-  auto isSameBitWidth = operandTypeBitWidth == resultTypeBitWidth;
+  unsigned operandTypeBitWidth = operandElemTy.getIntOrFloatBitWidth();
+  unsigned resultTypeBitWidth = resultElemTy.getIntOrFloatBitWidth();
+  bool isSameBitWidth = operandTypeBitWidth == resultTypeBitWidth;
 
   if (requireSameBitWidth) {
     if (!isSameBitWidth) {
       return op->emitOpError(
                  "expected the same bit widths for operand type and result "
                  "type, but provided ")
-             << operandType << " and " << resultType;
+             << operandElemTy << " and " << resultElemTy;
     }
     return success();
   }
 
   if (isSameBitWidth) {
     return op->emitOpError(
                "expected the different bit widths for operand type and result "
                "type, but provided ")
-           << operandType << " and " << resultType;
+           << operandElemTy << " and " << resultElemTy;
   }
   return success();
 }
 
 template <typename MemoryOpTy>
 static LogicalResult verifyMemoryAccessAttribute(MemoryOpTy memoryOp) {
   // ODS checks for attributes values. Just need to verify that if the
   // memory-access attribute is Aligned, then the alignment attribute must be
if (index >= totalSrcElements &&
       return emitOpError("component selector ")
              << index << " out of range: expected to be in [0, "
              << totalSrcElements << ") or 0xffffffff";
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spirv.KHR.CooperativeMatrixLength
+//===----------------------------------------------------------------------===//
+
+LogicalResult spirv::KHRCooperativeMatrixLengthOp::verify() {
+  if (!isa<spirv::CooperativeMatrixType>(getCooperativeMatrixType())) {
+    return emitOpError(
+               "type attribute must be a '!spirv.coopmatrix' type, found ")
+           << getCooperativeMatrixType() << " instead";
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// spirv.NV.CooperativeMatrixLength
+//===----------------------------------------------------------------------===//
+
+LogicalResult spirv::NVCooperativeMatrixLengthOp::verify() {
+  if (!isa<spirv::CooperativeMatrixNVType>(getCooperativeMatrixType())) {
+    return emitOpError(
+               "type attribute must be a '!spirv.NV.coopmatrix' type, found ")
+           << getCooperativeMatrixType() << " instead";
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // spirv.NV.CooperativeMatrixLoad
 //===----------------------------------------------------------------------===//
 
 ParseResult spirv::NVCooperativeMatrixLoadOp::parse(OpAsmParser &parser,
                                                     OperationState &result) {
   SmallVector<OpAsmParser::UnresolvedOperand, 3> operandInfo;
   Type strideType = parser.getBuilder().getIntegerType(32);
   Type columnMajorType = parser.getBuilder().getIntegerType(1);
void spirv::NVCooperativeMatrixLoadOp::print(OpAsmPrinter &printer) {
   printer << " " << getPointer() << ", " << getStride() << ", "
           << getColumnmajor();
   // Print optional memory access attribute.
   if (auto memAccess = getMemoryAccess())
     printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"]";
   printer << " : " << getPointer().getType() << " as " << getType();
 }
 
-static LogicalResult verifyPointerAndCoopMatrixType(Operation *op, Type pointer,
-                                                    Type coopMatrix) {
+static LogicalResult
+verifyPointerAndCoopMatrixNVType(Operation *op, Type pointer, Type coopMatrix) {
   Type pointeeType = llvm::cast<spirv::PointerType>(pointer).getPointeeType();
   if (!llvm::isa<spirv::ScalarType>(pointeeType) &&
       !llvm::isa<VectorType>(pointeeType))
     return op->emitError(
                "Pointer must point to a scalar or vector type but provided ")
            << pointeeType;
   spirv::StorageClass storage =
       llvm::cast<spirv::PointerType>(pointer).getStorageClass();
   if (storage != spirv::StorageClass::Workgroup &&
       storage != spirv::StorageClass::StorageBuffer &&
       storage != spirv::StorageClass::PhysicalStorageBuffer)
     return op->emitError(
                "Pointer storage class must be Workgroup, StorageBuffer or "
                "PhysicalStorageBufferEXT but provided ")
            << stringifyStorageClass(storage);
   return success();
 }
 
 LogicalResult spirv::NVCooperativeMatrixLoadOp::verify() {
-  return verifyPointerAndCoopMatrixType(*this, getPointer().getType(),
+  return verifyPointerAndCoopMatrixNVType(*this, getPointer().getType(),
                                           getResult().getType());
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.NV.CooperativeMatrixStore
 //===----------------------------------------------------------------------===//
 
 ParseResult spirv::NVCooperativeMatrixStoreOp::parse(OpAsmParser &parser,
                                                      OperationState &result) {
printer << " " << getPointer() << ", " << getObject() << ", " << getStride()
           << ", " << getColumnmajor();
   // Print optional memory access attribute.
   if (auto memAccess = getMemoryAccess())
     printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"]";
   printer << " : " << getPointer().getType() << ", " << getOperand(1).getType();
 }
 
 LogicalResult spirv::NVCooperativeMatrixStoreOp::verify() {
-  return verifyPointerAndCoopMatrixType(*this, getPointer().getType(),
+  return verifyPointerAndCoopMatrixNVType(*this, getPointer().getType(),
                                           getObject().getType());
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.NV.CooperativeMatrixMulAdd
 //===----------------------------------------------------------------------===//
 
 static LogicalResult
-verifyCoopMatrixMulAdd(spirv::NVCooperativeMatrixMulAddOp op) {
+verifyCoopMatrixMulAddNV(spirv::NVCooperativeMatrixMulAddOp op) {
   if (op.getC().getType() != op.getResult().getType())
     return op.emitOpError("result and third operand must have the same type");
   auto typeA = llvm::cast<spirv::CooperativeMatrixNVType>(op.getA().getType());
   auto typeB = llvm::cast<spirv::CooperativeMatrixNVType>(op.getB().getType());
   auto typeC = llvm::cast<spirv::CooperativeMatrixNVType>(op.getC().getType());
   auto typeR =
       llvm::cast<spirv::CooperativeMatrixNVType>(op.getResult().getType());
   if (typeA.getRows() != typeR.getRows() ||
return op.emitOpError(
         "matrix A and B non-integer element types must match");
   }
   if (typeR.getElementType() != typeC.getElementType())
     return op.emitOpError("matrix accumulator element type must match");
   return success();
 }
 
 LogicalResult spirv::NVCooperativeMatrixMulAddOp::verify() {
-  return verifyCoopMatrixMulAdd(*this);
+  return verifyCoopMatrixMulAddNV(*this);
 }
 
+//===----------------------------------------------------------------------===//
+// spirv.INTEL.JointMatrixLoad
+//===----------------------------------------------------------------------===//
+
 static LogicalResult
 verifyPointerAndJointMatrixType(Operation *op, Type pointer, Type jointMatrix) {
   Type pointeeType = llvm::cast<spirv::PointerType>(pointer).getPointeeType();
   if (!llvm::isa<spirv::ScalarType>(pointeeType) &&
       !llvm::isa<VectorType>(pointeeType))
     return op->emitError(
                "Pointer must point to a scalar or vector type but provided ")
            << pointeeType;
   spirv::StorageClass storage =
       llvm::cast<spirv::PointerType>(pointer).getStorageClass();
   if (storage != spirv::StorageClass::Workgroup &&
       storage != spirv::StorageClass::CrossWorkgroup &&
       storage != spirv::StorageClass::UniformConstant &&
       storage != spirv::StorageClass::Generic)
     return op->emitError("Pointer storage class must be Workgroup or "
                          "CrossWorkgroup but provided ")
            << stringifyStorageClass(storage);
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// spirv.INTEL.JointMatrixLoad
-//===----------------------------------------------------------------------===//
-
 LogicalResult spirv::INTELJointMatrixLoadOp::verify() {
   return verifyPointerAndJointMatrixType(*this, getPointer().getType(),
                                          getResult().getType());
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.INTEL.JointMatrixStore
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp

 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 
+#include <cstdint>
 #include <iterator>
 
 using namespace mlir;
 using namespace mlir::spirv;
 
 //===----------------------------------------------------------------------===//
 // ArrayType
 //===----------------------------------------------------------------------===//

 
 //===----------------------------------------------------------------------===//
 // CompositeType
 //===----------------------------------------------------------------------===//
 
 bool CompositeType::classof(Type type) {
   if (auto vectorType = llvm::dyn_cast<VectorType>(type))
     return isValid(vectorType);
-  return llvm::isa<spirv::ArrayType, spirv::CooperativeMatrixNVType,
-                   spirv::JointMatrixINTELType, spirv::MatrixType,
-                   spirv::RuntimeArrayType, spirv::StructType>(type);
+  return llvm::isa<spirv::ArrayType, spirv::CooperativeMatrixType,
+                   spirv::CooperativeMatrixNVType, spirv::JointMatrixINTELType,
+                   spirv::MatrixType, spirv::RuntimeArrayType,
+                   spirv::StructType>(type);
 }
 
 bool CompositeType::isValid(VectorType type) {
   switch (type.getNumElements()) {
   case 2:
   case 3:
   case 4:
   case 8:
   case 16:
     break;
   default:
     return false;
   }
   return type.getRank() == 1 && llvm::isa<ScalarType>(type.getElementType());
 }
 
 Type CompositeType::getElementType(unsigned index) const {
   return TypeSwitch<Type, Type>(*this)
-      .Case<ArrayType, CooperativeMatrixNVType, JointMatrixINTELType,
-            RuntimeArrayType, VectorType>(
+      .Case<ArrayType, CooperativeMatrixType, CooperativeMatrixNVType,
+            JointMatrixINTELType, RuntimeArrayType, VectorType>(
           [](auto type) { return type.getElementType(); })
       .Case<MatrixType>([](MatrixType type) { return type.getColumnType(); })
       .Case<StructType>(
           [index](StructType type) { return type.getElementType(index); })
       .Default(
           [](Type) -> Type { llvm_unreachable("invalid composite type"); });
 }
 
 unsigned CompositeType::getNumElements() const {
   if (auto arrayType = llvm::dyn_cast<ArrayType>(*this))
     return arrayType.getNumElements();
   if (auto matrixType = llvm::dyn_cast<MatrixType>(*this))
     return matrixType.getNumColumns();
   if (auto structType = llvm::dyn_cast<StructType>(*this))
     return structType.getNumElements();
   if (auto vectorType = llvm::dyn_cast<VectorType>(*this))
     return vectorType.getNumElements();
-  if (llvm::isa<CooperativeMatrixNVType>(*this)) {
+  if (llvm::isa<CooperativeMatrixType, CooperativeMatrixNVType>(*this)) {
     llvm_unreachable(
-        "invalid to query number of elements of spirv::CooperativeMatrix type");
+        "invalid to query number of elements of spirv Cooperative Matrix type");
   }
   if (llvm::isa<JointMatrixINTELType>(*this)) {
     llvm_unreachable(
         "invalid to query number of elements of spirv::JointMatrix type");
   }
   if (llvm::isa<RuntimeArrayType>(*this)) {
     llvm_unreachable(
         "invalid to query number of elements of spirv::RuntimeArray type");
   }
   llvm_unreachable("invalid composite type");
 }
 
 bool CompositeType::hasCompileTimeKnownNumElements() const {
-  return !llvm::isa<CooperativeMatrixNVType, JointMatrixINTELType,
-              RuntimeArrayType>(*this);
+  return !llvm::isa<CooperativeMatrixType, CooperativeMatrixNVType,
+                    JointMatrixINTELType, RuntimeArrayType>(*this);
 }
 
 void CompositeType::getExtensions(
     SPIRVType::ExtensionArrayRefVector &extensions,
     std::optional<StorageClass> storage) {
   TypeSwitch<Type>(*this)
-      .Case<ArrayType, CooperativeMatrixNVType, JointMatrixINTELType,
-            MatrixType, RuntimeArrayType, StructType>(
+      .Case<ArrayType, CooperativeMatrixType, CooperativeMatrixNVType,
+            JointMatrixINTELType, MatrixType, RuntimeArrayType, StructType>(
           [&](auto type) { type.getExtensions(extensions, storage); })
       .Case<VectorType>([&](VectorType type) {
         return llvm::cast<ScalarType>(type.getElementType())
             .getExtensions(extensions, storage);
       })
       .Default([](Type) { llvm_unreachable("invalid composite type"); });
 }
 
 void CompositeType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     std::optional<StorageClass> storage) {
   TypeSwitch<Type>(*this)
-      .Case<ArrayType, CooperativeMatrixNVType, JointMatrixINTELType,
-            MatrixType, RuntimeArrayType, StructType>(
+      .Case<ArrayType, CooperativeMatrixType, CooperativeMatrixNVType,
+            JointMatrixINTELType, MatrixType, RuntimeArrayType, StructType>(
           [&](auto type) { type.getCapabilities(capabilities, storage); })
       .Case<VectorType>([&](VectorType type) {
         auto vecSize = getNumElements();
         if (vecSize == 8 || vecSize == 16) {
           static const Capability caps[] = {Capability::Vector16};
           ArrayRef<Capability> ref(caps, std::size(caps));
           capabilities.push_back(ref);
         }
if (auto vectorType = llvm::dyn_cast<VectorType>(*this)) {
     if (!elementSize)
       return std::nullopt;
     return *elementSize * vectorType.getNumElements();
   }
   return std::nullopt;
 }
 
 //===----------------------------------------------------------------------===//
+// CooperativeMatrixType
+//===----------------------------------------------------------------------===//
+
+struct spirv::detail::CooperativeMatrixTypeStorage final : TypeStorage {
+  using KeyTy =
+      std::tuple<Type, uint32_t, uint32_t, Scope, CooperativeMatrixUseKHR>;
+
+  static CooperativeMatrixTypeStorage *
+  construct(TypeStorageAllocator &allocator, const KeyTy &key) {
+    return new (allocator.allocate<CooperativeMatrixTypeStorage>())
+        CooperativeMatrixTypeStorage(key);
+  }
+
+  bool operator==(const KeyTy &key) const {
+    return key == KeyTy(elementType, rows, columns, scope, use);
+  }
+
+  CooperativeMatrixTypeStorage(const KeyTy &key)
+      : elementType(std::get<0>(key)), rows(std::get<1>(key)),
+        columns(std::get<2>(key)), scope(std::get<3>(key)),
+        use(std::get<4>(key)) {}
+
+  Type elementType;
+  uint32_t rows;
+  uint32_t columns;
+  Scope scope;
+  CooperativeMatrixUseKHR use;
+};
+
+CooperativeMatrixType CooperativeMatrixType::get(Type elementType,
+                                                 uint32_t rows,
+                                                 uint32_t columns, Scope scope,
+                                                 CooperativeMatrixUseKHR use) {
+  return Base::get(elementType.getContext(), elementType, rows, columns, scope,
+                   use);
+}
+
+Type CooperativeMatrixType::getElementType() const {
+  return getImpl()->elementType;
+}
+
+uint32_t CooperativeMatrixType::getRows() const { return getImpl()->rows; }
+
+uint32_t CooperativeMatrixType::getColumns() const {
+  return getImpl()->columns;
+}
+
+Scope CooperativeMatrixType::getScope() const { return getImpl()->scope; }
+
+CooperativeMatrixUseKHR CooperativeMatrixType::getUse() const {
+  return getImpl()->use;
+}
+
+void CooperativeMatrixType::getExtensions(
+    SPIRVType::ExtensionArrayRefVector &extensions,
+    std::optional<StorageClass> storage) {
+  llvm::cast<SPIRVType>(getElementType()).getExtensions(extensions, storage);
+  static constexpr Extension exts[] = {Extension::SPV_KHR_cooperative_matrix};
+  extensions.push_back(exts);
+}
+
+void CooperativeMatrixType::getCapabilities(
+    SPIRVType::CapabilityArrayRefVector &capabilities,
+    std::optional<StorageClass> storage) {
+  llvm::cast<SPIRVType>(getElementType())
+      .getCapabilities(capabilities, storage);
+  static constexpr Capability caps[] = {Capability::CooperativeMatrixKHR};
+  capabilities.push_back(caps);
+}
+
+//===----------------------------------------------------------------------===//
 // CooperativeMatrixNVType
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::CooperativeMatrixNVTypeStorage : public TypeStorage {
   using KeyTy = std::tuple<Type, Scope, unsigned, unsigned>;
 
   static CooperativeMatrixNVTypeStorage *
   construct(TypeStorageAllocator &allocator, const KeyTy &key) {
void MatrixType::getCapabilities(
   llvm::cast<SPIRVType>(getColumnType()).getCapabilities(capabilities, storage);
 }
 
 //===----------------------------------------------------------------------===//
 // SPIR-V Dialect
 //===----------------------------------------------------------------------===//
 
 void SPIRVDialect::registerTypes() {
-  addTypes<ArrayType, CooperativeMatrixNVType, ImageType, JointMatrixINTELType,
-           MatrixType, PointerType, RuntimeArrayType, SampledImageType,
-           StructType>();
+  addTypes<ArrayType, CooperativeMatrixType, CooperativeMatrixNVType, ImageType,
+           JointMatrixINTELType, MatrixType, PointerType, RuntimeArrayType,
+           SampledImageType, StructType>();
 }

mlir/test/Dialect/SPIRV/IR/cast-ops.mlir

 func.func @convert_f_to_u_vector(%arg0 : vector<3xf32>) -> vector<3xi32> {
   // CHECK: {{%.*}} = spirv.ConvertFToU {{%.*}} : vector<3xf32> to vector<3xi32>
   %0 = spirv.ConvertFToU %arg0 : vector<3xf32> to vector<3xi32>
   spirv.ReturnValue %0 : vector<3xi32>
 }
 
 // -----
 
+func.func @convert_f_to_u.coopmatrix(%arg0 : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixB>) {
+  // CHECK: {{%.*}} = spirv.ConvertFToU {{%.*}} : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixB> to !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
+  %0 = spirv.ConvertFToU %arg0 : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixB> to !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
+  spirv.Return
+}
+
+// -----
+
 func.func @convert_f_to_u_NV.coopmatrix(%arg0 : !spirv.NV.coopmatrix<8x16xf32, Subgroup>) {
   // CHECK: {{%.*}} = spirv.ConvertFToU {{%.*}} : !spirv.NV.coopmatrix<8x16xf32, Subgroup> to !spirv.NV.coopmatrix<8x16xi32, Subgroup>
   %0 = spirv.ConvertFToU %arg0 : !spirv.NV.coopmatrix<8x16xf32, Subgroup> to !spirv.NV.coopmatrix<8x16xi32, Subgroup>
   spirv.Return
 }
 
 // -----
 

 func.func @f_convert_vector(%arg0 : vector<3xf32>) -> vector<3xf64> {
   // CHECK: {{%.*}} = spirv.FConvert {{%.*}} : vector<3xf32> to vector<3xf64>
   %0 = spirv.FConvert %arg0 : vector<3xf32> to vector<3xf64>
   spirv.ReturnValue %0 : vector<3xf64>
 }
 
 // -----
 
-func.func @f_convert_coop_matrix(%arg0 : !spirv.NV.coopmatrix<8x16xf32, Subgroup>) {
+func.func @f_convert_coop_matrix(%arg0 : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixA>) {
+  // CHECK: {{%.*}} = spirv.FConvert {{%.*}} : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixA> to !spirv.coopmatrix<8x16xf64, Subgroup, MatrixA>
+  %0 = spirv.FConvert %arg0 : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixA> to !spirv.coopmatrix<8x16xf64, Subgroup, MatrixA>
+  spirv.Return
+}
+
+// -----
+
+func.func @f_convert_coop_matrix_nv(%arg0 : !spirv.NV.coopmatrix<8x16xf32, Subgroup>) {
   // CHECK: {{%.*}} = spirv.FConvert {{%.*}} : !spirv.NV.coopmatrix<8x16xf32, Subgroup> to !spirv.NV.coopmatrix<8x16xf64, Subgroup>
   %0 = spirv.FConvert %arg0 : !spirv.NV.coopmatrix<8x16xf32, Subgroup> to !spirv.NV.coopmatrix<8x16xf64, Subgroup>
   spirv.Return
 }
 
 // -----
 
 func.func @f_convert_vector(%arg0 : f32) -> f32 {
   // expected-error @+1 {{expected the different bit widths for operand type and result type, but provided 'f32' and 'f32'}}
   %0 = spirv.FConvert %arg0 : f32 to f32
   spirv.ReturnValue %0 : f32
 }
 
 // -----
 
+func.func @f_convert_coop_matrix_to_nv_coop_matrix(%arg0 : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixAcc>) {
+  // expected-error @+1 {{incompatible operand and result types}}
+  %0 = spirv.FConvert %arg0 : !spirv.coopmatrix<8x16xf32, Subgroup, MatrixAcc> to !spirv.NV.coopmatrix<8x16xf64, Subgroup>
+  spirv.Return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spirv.SConvert
 //===----------------------------------------------------------------------===//
 
 func.func @s_convert_scalar(%arg0 : i32) -> i64 {
   // CHECK: {{%.*}} = spirv.SConvert {{%.*}} : i32 to i64
   %0 = spirv.SConvert %arg0 : i32 to i64
   spirv.ReturnValue %0 : i64

mlir/test/Dialect/SPIRV/IR/cooperative-matrix-ops.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect --split-input-file --verify-diagnostics %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// CooperativeMatrix (KHR)
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @cooperative_matrix_length
+spirv.func @cooperative_matrix_length() -> i32 "None" {
+  // CHECK: {{%.*}} = spirv.KHR.CooperativeMatrixLength : !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+  %0 = spirv.KHR.CooperativeMatrixLength : !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+  spirv.ReturnValue %0 : i32
+}
+
+// -----
+
+spirv.func @cooperative_matrix_length_wrong_matrix() -> i32 "None" {
+  // expected-error @+1 {{'spirv.KHR.CooperativeMatrixLength' op type attribute must be a '!spirv.coopmatrix'}}
+  %0 = spirv.KHR.CooperativeMatrixLength : !spirv.NV.coopmatrix<8x16xi32, Subgroup>
+  spirv.ReturnValue %0 : i32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// NV.CooperativeMatrix
+//===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @cooperative_matrix_load
 spirv.func @cooperative_matrix_load(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
   // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLoad {{%.*}}, {{%.*}}, {{%.*}} : !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<16x8xi32, Workgroup>
   %0 = spirv.NV.CooperativeMatrixLoad %ptr, %stride, %b : !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<16x8xi32, Workgroup>
   spirv.Return
 }
 
-// -----
 // CHECK-LABEL: @cooperative_matrix_load_memaccess
 spirv.func @cooperative_matrix_load_memaccess(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
   // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLoad {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup>
   %0 = spirv.NV.CooperativeMatrixLoad %ptr, %stride, %b ["Volatile"] : !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup>
   spirv.Return
 }
 
 // CHECK-LABEL: @cooperative_matrix_load_diff_ptr_type

 
 // -----
 
 spirv.func @cooperative_matrix_load_memaccess(%ptr : !spirv.ptr<i32, Function>, %stride : i32, %b : i1) "None" {
   // expected-error @+1 {{Pointer storage class must be Workgroup, StorageBuffer or PhysicalStorageBufferEXT}}
   %0 = spirv.NV.CooperativeMatrixLoad %ptr, %stride, %b : !spirv.ptr<i32, Function> as !spirv.NV.coopmatrix<8x16xi32, Subgroup>
   spirv.Return
 }
+
+// -----
+
+spirv.func @cooperative_matrix_length_wrong_matrix() -> i32 "None" {
+  // expected-error @+1 {{'spirv.NV.CooperativeMatrixLength' op type attribute must be a '!spirv.NV.coopmatrix'}}
+  %0 = spirv.NV.CooperativeMatrixLength : !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
+  spirv.ReturnValue %0 : i32
+}

mlir/test/Dialect/SPIRV/IR/types.mlir

 //   struct a { struct b *bPtr; };
 //   struct b { struct a *aPtr, struct b *bPtr; };
 // CHECK: func private @id_struct_recursive(!spirv.struct<a10, (!spirv.ptr<!spirv.struct<b10, (!spirv.ptr<!spirv.struct<a10>, Uniform>, !spirv.ptr<!spirv.struct<b10>, Uniform>)>, Uniform>)>)
 func.func private @id_struct_recursive(!spirv.struct<a10, (!spirv.ptr<!spirv.struct<b10, (!spirv.ptr<!spirv.struct<a10>, Uniform>, !spirv.ptr<!spirv.struct<b10>, Uniform>)>, Uniform>)>) -> ()
 
 // -----
 
 //===----------------------------------------------------------------------===//
-// CooperativeMatrix
+// CooperativeMatrix (KHR)
 //===----------------------------------------------------------------------===//
 
-// CHECK: func private @coop_matrix_type(!spirv.NV.coopmatrix<8x16xi32, Subgroup>, !spirv.NV.coopmatrix<8x8xf32, Workgroup>)
-func.func private @coop_matrix_type(!spirv.NV.coopmatrix<8x16xi32, Subgroup>, !spirv.NV.coopmatrix<8x8xf32, Workgroup>) -> ()
+// CHECK-LABEL: func private @coop_matrix_types
+// CHECK-SAME:    !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
+// CHECK-SAME:    !spirv.coopmatrix<8x8xf32, Workgroup, MatrixB>
+// CHECK-SAME:    !spirv.coopmatrix<4x8xf32, Workgroup, MatrixAcc>
+func.func private @coop_matrix_types(!spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>,
+                                     !spirv.coopmatrix<8x8xf32, Workgroup, MatrixB>,
+                                     !spirv.coopmatrix<4x8xf32, Workgroup, MatrixAcc>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected valid keyword}}
+func.func private @missing_scope(!spirv.coopmatrix<8x8xi32, >) -> ()
+
+// -----
+
+// expected-error @+1 {{expected ','}}
+func.func private @missing_use(!spirv.coopmatrix<8x16xi32, Subgroup>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected valid keyword}}
+func.func private @missing_use2(!spirv.coopmatrix<8x8xi32, Subgroup,>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected row and column count}}
+func.func private @missing_count(!spirv.coopmatrix<8xi32, Subgroup, MatrixA>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected row and column count}}
+func.func private @too_many_dims(!spirv.coopmatrix<8x16x32xi32, Subgroup, MatrixB>) -> ()
+
+// -----
+
+// expected-error @+1 {{invalid use <id> attribute specification: Subgroup}}
+func.func private @use_not_integer(!spirv.coopmatrix<8x8xi32, Subgroup, Subgroup>) -> ()
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// NV.CooperativeMatrix
+//===----------------------------------------------------------------------===//
+
+// CHECK: func private @nv_coop_matrix_type(!spirv.NV.coopmatrix<8x16xi32, Subgroup>, !spirv.NV.coopmatrix<8x8xf32, Workgroup>)
+func.func private @nv_coop_matrix_type(!spirv.NV.coopmatrix<8x16xi32, Subgroup>, !spirv.NV.coopmatrix<8x8xf32, Workgroup>) -> ()
 
 // -----
 
 // expected-error @+1 {{expected ','}}
 func.func private @missing_scope(!spirv.NV.coopmatrix<8x16xi32>) -> ()
 
 // -----
 

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

os << tabs
        << formatv("{0}.push_back(opBuilder.getNamedAttr(\"{1}\", "
                   "opBuilder.getArrayAttr(attrListElems)));\n",
                   attrList, attrName);
   } else if (attr.getAttrDefName() == "I32Attr") {
     os << tabs
        << formatv("{0}.push_back(opBuilder.getNamedAttr(\"{1}\", "
                   "opBuilder.getI32IntegerAttr({2}[{3}++])));\n",
                   attrList, attrName, words, wordIndex);
-  } else if (attr.isEnumAttr() || attr.getAttrDefName() == "TypeAttr") {
+  } else if (attr.isEnumAttr() || attr.isTypeAttr()) {
     os << tabs
        << formatv("{0}.push_back(opBuilder.getNamedAttr(\"{1}\", "
                   "TypeAttr::get(getType({2}[{3}++]))));\n",
                   attrList, attrName, words, wordIndex);
   } else {
     PrintFatalError(
         loc, llvm::Twine(
                  "unhandled attribute type in deserialization generation : '") +
-                 attr.getAttrDefName() + llvm::Twine("'"));
+                 attrName + llvm::Twine("'"));
   }
 }
 
 /// Generates the code to deserialize the result of an SPIRV_Op `op` into
 /// `os`. The generated code gets the type of the result specified at
 /// `words`[`wordIndex`], the SSA ID for the result at position `wordIndex` + 1
 /// and updates the `resultType` and `valueID` with the parsed type and SSA ID,
 /// respectively.