[mlir][spirv] Add remaining cooperative matrix instructions.

[ThomasRaoux]

Summary

Those are changes on top of D80043 which add cooperative matrix basic infrastructure. This adds support for cooperative matrix support for arithmetic and cast instructions. It also adds cooperative matrix store, muladd and matrixlength instructions which are part of the extension.

[antiagainst]

Nice! Thanks Thomas!

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

3006

Nit: can this be ordered after SPV_Integer?

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

56

I think we can add

let assemblyFormat = "attr-dict `:` type($result)";

So that we don't need to manually write the parser and printer?

133

I think we need verification on $pointer. At the moment every pointer will be accepted.

175

ssa-use

176

One general rule in MLIR regarding assembly is that it should be parsable on its own. I think we need to give at least two types (for $a and $b) for this? Otherwise by only looking at this op, I'm not able to tell what the type is for $a and $b.

203

I see we don't do any validation on these cooperative matrix ops. Do you plan to add them later?

Here we should at least additionally check that 1) $a, $b, $c, and $result has consistent dimensions: MxK, KxN, MxN, MxN, and 2) they are at the same scope. Basically going through the description and verify what's written there is checked. :)

204

Similarly here I think we can do

let assemblyFormat = "$a, $b, $c `:` type($a), type($b), type($c, $result)

(type($c) is kinda redundant given that we can deduce it but I don't care that much. Also it makes the mapping very clear.)

265

Similarly we need to verify the requirements on $pointer.

[ThomasRaoux]

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

56

Here type is not the type of $result so this syntax wouldn't work. Note sure if there is a way to get it to pick up type for the argument?

133

I added some verification for both load and store pointer type however I didn't find an easy way to test the verification as the invalid IR would fail parsing.

176

My bad I had wrongly assumed the type were matching. I added 2 types and a result type even though the result type could be deducted from the operand types.
I wasn't able to get the assemblyFormat tow work as I could find a way to make the type of c the same as type of result. (type($c, $result) doesn't work) I haven't look very deeply though I can investigate more if you want. Fow now I left the custom parsing/printing functions.

I also added a verify function for muladd and some tests along with it. I was planning to add those later to keep the patch small but it is better to add it now so that I don't forget.

[antiagainst]

Cool! Sorry just a few more nits.

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

56

Ah you are right. Actually the op takes in the id for a type directly. Sorry missed that.

176

SGTM. The assembly format is a bit opaque; River added support for it so he knows all the nitty gritty. Feel free to ping him later if you have questions there. (But he is OOO ATM I think.)

179

This need to be updated right?

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2672

Nit: ...And... ?

2677

s/snd/and/
s/coop/cooperative/

For error messages it's better to be clear; so spelling out things completely is preferrable.

2766

Nit: In general MLIR does not use capitalized variable names. So I'd suggest to s/M/coopMatrix/ or something.

2777

result and the third ... ?

No need to have the trailing dot here.

2782

We can do unconditional cast in the above and omit the check here. These are guaranteed by ODS. Normally the constraints in ODS does not need to be repeated here. We just need to check additional constraints using C++ here.

2789

In general error messages should stay all lower-case in MLIR. I think that's also the case for Clang/etc. So here: "matrix ..."

Besides, it's better to be consistent so either "matrix .. mismatch" or "matrix ... must match" for this and the following two cases.

[ThomasRaoux]

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

176

Sounds good. I'll chat with him when I get a chance to see if this can be improved.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2738

nit: Drop trivial braces.

2746

This looks like it should use the declarative format:

https://mlir.llvm.org/docs/OpDefinitions/#declarative-assembly-format

2756

Same here and a few others.

[ThomasRaoux]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2738

I'll send a patch to address this.

2746

River, I wasn't able to get a syntax working for this instruction as the type itself is an argument and is unrelated to the destination. Do yo have any suggestion to make it work?

2756

Here also I couldn't get the automatic print/parser to work as I have an instruction
r = matmuladd a, b, c : type(a), type(b) -> type(r and c)
I wasn't able to express the fact that c and r must be the same type. Is there a way to make it work?

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2746

If the type of the result is inferrable, you don't have to specify it. Inferable types are either via traits(e.g., SameOperands) or with buildable types(i.e. statically constructible like index, i32, etc.). In this case, the result type is SPV_Int32(i.e. just an alias i32) which is a buildable type. To get the same syntax, all you really need to do it seems is something like:

attr-dict `:` $type

