[mlir][spirv] Enable composite instructions for cooperative matrix type.

[ThomasRaoux]

Summary

Enable inset/extract/construct composite ops as well as access chain for cooperative matrix. ConstantComposite requires more change and will be done in a separate patch.
Also fix the getNumElements function for coopMatrix per feedback from Jeff Bolz. The number of element is implementation dependent so it cannot be known at compile time.

[antiagainst]

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

143

Could you add a comment to getNumElements to mention that this method being false is a pre-condition? (With that actually it would be nice to rename this as hasCompileTimeKnownNumElements or something so that we don't always use it as the negate form. But a minor issue.)

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

100

Can we return %1 here so we can check the return type too? Right now it's hidden.

mlir/test/Dialect/SPIRV/composite-ops.mlir

25

I think this should only take one operand.

[ThomasRaoux]

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

143

Added a comment about getNumElements and change the function name.

mlir/test/Dialect/SPIRV/composite-ops.mlir

25

Thanks for catching that, changed it and added a negative test.

[antiagainst]

Nice, thanks! Just one more nit.

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

96

Note that according to the spec: "Cooperative matrix types (or types containing them) can only be allocated in Function or Private storage classes." So here we cannot have a pointer pointing to a coop matrix in a StorageBuffer. We should use Function/Private here. Although this is test; still would be nice to be valid. So later when we enforce this rule these tests can still pass as-is. :)

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

98

Same here.

[ThomasRaoux]

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

96

Good point, I changed it.

[rriddle]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

1138

nit: Merge this if into the else.

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

def SPV_AnyRTArray : DialectType<SPIRV_Dialect, SPV_IsRTArrayType,
                                  "any SPIR-V runtime array type">;
 def SPV_AnyStruct : DialectType<SPIRV_Dialect, SPV_IsStructType,
                                 "any SPIR-V struct type">;
 
 def SPV_Numerical : AnyTypeOf<[SPV_Integer, SPV_Float]>;
 def SPV_Scalar : AnyTypeOf<[SPV_Numerical, SPV_Bool]>;
 def SPV_Aggregate : AnyTypeOf<[SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct]>;
 def SPV_Composite :
-    AnyTypeOf<[SPV_Vector, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct]>;
+    AnyTypeOf<[SPV_Vector, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct,
+               SPV_AnyCooperativeMatrix]>;
 def SPV_Type : AnyTypeOf<[
     SPV_Void, SPV_Bool, SPV_Integer, SPV_Float, SPV_Vector,
-    SPV_AnyPtr, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct
+    SPV_AnyPtr, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct,
+    SPV_AnyCooperativeMatrix
   ]>;
 
 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">;

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

 public:
   using SPIRVType::SPIRVType;
 
   static bool classof(Type type);
 
   /// Returns true if the given vector type is valid for the SPIR-V dialect.
   static bool isValid(VectorType);
 
+  /// Return the number of elements of the type. This should only be called if
+  /// hasCompileTimeKnownNumElements is true.
   unsigned getNumElements() const;
 
   Type getElementType(unsigned) const;
 
+  /// Return true if the number of elements is known at compile time and is not
+  /// implementation dependent.
+  bool hasCompileTimeKnownNumElements() const;
+
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
 };
 
 // SPIR-V array type
 class ArrayType : public Type::TypeBase<ArrayType, CompositeType,
public:
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
 };
 
 // SPIR-V cooperative matrix type
 class CooperativeMatrixNVType
