[mlir][spirv] Add cooperative matrix store op

[kuhar]

Summary

Implement cooperative matrix store for the SPV_KHR_cooperative_matrix
extension: https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_cooperative_matrix.html.

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

 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_OpCooperativeMatrixLoadKHR   : I32EnumAttrCase<"OpCooperativeMatrixLoadKHR", 4457>;
+def SPIRV_OC_OpCooperativeMatrixStoreKHR  : I32EnumAttrCase<"OpCooperativeMatrixStoreKHR", 4458>;
 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>;
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_OpTypeCooperativeMatrixKHR,
-      SPIRV_OC_OpCooperativeMatrixLoadKHR, SPIRV_OC_OpCooperativeMatrixLengthKHR,
-      SPIRV_OC_OpTypeCooperativeMatrixNV,
-      SPIRV_OC_OpCooperativeMatrixLoadNV, SPIRV_OC_OpCooperativeMatrixStoreNV,
-      SPIRV_OC_OpCooperativeMatrixMulAddNV, SPIRV_OC_OpCooperativeMatrixLengthNV,
+      SPIRV_OC_OpSUDotAccSat,
+      SPIRV_OC_OpTypeCooperativeMatrixKHR, SPIRV_OC_OpCooperativeMatrixLoadKHR,
+      SPIRV_OC_OpCooperativeMatrixStoreKHR, 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,
       SPIRV_OC_OpTypejointMatrixWorkItemLengthINTEL,
 

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

let arguments = (ins
     OptionalAttr<SPIRV_MemoryAccessAttr>:$memory_operand
   );
 
   let results = (outs
     SPIRV_AnyCooperativeMatrix:$result
   );
 }
 