(You would have seen this if you read the section that I sent you : https://mlir.llvm.org/docs/OpDefinitions/#type-inference)

2756

In https://mlir.llvm.org/docs/OpDefinitions/#type-inference, it lists several traits that are supported for listing type equality constraints. You can do simple equality using AllTypesMatch<["c", "r"]> where r and c are the names of the operands/results/etc. With that types can be inferred for all as long as one can be inferred. I would expect the following to work:

def SPV_CooperativeMatrixMulAddNVOp : SPV_Op<..., [..., AllTypesMatch<"c", "result">]> {
  ...
  let assemblyFormat = "operands `:` type($a), type($b) -> type($c)";
}

mlir/include/mlir/Dialect/SPIRV/SPIRVArithmeticOps.td

 include "mlir/Dialect/SPIRV/SPIRVBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class SPV_ArithmeticBinaryOp<string mnemonic, Type type,
                        list<OpTrait> traits = []> :
       // Operands type same as result type.
       SPV_BinaryOp<mnemonic, type, type,
                    !listconcat(traits,
-                               [NoSideEffect, SameOperandsAndResultType])>;
+                               [NoSideEffect, SameOperandsAndResultType])> {
+  // In addition to normal types arithmetic instructions can support cooperative
+  // matrix.
+  let arguments = (ins
+    SPV_ScalarOrVectorOrCoopMatrixOf<type>:$operand1,
+    SPV_ScalarOrVectorOrCoopMatrixOf<type>:$operand2
+  );
+
+  let results = (outs
+    SPV_ScalarOrVectorOrCoopMatrixOf<type>:$result
+  );
+}
 
 class SPV_ArithmeticUnaryOp<string mnemonic, Type type,
                             list<OpTrait> traits = []> :
       // Operand type same as result type.
       SPV_UnaryOp<mnemonic, type, type,
                    !listconcat(traits,
                                [NoSideEffect, SameOperandsAndResultType])>;
 

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

 def SPV_IsRTArrayType : CPred<"$_self.isa<::mlir::spirv::RuntimeArrayType>()">;
 def SPV_IsStructType : CPred<"$_self.isa<::mlir::spirv::StructType>()">;
 
 // See https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html#_types
 // for the definition of the following types and type categories.
 
 def SPV_Void : TypeAlias<NoneType, "void">;
 def SPV_Bool : TypeAlias<I1, "bool">;
+def SPV_Int32 : TypeAlias<I32, "Int32">;
 def SPV_Integer : AnyIntOfWidths<[8, 16, 32, 64]>;
 def SPV_Float : FloatOfWidths<[16, 32, 64]>;
 def SPV_Float16or32 : FloatOfWidths<[16, 32]>;
 def SPV_Vector : VectorOfLengthAndType<[2, 3, 4],
                                        [SPV_Bool, SPV_Integer, SPV_Float]>;
 // Component type check is done in the type parser for the following SPIR-V
 // dialect-specific types so we use "Any" here.
 def SPV_AnyPtr : DialectType<SPIRV_Dialect, SPV_IsPtrType,
def SPV_Composite :
     AnyTypeOf<[SPV_Vector, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct]>;
 def SPV_Type : AnyTypeOf<[
     SPV_Void, SPV_Bool, SPV_Integer, SPV_Float, SPV_Vector,
     SPV_AnyPtr, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct
   ]>;
 
 def SPV_SignlessOrUnsignedInt : SignlessOrUnsignedIntOfWidths<[8, 16, 32, 64]>;
 
+class SPV_CoopMatrixOfType<list<Type> allowedTypes> :
+  ContainerType<AnyTypeOf<allowedTypes>, SPV_IsCooperativeMatrixType,
+    "$_self.cast<::mlir::spirv::CooperativeMatrixNVType>().getElementType()",
+    "Cooperative Matrix">;
+
 class SPV_ScalarOrVectorOf<Type type> :
     AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4], [type]>]>;
 