-    : public Type::TypeBase<CooperativeMatrixNVType, SPIRVType,
+    : public Type::TypeBase<CooperativeMatrixNVType, CompositeType,
                             detail::CooperativeMatrixTypeStorage> {
 public:
   using Base::Base;
 
   static bool kindof(unsigned kind) {
     return kind == TypeKind::CooperativeMatrix;
   }
 

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

getElementType(Type type, ArrayRef<int32_t> indices,
                function_ref<InFlightDiagnostic(StringRef)> emitErrorFn) {
   if (indices.empty()) {
     emitErrorFn("expected at least one index for spv.CompositeExtract");
     return nullptr;
   }
 
   for (auto index : indices) {
     if (auto cType = type.dyn_cast<spirv::CompositeType>()) {
-      if (index < 0 || static_cast<uint64_t>(index) >= cType.getNumElements()) {
+      if (cType.hasCompileTimeKnownNumElements() &&
+          (index < 0 ||
+           static_cast<uint64_t>(index) >= cType.getNumElements())) {
         emitErrorFn("index ") << index << " out of bounds for " << type;
         return nullptr;
       }
       type = cType.getElementType(index);
     } else {
       emitErrorFn("cannot extract from non-composite type ")
           << type << " with index " << index;
       return nullptr;
static ParseResult parseCompositeConstructOp(OpAsmParser &parser,
   }
   auto cType = type.dyn_cast<spirv::CompositeType>();
   if (!cType) {
     return parser.emitError(
                loc, "result type must be a composite type, but provided ")
            << type;
   }
 
-  if (operands.size() != cType.getNumElements()) {
+  if (cType.hasCompileTimeKnownNumElements() &&
+      operands.size() != cType.getNumElements()) {
     return parser.emitError(loc, "has incorrect number of operands: expected ")
            << cType.getNumElements() << ", but provided " << operands.size();
   }
   // TODO: Add support for constructing a vector type from the vector operands.
   // According to the spec: "for constructing a vector, the operands may
   // also be vectors with the same component type as the Result Type component
   // type".
   SmallVector<Type, 4> elementTypes;
-  elementTypes.reserve(cType.getNumElements());
-  for (auto index : llvm::seq<uint32_t>(0, cType.getNumElements())) {
+  elementTypes.reserve(operands.size());
+  for (auto index : llvm::seq<uint32_t>(0, operands.size())) {
     elementTypes.push_back(cType.getElementType(index));
   }
   state.addTypes(type);
   return parser.resolveOperands(operands, elementTypes, loc, state.operands);
 }
 
 static void print(spirv::CompositeConstructOp compositeConstructOp,
                   OpAsmPrinter &printer) {
   printer << spirv::CompositeConstructOp::getOperationName() << " "
           << compositeConstructOp.constituents() << " : "
           << compositeConstructOp.getResult().getType();
 }
 
 static LogicalResult verify(spirv::CompositeConstructOp compositeConstructOp) {
   auto cType = compositeConstructOp.getType().cast<spirv::CompositeType>();
-
   SmallVector<Value, 4> constituents(compositeConstructOp.constituents());
-  if (constituents.size() != cType.getNumElements()) {
+
+  if (cType.isa<spirv::CooperativeMatrixNVType>()) {
+    if (constituents.size() != 1)
+      return compositeConstructOp.emitError(
+                 "has incorrect number of operands: expected ")
+             << "1, but provided " << constituents.size();
+  } else {
+    if (constituents.size() != cType.getNumElements())
       return compositeConstructOp.emitError(
                  "has incorrect number of operands: expected ")
              << cType.getNumElements() << ", but provided "
              << constituents.size();
   }
 
   for (auto index : llvm::seq<uint32_t>(0, constituents.size())) {
     if (constituents[index].getType() != cType.getElementType(index)) {
       return compositeConstructOp.emitError(
                  "operand type mismatch: expected operand type ")
              << cType.getElementType(index) << ", but provided "
              << constituents[index].getType();

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

Type CompositeType::getElementType(unsigned index) const {
   }
 }
 
 unsigned CompositeType::getNumElements() const {
   switch (getKind()) {
   case spirv::TypeKind::Array:
     return cast<ArrayType>().getNumElements();
   case spirv::TypeKind::CooperativeMatrix:
-    return cast<CooperativeMatrixNVType>().getRows() *
-           cast<CooperativeMatrixNVType>().getColumns();
+    llvm_unreachable(
+        "invalid to query number of elements of spirv::CooperativeMatrix type");
   case spirv::TypeKind::RuntimeArray:
     llvm_unreachable(
         "invalid to query number of elements of spirv::RuntimeArray type");
   case spirv::TypeKind::Struct:
     return cast<StructType>().getNumElements();
   case StandardTypes::Vector:
     return cast<VectorType>().getNumElements();
   default:
     llvm_unreachable("invalid composite type");
   }
 }
 
+bool CompositeType::hasCompileTimeKnownNumElements() const {
+  switch (getKind()) {
+  case TypeKind::CooperativeMatrix:
+  case TypeKind::RuntimeArray:
+    return false;
+  default:
+    return true;
+  }
+}
+
 void CompositeType::getExtensions(
     SPIRVType::ExtensionArrayRefVector &extensions,
     Optional<StorageClass> storage) {
   switch (getKind()) {
   case spirv::TypeKind::Array:
     cast<ArrayType>().getExtensions(extensions, storage);
     break;
   case spirv::TypeKind::CooperativeMatrix:

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

spv.module Logical GLSL450 requires #spv.vce<v1.0, [CooperativeMatrixNV], [SPV_NV_cooperative_matrix]> {
   }
 
   // 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
   }
+
+  // CHECK-LABEL: @cooperative_matrix_access_chain
+  spv.func @cooperative_matrix_access_chain(%a : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>) -> !spv.ptr<f32, Function> "None" {
+    %0 = spv.constant 0: i32
+    // CHECK: {{%.*}} = spv.AccessChain {{%.*}}[{{%.*}}] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>
+    %1 = spv.AccessChain %a[%0] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>
+    spv.ReturnValue %1 : !spv.ptr<f32, Function>
+  }
 }

mlir/test/Dialect/SPIRV/composite-ops.mlir

 func @composite_construct_struct(%arg0: vector<3xf32>, %arg1: !spv.array<4xf32>, %arg2 : !spv.struct<f32>) -> !spv.struct<vector<3xf32>, !spv.array<4xf32>, !spv.struct<f32>> {
   // CHECK: spv.CompositeConstruct %arg0, %arg1, %arg2 : !spv.struct<vector<3xf32>, !spv.array<4 x f32>, !spv.struct<f32>>
   %0 = spv.CompositeConstruct %arg0, %arg1, %arg2 : !spv.struct<vector<3xf32>, !spv.array<4xf32>, !spv.struct<f32>>
   return %0: !spv.struct<vector<3xf32>, !spv.array<4xf32>, !spv.struct<f32>>
 }
 
 // -----
 
+func @composite_construct_coopmatrix(%arg0 : f32) -> !spv.coopmatrix<8x16xf32, Subgroup> {
+  // CHECK: spv.CompositeConstruct {{%.*}} : !spv.coopmatrix<8x16xf32, Subgroup>
+  %0 = spv.CompositeConstruct %arg0 : !spv.coopmatrix<8x16xf32, Subgroup>
+  return %0: !spv.coopmatrix<8x16xf32, Subgroup>
+}
+
+// -----
+
 func @composite_construct_empty_struct() -> !spv.struct<> {
   // CHECK: spv.CompositeConstruct : !spv.struct<>
   %0 = spv.CompositeConstruct : !spv.struct<>
   return %0: !spv.struct<>
 }
 
 // -----
 

 func @composite_construct_invalid_operand_type(%arg0: f32, %arg1: f32, %arg2 : f32) -> vector<3xi32> {
   // expected-error @+1 {{operand type mismatch: expected operand type 'i32', but provided 'f32'}}
   %0 = "spv.CompositeConstruct" (%arg0, %arg1, %arg2) : (f32, f32, f32) -> vector<3xi32>
   return %0: vector<3xi32>
 }
 
 // -----
 
+func @composite_construct_coopmatrix(%arg0 : f32, %arg1 : f32) -> !spv.coopmatrix<8x16xf32, Subgroup> {
+  // expected-error @+1 {{has incorrect number of operands: expected 1, but provided 2}}
+  %0 = spv.CompositeConstruct %arg0, %arg1 : !spv.coopmatrix<8x16xf32, Subgroup>
+  return %0: !spv.coopmatrix<8x16xf32, Subgroup>
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spv.CompositeExtractOp
 //===----------------------------------------------------------------------===//
 
 func @composite_extract_array(%arg0: !spv.array<4xf32>) -> f32 {
   // CHECK: {{%.*}} = spv.CompositeExtract {{%.*}}[1 : i32] : !spv.array<4 x f32>
   %0 = spv.CompositeExtract %arg0[1 : i32] : !spv.array<4xf32>
   return %0: f32

 func @composite_extract_vector(%arg0 : vector<4xf32>) -> f32 {
   // CHECK: {{%.*}} = spv.CompositeExtract {{%.*}}[1 : i32] : vector<4xf32>
   %0 = spv.CompositeExtract %arg0[1 : i32] : vector<4xf32>
   return %0 : f32
 }
 
 // -----
 
+func @composite_extract_coopmatrix(%arg0 : !spv.coopmatrix<8x16xf32, Subgroup>) -> f32 {
+  // CHECK: {{%.*}} = spv.CompositeExtract {{%.*}}[2 : i32] : !spv.coopmatrix<8x16xf32, Subgroup>
+  %0 = spv.CompositeExtract %arg0[2 : i32] : !spv.coopmatrix<8x16xf32, Subgroup>
+  return %0 : f32
+}
+
+// -----
+
 func @composite_extract_no_ssa_operand() -> () {
   // expected-error @+1 {{expected SSA operand}}
   %0 = spv.CompositeExtract [4 : i32, 1 : i32] : !spv.array<4x!spv.array<4xf32>>
   return
 }
 
 // -----
 

 func @composite_insert_struct(%arg0: !spv.struct<!spv.array<4xf32>, f32>, %arg1: !spv.array<4xf32>) -> !spv.struct<!spv.array<4xf32>, f32> {
   // CHECK: {{%.*}} = spv.CompositeInsert {{%.*}}, {{%.*}}[0 : i32] : !spv.array<4 x f32> into !spv.struct<!spv.array<4 x f32>, f32>
   %0 = spv.CompositeInsert %arg1, %arg0[0 : i32] : !spv.array<4xf32> into !spv.struct<!spv.array<4xf32>, f32>
   return %0: !spv.struct<!spv.array<4xf32>, f32>
 }
 
 // -----
 
+func @composite_insert_coopmatrix(%arg0: !spv.coopmatrix<8x16xi32, Subgroup>, %arg1: i32) -> !spv.coopmatrix<8x16xi32, Subgroup> {
+  // CHECK: {{%.*}} = spv.CompositeInsert {{%.*}}, {{%.*}}[5 : i32] : i32 into !spv.coopmatrix<8x16xi32, Subgroup>
+  %0 = spv.CompositeInsert %arg1, %arg0[5 : i32] : i32 into !spv.coopmatrix<8x16xi32, Subgroup>
+  return %0: !spv.coopmatrix<8x16xi32, Subgroup>
+}
+
+// -----
+
 func @composite_insert_no_indices(%arg0: !spv.array<4xf32>, %arg1: f32) -> !spv.array<4xf32> {
   // expected-error @+1 {{expected at least one index}}
   %0 = spv.CompositeInsert %arg1, %arg0[] : f32 into !spv.array<4xf32>
   return %0: !spv.array<4xf32>
 }
 
 // -----
 

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

 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
 }
 
 // -----
 
+// CHECK-LABEL: @cooperative_matrix_access_chain
+spv.func @cooperative_matrix_access_chain(%a : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>) -> !spv.ptr<f32, Function> "None" {
+  %0 = spv.constant 0: i32
+  // CHECK: {{%.*}} = spv.AccessChain {{%.*}}[{{%.*}}] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>
+  %1 = spv.AccessChain %a[%0] : !spv.ptr<!spv.coopmatrix<8x16xf32, Subgroup>, Function>
+  spv.ReturnValue %1 : !spv.ptr<f32, Function>
+}
+
+// -----
+
 spv.func @cooperative_matrix_muladd(%a : !spv.coopmatrix<16x16xi32, Subgroup>, %b : !spv.coopmatrix<16x8xi32, Subgroup>, %c : !spv.coopmatrix<8x8xi32, Subgroup>) "None" {
   // expected-error @+1 {{'spv.CooperativeMatrixMulAddNV' op matrix size must match}}
   %r = spv.CooperativeMatrixMulAddNV %a, %b, %c : !spv.coopmatrix<16x16xi32, Subgroup>, !spv.coopmatrix<16x8xi32, Subgroup> -> !spv.coopmatrix<8x8xi32, Subgroup>
   spv.Return
 }
 
 // -----