+// -----
+
+def SPIRV_KHRCooperativeMatrixStoreOp : SPIRV_KhrVendorOp<"CooperativeMatrixStore", []> {
+  let summary = "Stores a cooperative matrix through a pointer";
+
+  let description = [{
+    Store a cooperative matrix through a pointer.
+    Pointer is a pointer. Its type must be an OpTypePointer whose Type operand
+    is a scalar or vector type. If the Shader capability was declared, Pointer
+    must point into an array and any ArrayStride decoration on Pointer is
+    ignored.
+
+    Object is the object to store. Its type must be an
+    OpTypeCooperativeMatrixKHR.
+
+    MemoryLayout specifies how matrix elements are laid out in memory. It must
+    come from a 32-bit integer constant instruction whose value corresponds to a
+    Cooperative Matrix Layout. See the Cooperative Matrix Layout table for a
+    description of the layouts and detailed layout-specific rules.
+
+    Stride further qualifies how matrix elements are laid out in memory. It must
+    be a scalar integer type and its exact semantics depend on MemoryLayout.
+
+    Memory Operand must be a Memory Operand literal. If not present, it is the
+    same as specifying None.
+
+    NOTE: In earlier versions of the SPIR-V spec, 'Memory Operand' was known
+    as 'Memory Access'.
+
+    For a given dynamic instance of this instruction, all operands of this
+    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.
+
+    ``` {.ebnf}
+     coop-matrix-store-op ::= `spirv.KHR.CooperativeMatrixStore `
+                              ssa-use `, ` ssa-use `, `
+                              ssa-use `, ` cooperative-matrix-layout `, `
+                              (`[` memory-operand `]`)? `:`
+                              pointer-type `,` coop-matrix-type
+    ```
+
+    #### Example:
+
+    ```
+      spirv.KHR.CooperativeMatrixStore %ptr, %obj, %stride :
+        !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_6>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[SPV_KHR_cooperative_matrix]>,
+    Capability<[SPIRV_C_CooperativeMatrixKHR]>
+  ];
+
+  let arguments = (ins
+    SPIRV_AnyPtr:$pointer,
+    SPIRV_AnyCooperativeMatrix:$object,
+    SPIRV_Integer:$stride,
+    SPIRV_KHR_CooperativeMatrixLayoutAttr:$matrix_layout,
+    OptionalAttr<SPIRV_MemoryAccessAttr>:$memory_operand
+  );
+
+  let results = (outs);
+}
+
 //===----------------------------------------------------------------------===//
 // SPV_NV_cooperative_matrix extension ops.
 //===----------------------------------------------------------------------===//
 
 // -----
 
 def SPIRV_NVCooperativeMatrixLengthOp : SPIRV_NvVendorOp<"CooperativeMatrixLength",
   [Pure]> {
let description = [{
     Memory Access must be a Memory Access literal. If not present, it is the
     same as specifying None.
 
     ``` {.ebnf}
     coop-matrix-store-op ::= `spirv.NV.CooperativeMatrixStore `
                               ssa-use `, ` ssa-use `, `
                               ssa-use `, ` ssa-use `, `
                               (`[` memory-access `]`)? `:`
-                              pointer-type `,` spirv-element-type
+                              pointer-type `,` coop-matrix-type
     ```
 
     For example:
 
     ```
       spirv.NV.CooperativeMatrixStore %arg0, %arg2, %arg1, %arg3 :
         !spirv.ptr<i32, StorageBuffer>, !spirv.NV.coopmatrix<16x8xi32, Workgroup>
     ```

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

 }
 
 LogicalResult spirv::KHRCooperativeMatrixLoadOp::verify() {
   return verifyPointerAndCoopMatrixType(*this, getPointer().getType(),
                                         getResult().getType());
 }
 
 //===----------------------------------------------------------------------===//
+// spirv.KHR.CooperativeMatrixStore
+//===----------------------------------------------------------------------===//
+
+ParseResult spirv::KHRCooperativeMatrixStoreOp::parse(OpAsmParser &parser,
+                                                      OperationState &result) {
+  std::array<OpAsmParser::UnresolvedOperand, 3> operandInfo = {};
+  for (auto &op : operandInfo) {
+    if (parser.parseOperand(op) || parser.parseComma())
+      return failure();
+  }
+
+  spirv::CooperativeMatrixLayoutKHR layout;
+  if (::parseEnumKeywordAttr<spirv::CooperativeMatrixLayoutKHRAttr>(
+          layout, parser, result, kKhrCooperativeMatrixLayoutAttrName)) {
+    return failure();
+  }
+
+  if (parseMemoryAccessAttributes(parser, result, kMemoryOperandAttrName))
+    return failure();
+
+  Type ptrType;
+  Type objectType;
+  if (parser.parseColon() || parser.parseType(ptrType) || parser.parseComma() ||
+      parser.parseType(objectType)) {
+    return failure();
+  }
+
+  Type strideType = parser.getBuilder().getIntegerType(32);
+  if (parser.resolveOperands(operandInfo, {ptrType, objectType, strideType},
+                             parser.getNameLoc(), result.operands)) {
+    return failure();
+  }
+
+  return success();
+}
+
+void spirv::KHRCooperativeMatrixStoreOp::print(OpAsmPrinter &printer) {
+  printer << " " << getPointer() << ", " << getObject() << ", " << getStride()
+          << ", " << getMatrixLayout();
+
+  // Print optional memory operand attribute.
+  if (auto memOperand = getMemoryOperand())
+    printer << " [\"" << *memOperand << "\"]";
+  printer << " : " << getPointer().getType() << ", " << getObject().getType();
+}
+
+LogicalResult spirv::KHRCooperativeMatrixStoreOp::verify() {
+  return verifyPointerAndCoopMatrixType(*this, getPointer().getType(),
+                                        getObject().getType());
+}
+
+//===----------------------------------------------------------------------===//
 // 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";

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

 spirv.func @cooperative_matrix_load_function(%ptr : !spirv.ptr<i32, Function>, %stride : i32) "None" {
   // CHECK:      {{%.*}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, RowMajor :
   // CHECK-SAME:   !spirv.ptr<i32, Function> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixAcc>
   %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, RowMajor :
     !spirv.ptr<i32, Function> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixAcc>
   spirv.Return
 }
 
+// CHECK-LABEL: @cooperative_matrix_store
+spirv.func @cooperative_matrix_store(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32,
+                                     %m : !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>) "None" {
+  // CHECK:      spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, {{%.*}}, RowMajor :
+  // CHECK-SAME:   !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>
+  spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, RowMajor :
+    !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>
+  spirv.Return
+}
+
+// CHECK-LABEL: @cooperative_matrix_store_memoperand
+spirv.func @cooperative_matrix_store_memoperand(%ptr : !spirv.ptr<i32, StorageBuffer>,
+                                                %m : !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>,
+                                                %stride : i32) "None" {
+  // CHECK:       spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, {{%.*}}, ColumnMajor ["Volatile"] :
+  // CHECK-SAME:    !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
+  spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, ColumnMajor ["Volatile"] :
+    !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Subgroup, MatrixB>
+  spirv.Return
+}
+
 // -----
 
 spirv.func @cooperative_matrix_load_bad_ptr(%ptr : !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer>, %stride : i32) "None" {
   // expected-error @+1 {{Pointer must point to a scalar or vector type}}
   %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, ColumnMajor :
     !spirv.ptr<!spirv.struct<(f32 [0])>, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
   spirv.Return
 }
spirv.func @cooperative_matrix_load_bad_result(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {
   // expected-error @+1 {{op result #0 must be any SPIR-V cooperative matrix type}}
   %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, ColumnMajor :
     !spirv.ptr<i32, StorageBuffer> as !spirv.NV.coopmatrix<8x16xi32, Subgroup>
   spirv.Return
 }
 
 // -----
 
+spirv.func @cooperative_matrix_store_missing_attr(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32,
+                                                  %m : !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>) "None" {
+  // expected-error @+1 {{expected ','}}
+  spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride :
+    !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>
+  spirv.Return
+}
+
+// -----
+
+spirv.func @cooperative_matrix_store_missing_attr(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32,
+                                                  %m : !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>) "None" {
+  // expected-error @+1 {{expected valid keyword}}
+  spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, :
+    !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Workgroup, MatrixA>
+  spirv.Return
+}
+
+// -----
+
+spirv.func @cooperative_matrix_store_bad_object_type(%ptr : !spirv.ptr<i32, StorageBuffer>,
+                                                     %stride : i32) "None" {
+  // expected-error @+1 {{op operand #1 must be any SPIR-V cooperative matrix type}}
+  spirv.KHR.CooperativeMatrixStore %ptr, %stride, %stride, RowMajor :
+    !spirv.ptr<i32, StorageBuffer>, i32
+  spirv.Return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // 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>