+class SPV_ScalarOrVectorOrCoopMatrixOf<Type type> :
+    AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4], [type]>,
+               SPV_CoopMatrixOfType<[type]>]>;
+
 def SPV_ScalarOrVector : AnyTypeOf<[SPV_Scalar, SPV_Vector]>;
 def SPV_ScalarOrVectorOrPtr : AnyTypeOf<[SPV_ScalarOrVector, SPV_AnyPtr]>;
 
 class SPV_Vec4<Type type> : VectorOfLengthAndType<[4], [type]>;
 def SPV_IntVec4 : SPV_Vec4<SPV_Integer>;
 def SPV_IOrUIVec4 : SPV_Vec4<SPV_SignlessOrUnsignedInt>;
 def SPV_Int32Vec4 : SPV_Vec4<AnyI32>;
 

 def SPV_OC_OpGroupNonUniformUMin       : I32EnumAttrCase<"OpGroupNonUniformUMin", 354>;
 def SPV_OC_OpGroupNonUniformFMin       : I32EnumAttrCase<"OpGroupNonUniformFMin", 355>;
 def SPV_OC_OpGroupNonUniformSMax       : I32EnumAttrCase<"OpGroupNonUniformSMax", 356>;
 def SPV_OC_OpGroupNonUniformUMax       : I32EnumAttrCase<"OpGroupNonUniformUMax", 357>;
 def SPV_OC_OpGroupNonUniformFMax       : I32EnumAttrCase<"OpGroupNonUniformFMax", 358>;
 def SPV_OC_OpSubgroupBallotKHR         : I32EnumAttrCase<"OpSubgroupBallotKHR", 4421>;
 def SPV_OC_OpTypeCooperativeMatrixNV   : I32EnumAttrCase<"OpTypeCooperativeMatrixNV", 5358>;
 def SPV_OC_OpCooperativeMatrixLoadNV   : I32EnumAttrCase<"OpCooperativeMatrixLoadNV", 5359>;
+def SPV_OC_OpCooperativeMatrixStoreNV  : I32EnumAttrCase<"OpCooperativeMatrixStoreNV", 5360>;
+def SPV_OC_OpCooperativeMatrixMulAddNV : I32EnumAttrCase<"OpCooperativeMatrixMulAddNV", 5361>;
+def SPV_OC_OpCooperativeMatrixLengthNV : I32EnumAttrCase<"OpCooperativeMatrixLengthNV", 5362>;
 
 def SPV_OpcodeAttr :
     SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpNop, SPV_OC_OpUndef, SPV_OC_OpSourceContinued, SPV_OC_OpSource,
       SPV_OC_OpSourceExtension, SPV_OC_OpName, SPV_OC_OpMemberName, SPV_OC_OpString,
       SPV_OC_OpLine, SPV_OC_OpExtension, SPV_OC_OpExtInstImport, SPV_OC_OpExtInst,
       SPV_OC_OpMemoryModel, SPV_OC_OpEntryPoint, SPV_OC_OpExecutionMode,
       SPV_OC_OpCapability, SPV_OC_OpTypeVoid, SPV_OC_OpTypeBool, SPV_OC_OpTypeInt,
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpReturnValue, SPV_OC_OpUnreachable, SPV_OC_OpNoLine,
       SPV_OC_OpModuleProcessed, SPV_OC_OpGroupNonUniformElect,
       SPV_OC_OpGroupNonUniformBallot, SPV_OC_OpGroupNonUniformIAdd,
       SPV_OC_OpGroupNonUniformFAdd, SPV_OC_OpGroupNonUniformIMul,
       SPV_OC_OpGroupNonUniformFMul, SPV_OC_OpGroupNonUniformSMin,
       SPV_OC_OpGroupNonUniformUMin, SPV_OC_OpGroupNonUniformFMin,
       SPV_OC_OpGroupNonUniformSMax, SPV_OC_OpGroupNonUniformUMax,
       SPV_OC_OpGroupNonUniformFMax, SPV_OC_OpSubgroupBallotKHR,
-      SPV_OC_OpTypeCooperativeMatrixNV, SPV_OC_OpCooperativeMatrixLoadNV
+      SPV_OC_OpTypeCooperativeMatrixNV, SPV_OC_OpCooperativeMatrixLoadNV,
+      SPV_OC_OpCooperativeMatrixStoreNV, SPV_OC_OpCooperativeMatrixMulAddNV,
+      SPV_OC_OpCooperativeMatrixLengthNV
     ]>;
 
 // End opcode section. Generated from SPIR-V spec; DO NOT MODIFY!
 
 //===----------------------------------------------------------------------===//
 // SPIR-V op definitions
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Dialect/SPIRV/SPIRVCastOps.td

 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class SPV_CastOp<string mnemonic, Type resultType, Type operandType,
                  list<OpTrait> traits = []> :
       SPV_Op<mnemonic,
              !listconcat(traits,
                          [NoSideEffect, SameOperandsAndResultShape])> {
   let arguments = (ins
-    SPV_ScalarOrVectorOf<operandType>:$operand
+    SPV_ScalarOrVectorOrCoopMatrixOf<operandType>:$operand
   );
 
   let results = (outs
-    SPV_ScalarOrVectorOf<resultType>:$result
+    SPV_ScalarOrVectorOrCoopMatrixOf<resultType>:$result
   );
 
   let parser = [{ return mlir::impl::parseCastOp(parser, result); }];
   let printer = [{ mlir::impl::printCastOp(this->getOperation(), p); }];
   let verifier = [{ return verifyCastOp(this->getOperation()); }];
 }
 
 // -----

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPIRV_COOPERATIVE_MATRIX_OPS
 #define SPIRV_COOPERATIVE_MATRIX_OPS
 
 // -----
 
+def SPV_CooperativeMatrixLengthNVOp : SPV_Op<"CooperativeMatrixLengthNV",
+  [NoSideEffect]> {
+  let summary = "See extension SPV_NV_cooperative_matrix";
+
+  let description = [{
+    Number of components of a cooperative matrix type accessible to each
+    invocation when treated as a composite.
+
+    Result Type must be an OpTypeInt with 32-bit Width and 0 Signedness.
+
+    Type is a cooperative matrix type.
+
+    ``` {.ebnf}
+    cooperative-matrix-length-op ::= ssa-id `=` `spv.CooperativeMatrixLengthNV
+                                    ` : ` cooperative-matrix-type
+    ```
+
+    For example:
+
+    ```
+    %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<Subgroup, i32, 8, 16>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[SPV_NV_cooperative_matrix]>,
+    Capability<[SPV_C_CooperativeMatrixNV]>
+  ];
+
+  let arguments = (ins
+    TypeAttr:$type
+  );
+
+  let results = (outs
+    SPV_Int32:$result
+  );
+  let verifier = [{ return success(); }];
+}
+
+// -----
+
 def SPV_CooperativeMatrixLoadNVOp : SPV_Op<"CooperativeMatrixLoadNV", []> {
   let summary = "See extension SPV_NV_cooperative_matrix";
 
   let description = [{
     Load a cooperative matrix through a pointer.
 
     Result Type is the type of the loaded object. It must be a cooperative
     matrix type.
let description = [{
     instruction must be the same for all invocations in a given scope instance
     (where the scope is the scope the cooperative matrix type was created with).
     All invocations in a given scope instance must be active or all must be
     inactive.
 
     ### Custom assembly form
 
     ``` {.ebnf}
-    cooperative-matrix-op ::= ssa-id `=` `spv.CooperativeMatrixLoadNV`
-                              storage-class ssa-use (`[` memory-access `]`)? `
-                              : ` cooperative-matrix-type
+    cooperative-matrixload-op ::= ssa-id `=` `spv.CooperativeMatrixLoadNV`
+                              storage-class ssa-use `,` ssa-use `,` ssa-use
+                              (`[` memory-access `]`)? ` : `
+                              cooperative-matrix-type
     ```
 
     For example:
 
     ```
     %0 = spv.CooperativeMatrixLoadNV "StorageBuffer" %ptr, %stride, %colMajor
          : !spv.coopmatrix<i32, Workgroup, 16, 8>
     ```
let arguments = (ins
     SPV_Bool:$columnmajor,
     OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
   );
 
   let results = (outs
     SPV_AnyCooperativeMatrix:$result
   );
 
   let verifier = [{ return success(); }];
 }
 
 // -----
 
+def SPV_CooperativeMatrixMulAddNVOp : SPV_Op<"CooperativeMatrixMulAddNV",
+  [NoSideEffect, SameOperandsAndResultType]> {
+  let summary = "See extension SPV_NV_cooperative_matrix";
+
+  let description = [{
+    Linear-algebraic matrix multiply of A by B and then component-wise add C.
+    The order of the operations is implementation-dependent. The internal
+    precision of floating-point operations is defined by the client API.
+    Integer operations are performed at the precision of the Result Type and are
+    exact unless there is overflow or underflow, in which case the result is
+    undefined.
+
+    Result Type must be a cooperative matrix type with M rows and N columns.
+
+    A is a cooperative matrix with M rows and K columns.
+
+    B is a cooperative matrix with K rows and N columns.
+
+    C is a cooperative matrix with M rows and N columns.
+
+    The values of M, N, and K must be consistent across the result and operands.
+    This is referred to as an MxNxK matrix multiply.
+
+    A, B, C, and Result Type must have the same scope, and this defines the
+    scope of the operation. A, B, C, and Result Type need not necessarily have
+    the same component type, this is defined by the client API.
+
+    If the Component Type of any matrix operand is an integer type, then its
+    components are treated as signed if its Component Type has Signedness of 1
+    and are treated as unsigned otherwise.
+
+    For a given dynamic instance of this instruction, all invocations in a given
+    scope instance must be active or all must be inactive (where the scope is
+    the scope of the operation).
+
+    ``` {.ebnf}
+    cooperative-matrixmuladd-op ::= ssa-id `=` `spv.CooperativeMatrixMulAddNV`
+                              sssa-use `,` ssa-use `,` ssa-use ` : `
+                              cooperative-matrix-type
+    ```
+    For example:
+
+    ```
+    %0 = spv.CooperativeMatrixMulAddNV %arg0, %arg1, %arg2,  :
+      !spv.coopmatrix<Subgroup, i32, 8, 16>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[SPV_NV_cooperative_matrix]>,
+    Capability<[SPV_C_CooperativeMatrixNV]>
+  ];
+
+  let arguments = (ins
+    SPV_AnyCooperativeMatrix:$a,
+    SPV_AnyCooperativeMatrix:$b,
+    SPV_AnyCooperativeMatrix:$c
+  );
+
+  let results = (outs
+    SPV_AnyCooperativeMatrix:$result
+  );
+
+  let verifier = [{ return success(); }];
+}
+
+// -----
+
+def SPV_CooperativeMatrixStoreNVOp : SPV_Op<"CooperativeMatrixStoreNV", []> {
+  let summary = "See extension SPV_NV_cooperative_matrix";
+
+  let description = [{
+    Store a cooperative matrix through a pointer.
+
+    Pointer is a pointer into an array. Its type must be an OpTypePointer whose
+    Type operand is a scalar or vector type. The storage class of Pointer must
+    be Workgroup, StorageBuffer, or (if SPV_EXT_physical_storage_buffer is
+    supported) PhysicalStorageBufferEXT.
+
+    Object is the object to store. Its type must be an
+    OpTypeCooperativeMatrixNV.
+
+    Stride is the number of elements in the array in memory between the first
+    component of consecutive rows (or columns) in the result. It must be a
+    scalar integer type.
+
+    ColumnMajor indicates whether the values stored to memory are arranged in
+    column-major or row-major order. It must be a boolean constant instruction,
+    with false indicating row major and true indicating column major.
+
+    Memory Access must be a Memory Access literal. If not present, it is the
+    same as specifying None.
+
+    ``` {.ebnf}
+    coop-matrix-store-op ::= `spv.CooperativeMatrixStoreNV ` 
+                              storage-class ssa-use `, ` ssa-use `, `
+                              ssa-use `, ` ssa-use `, `
+                  (`[` memory-access `]`)? `:` spirv-element-type
+    ```
+
+    For example:
+
+    ```
+      spv.CooperativeMatrixStoreNV "StorageBuffer" %arg0, %arg2, %arg1, %arg3 :
+        !spv.coopmatrix<Workgroup, i32, 16, 8>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[SPV_NV_cooperative_matrix]>,
+    Capability<[SPV_C_CooperativeMatrixNV]>
+  ];
+
+  let arguments = (ins
+    SPV_AnyPtr:$pointer,
+    SPV_AnyCooperativeMatrix:$object,
+    SPV_Integer:$stride,
+    SPV_Bool:$columnmajor,
+    OptionalAttr<SPV_MemoryAccessAttr>:$memory_access
+  );
+
+  let results = (outs);
+
+  let verifier = [{ return success(); }];
+}
+
+// -----
+
 #endif // SPIRV_COOPERATIVE_MATRIX_OPS

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static LogicalResult verifyCastOp(Operation *op,
   Type resultType = op->getResult(0).getType();
 
   // ODS checks that result type and operand type have the same shape.
   if (auto vectorType = operandType.dyn_cast<VectorType>()) {
     operandType = vectorType.getElementType();
     resultType = resultType.cast<VectorType>().getElementType();
   }
 
+  if (auto coopMatrixType =
+          operandType.dyn_cast<spirv::CooperativeMatrixNVType>()) {
+    operandType = coopMatrixType.getElementType();
+    resultType =
+        resultType.cast<spirv::CooperativeMatrixNVType>().getElementType();
+  }
+
   auto operandTypeBitWidth = operandType.getIntOrFloatBitWidth();
   auto resultTypeBitWidth = resultType.getIntOrFloatBitWidth();
   auto isSameBitWidth = operandTypeBitWidth == resultTypeBitWidth;
 
   if (requireSameBitWidth) {
     if (!isSameBitWidth) {
       return op->emitOpError(
                  "expected the same bit widths for operand type and result "
printer << spirv::CooperativeMatrixLoadNVOp::getOperationName() << " \"" << sc
           << "\" " << M.pointer() << ", " << M.stride() << ", "
           << M.columnmajor();
   // Print optional memory access attribute.
   if (auto memAccess = M.memory_access())
     printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"]";
   printer << " : " << M.getType();
 }
 
+//===----------------------------------------------------------------------===//
+// spv.CooperativeMatrixStoreNV
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseCooperativeMatrixStoreNVOp(OpAsmParser &parser,
+                                                   OperationState &state) {
+  spirv::StorageClass storageClass;
+  SmallVector<OpAsmParser::OperandType, 4> operandInfo;
+  Type strideType = parser.getBuilder().getIntegerType(32);
+  Type columnMajorType = parser.getBuilder().getIntegerType(1);
+  Type elementType;
+  if (parseEnumStrAttr(storageClass, parser) ||
+      parser.parseOperandList(operandInfo, 4) ||
+      parseMemoryAccessAttributes(parser, state) || parser.parseColon() ||
+      parser.parseType(elementType)) {
+    return failure();
+  }
+
+  auto ptrType = spirv::PointerType::get(
+      elementType.cast<spirv::CooperativeMatrixNVType>().getElementType(),
+      storageClass);
+  SmallVector<Type, 4> OperandType = {ptrType, elementType, strideType,
+                                      columnMajorType};
+  if (parser.resolveOperands(operandInfo, OperandType, parser.getNameLoc(),
+                             state.operands)) {
+    return failure();
+  }
+
+  return success();
+}
+
+static void print(spirv::CooperativeMatrixStoreNVOp M, OpAsmPrinter &printer) {
+  StringRef sc = stringifyStorageClass(
+      M.pointer().getType().cast<spirv::PointerType>().getStorageClass());
+  printer << spirv::CooperativeMatrixStoreNVOp::getOperationName() << " \""
+          << sc << "\" " << M.pointer() << ", " << M.object() << ", "
+          << M.stride() << ", " << M.columnmajor();
+  // Print optional memory access attribute.
+  if (auto memAccess = M.memory_access())
+    printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"]";
+  printer << " : " << M.getOperand(1).getType();
+}
+
+//===----------------------------------------------------------------------===//
+// spv.CooperativeMatrixLengthNV
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseCooperativeMatrixLengthNVOp(OpAsmParser &parser,
+                                                    OperationState &state) {
+  OpAsmParser::OperandType operandInfo;
+  Type dstType = parser.getBuilder().getIntegerType(32);
+  Type type;
+  if (parser.parseColonType(type)) {
+    return failure();
+  }
+  state.addAttribute(kTypeAttrName, TypeAttr::get(type));
+  state.addTypes(dstType);
+  return success();
+}
+
+static void print(spirv::CooperativeMatrixLengthNVOp M, OpAsmPrinter &printer) {
+  printer << M.getOperationName() << " : " << M.type();
+}
+
+//===----------------------------------------------------------------------===//
+// spv.CooperativeMatrixMulAddNV
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseCooperativeMatrixMulAddNVOp(OpAsmParser &parser,
+                                                    OperationState &state) {
+  SmallVector<OpAsmParser::OperandType, 3> ops;
+  Type type;
+  if (parser.parseOperandList(ops, 3) || parser.parseColonType(type) ||
+      parser.resolveOperands(ops, type, state.operands)) {
+    return failure();
+  }
+  state.addTypes(type);
+  return success();
+}
+
+static void print(spirv::CooperativeMatrixMulAddNVOp M, OpAsmPrinter &printer) {
+  printer << M.getOperationName() << ' ' << M.getOperand(0) << ", "
+          << M.getOperand(1) << ", " << M.getOperand(2) << ", "
+          << " : " << M.getOperand(0).getType();
+}
+
 namespace mlir {
 namespace spirv {
 
 // TableGen'erated operation interfaces for querying versions, extensions, and
 // capabilities.
 #include "mlir/Dialect/SPIRV/SPIRVAvailability.cpp.inc"
 
 // TablenGen'erated operation definitions.
 #define GET_OP_CLASSES
 #include "mlir/Dialect/SPIRV/SPIRVOps.cpp.inc"
 
 // TableGen'erated operation availability interface implementations.
 #include "mlir/Dialect/SPIRV/SPIRVOpAvailabilityImpl.inc"
 
 } // namespace spirv
 } // namespace mlir

mlir/test/Dialect/SPIRV/Serialization/cooperative-matrix.mlir

 // RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [CooperativeMatrixNV], [SPV_NV_cooperative_matrix]> {
   // CHECK-LABEL: @cooperative_matrix_load
   spv.func @cooperative_matrix_load(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
     // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} : !spv.coopmatrix<16x8xi32, Workgroup>
     %0 = spv.CooperativeMatrixLoadNV "StorageBuffer" %ptr, %stride, %b : !spv.coopmatrix<16x8xi32, Workgroup>
     spv.Return
   }
 
   // CHECK-LABEL: @cooperative_matrix_load_memaccess
   spv.func @cooperative_matrix_load_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
     // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
     %0 = spv.CooperativeMatrixLoadNV "StorageBuffer" %ptr, %stride, %b ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
     spv.Return
   }
+
+  // CHECK-LABEL: @cooperative_matrix_store
+  spv.func @cooperative_matrix_store(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %m : !spv.coopmatrix<16x8xi32, Workgroup>, %b : i1) "None" {
+    // CHECK: spv.CooperativeMatrixStoreNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} : !spv.coopmatrix<16x8xi32, Workgroup>
+    spv.CooperativeMatrixStoreNV "StorageBuffer" %ptr, %m, %stride, %b : !spv.coopmatrix<16x8xi32, Workgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_store_memaccess
+  spv.func @cooperative_matrix_store_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %m : !spv.coopmatrix<8x16xi32, Subgroup>, %stride : i32, %b : i1) "None" {
+    // CHECK: spv.CooperativeMatrixStoreNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.CooperativeMatrixStoreNV "StorageBuffer" %ptr, %m, %stride, %b ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_length
+  spv.func @cooperative_matrix_length() -> i32 "None" {
+    // CHECK: {{%.*}} = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
+    %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.ReturnValue %0 : i32
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_muladd
+  spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>, %c : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.CooperativeMatrixMulAddNV {{%.*}}, {{%.*}}, {{%.*}},  : !spv.coopmatrix<8x16xi32, Subgroup>
+    %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_add
+  spv.func @cooperative_matrix_add(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.IAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+    %r = spv.IAdd %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_sub
+  spv.func @cooperative_matrix_sub(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.ISub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+    %r = spv.ISub %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_sdiv
+  spv.func @cooperative_matrix_sdiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.SDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+    %r = spv.SDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_udiv
+  spv.func @cooperative_matrix_udiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.UDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+    %r = spv.UDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_fadd
+  spv.func @cooperative_matrix_fadd(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.FAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+    %r = spv.FAdd %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_fsub
+  spv.func @cooperative_matrix_fsub(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.FSub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+    %r = spv.FSub %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
+    spv.Return
+  }
+
+  // CHECK-LABEL: @cooperative_matrix_fdiv
+  spv.func @cooperative_matrix_fdiv(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
+    // CHECK: {{%.*}} = spv.FDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+    %r = spv.FDiv %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
+    spv.Return
+  }
 }

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

 // RUN: mlir-opt -allow-unregistered-dialect -split-input-file -verify-diagnostics %s | FileCheck %s
 
 // CHECK-LABEL: @cooperative_matrix_load
 spv.func @cooperative_matrix_load(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
   // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} : !spv.coopmatrix<16x8xi32, Workgroup>
   %0 = spv.CooperativeMatrixLoadNV "StorageBuffer" %ptr, %stride, %b : !spv.coopmatrix<16x8xi32, Workgroup>
   spv.Return
 }
 
 // -----
 // CHECK-LABEL: @cooperative_matrix_load_memaccess
 spv.func @cooperative_matrix_load_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
   // CHECK: {{%.*}} = spv.CooperativeMatrixLoadNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
   %0 = spv.CooperativeMatrixLoadNV "StorageBuffer" %ptr, %stride, %b ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
   spv.Return
 }
+
+// CHECK-LABEL: @cooperative_matrix_store
+spv.func @cooperative_matrix_store(%ptr : !spv.ptr<i32, StorageBuffer>, %stride : i32, %m : !spv.coopmatrix<8x16xi32, Workgroup>, %b : i1) "None" {
+  // CHECK: spv.CooperativeMatrixStoreNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Workgroup>
+  spv.CooperativeMatrixStoreNV "StorageBuffer" %ptr, %m, %stride, %b : !spv.coopmatrix<8x16xi32, Workgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_store_memaccess
+spv.func @cooperative_matrix_store_memaccess(%ptr : !spv.ptr<i32, StorageBuffer>, %m : !spv.coopmatrix<8x16xi32, Subgroup>, %stride : i32, %b : i1) "None" {
+  // CHECK: spv.CooperativeMatrixStoreNV "StorageBuffer" {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.CooperativeMatrixStoreNV "StorageBuffer" %ptr, %m, %stride, %b ["Volatile"] : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_length
+spv.func @cooperative_matrix_length() -> i32 "None" {
+  // CHECK: {{%.*}} = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
+  %0 = spv.CooperativeMatrixLengthNV : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.ReturnValue %0 : i32
+}
+
+// CHECK-LABEL: @cooperative_matrix_muladd
+spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>, %c : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.CooperativeMatrixMulAddNV {{%.*}}, {{%.*}}, {{%.*}},  : !spv.coopmatrix<8x16xi32, Subgroup>
+  %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_add
+spv.func @cooperative_matrix_add(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.IAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+  %r = spv.IAdd %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_sub
+spv.func @cooperative_matrix_sub(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.ISub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+  %r = spv.ISub %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_sdiv
+spv.func @cooperative_matrix_sdiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.SDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+  %r = spv.SDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_udiv
+spv.func @cooperative_matrix_udiv(%a : !spv.coopmatrix<8x16xi32, Subgroup>, %b : !spv.coopmatrix<8x16xi32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.UDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xi32, Subgroup>
+  %r = spv.UDiv %a, %b : !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_fadd
+spv.func @cooperative_matrix_fadd(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.FAdd {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+  %r = spv.FAdd %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_fsub
+spv.func @cooperative_matrix_fsub(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.FSub {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+  %r = spv.FSub %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
+  spv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_fdiv
+spv.func @cooperative_matrix_fdiv(%a : !spv.coopmatrix<8x16xf32, Subgroup>, %b : !spv.coopmatrix<8x16xf32, Subgroup>) "None" {
+  // CHECK: {{%.*}} = spv.FDiv {{%.*}}, {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+  %r = spv.FDiv %a, %b : !spv.coopmatrix<8x16xf32, Subgroup>
+  spv.Return
+}

mlir/test/Dialect/SPIRV/ops.mlir

 func @convert_f_to_u_vector(%arg0 : vector<3xf32>) -> vector<3xi32> {
   // CHECK: {{%.*}} = spv.ConvertFToU {{%.*}} : vector<3xf32> to vector<3xi32>
   %0 = spv.ConvertFToU %arg0 : vector<3xf32> to vector<3xi32>
   spv.ReturnValue %0 : vector<3xi32>
 }
 
 // -----
 
+func @convert_f_to_u_coopmatrix(%arg0 : !spv.coopmatrix<8x16xf32, Subgroup>) {
+  // CHECK: {{%.*}} = spv.ConvertFToU {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup> to !spv.coopmatrix<8x16xi32, Subgroup>
+  %0 = spv.ConvertFToU %arg0 : !spv.coopmatrix<8x16xf32, Subgroup> to !spv.coopmatrix<8x16xi32, Subgroup>
+  spv.Return
+}
+
+// -----
+
 func @convert_f_to_u_scalar_invalid(%arg0 : f16) -> i32 {
   // expected-error @+1 {{expected the same bit widths for operand type and result type, but provided 'f16' and 'i32'}}
   %0 = spv.ConvertFToU %arg0 : f16 to i32
   spv.ReturnValue %0 : i32
 }
 
 // -----
 

 func @f_convert_vector(%arg0 : vector<3xf32>) -> vector<3xf64> {
   // CHECK: {{%.*}} = spv.FConvert {{%.*}} : vector<3xf32> to vector<3xf64>
   %0 = spv.FConvert %arg0 : vector<3xf32> to vector<3xf64>
   spv.ReturnValue %0 : vector<3xf64>
 }
 
 // -----
 
+func @f_convert_coop_matrix(%arg0 : !spv.coopmatrix<8x16xf32, Subgroup>) {
+  // CHECK: {{%.*}} = spv.FConvert {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup> to !spv.coopmatrix<8x16xf64, Subgroup>
+  %0 = spv.FConvert %arg0 : !spv.coopmatrix<8x16xf32, Subgroup> to !spv.coopmatrix<8x16xf64, Subgroup>
+  spv.Return
+}
+
+// -----
+
 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 = spv.FConvert %arg0 : f32 to f32
   spv.ReturnValue %0 : f32
 }
 
 // -----
 

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

os << tabs
                   "attrElem.cast<IntegerAttr>().getValue().getZExtValue()));\n",
                   operandList);
     os << tabs << "  }\n";
   } else if (attr.isEnumAttr() || attr.getAttrDefName() == "I32Attr") {
     os << tabs
        << formatv("  {0}.push_back(static_cast<uint32_t>("
                   "attr.cast<IntegerAttr>().getValue().getZExtValue()));\n",
                   operandList);
+  } else if (attr.isEnumAttr() || attr.getAttrDefName() == "TypeAttr") {
+    os << tabs
+       << formatv("  {0}.push_back(static_cast<uint32_t>("
+                  "getTypeID(attr.cast<TypeAttr>().getValue())));\n",
+                  operandList);
   } else {
     PrintFatalError(
         loc,
         llvm::Twine(
             "unhandled attribute type in SPIR-V serialization generation : '") +
             attr.getAttrDefName() + llvm::Twine("'"));
   }
   os << tabs << "}\n";
os << tabs
        << formatv("{0}.push_back(opBuilder.getNamedAttr(\"{1}\", "
                   "opBuilder.getArrayAttr(attrListElems)));\n",
                   attrList, attrName);
   } else if (attr.isEnumAttr() || 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") {
+    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("'"));
   }
 }