Address feedback

Author: IgWod-IMG

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

@@ -18,21 +18,12 @@ include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
-class SPIRV_SameCoopMatrix<string lhs, string rhs> : PredOpTrait<
-  "cooperative matrix types match",
-  CPred<"(::llvm::isa<::mlir::spirv::CooperativeMatrixType>($" # lhs # ".getType()) "
-        "&& ::llvm::isa<::mlir::spirv::CooperativeMatrixType>($" # rhs # ".getType()))"
-        "? $" # lhs # ".getType() == $" # rhs # ".getType() : true">
->;
-
 class SPIRV_ArithmeticBinaryOp<string mnemonic, Type type,
                                list<Trait> traits = []> :
       // Operands type same as result type.
       SPIRV_BinaryOp<mnemonic, type, type,
                    !listconcat(traits,
-                               [Pure, SameOperandsAndResultType,
-                                SPIRV_SameCoopMatrix<"operand1", "operand2">,
-                                SPIRV_SameCoopMatrix<"operand2", "result">])> {
+                               [Pure, AllTypesMatch<["operand1", "operand2", "result"]>])> {
   // In addition to normal types arithmetic instructions can support cooperative
   // matrix.
   let arguments = (ins
@@ -51,8 +42,7 @@ class SPIRV_ArithmeticUnaryOp<string mnemonic, Type type,
       // Operand type same as result type.
       SPIRV_UnaryOp<mnemonic, type, type,
                    !listconcat(traits,
-                               [Pure, SameOperandsAndResultType,
-                                SPIRV_SameCoopMatrix<"operand", "result">])> {
+                               [Pure, AllTypesMatch<["operand", "result"]>])> {
   // In addition to normal types arithmetic instructions can support cooperative
   // matrix.
   let arguments = (ins

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

@@ -428,13 +428,12 @@ class CooperativeMatrixType
 
   CooperativeMatrixType cloneWith(std::optional<ArrayRef<int64_t>> shape,
                                   Type elementType) const {
-    if (shape == std::nullopt)
+    if (!shape)
       return get(elementType, getRows(), getColumns(), getScope(), getUse());
-    else {
-      assert(shape.value().size() == 2);
-      return get(elementType, shape.value()[0], shape.value()[1], getScope(),
-                 getUse());
-    }
+
+    assert(shape.value().size() == 2);
+    return get(elementType, shape.value()[0], shape.value()[1], getScope(),
+               getUse());
   }
 };
 

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

@@ -194,6 +194,19 @@ std::optional<int64_t> CompositeType::getSizeInBytes() {
 //===----------------------------------------------------------------------===//
 
 struct spirv::detail::CooperativeMatrixTypeStorage final : TypeStorage {
+  // In the specification dimensions of the Cooperative Matrix are 32-bit
+  // integers --- the initial implementation kept those values as such. However,
+  // the `ShapedType` expects the shape to be `int64_t`. We could keep the shape
+  // as 32-bits and expose it as int64_t through `getShape`, however, this
+  // method returns an `ArrayRef`, so returning `ArrayRef<int64_t>` having two
+  // 32-bits integers would require an extra logic and storage. So, we diverge
+  // from the spec and internally represent the dimensions as 64-bit integers,
+  // so we can easily return an `ArrayRef` from `getShape` without any extra
+  // logic. Alternatively, we could store both rows and columns (both 32-bits)
+  // and shape (64-bits), assigning rows and columns to shape whenever
+  // `getShape` is called. This would be at the cost of extra logic and storage.
+  // Note: Because `ArrayRef` is returned we cannot construct an object in
+  // `getShape` on the fly.
   using KeyTy =
       std::tuple<Type, int64_t, int64_t, Scope, CooperativeMatrixUseKHR>;
 
@@ -214,7 +227,7 @@ struct spirv::detail::CooperativeMatrixTypeStorage final : TypeStorage {
 
   Type elementType;
   // [#rows, #columns]
-  SmallVector<int64_t, 2> shape;
+  std::array<int64_t, 2> shape;
   Scope scope;
   CooperativeMatrixUseKHR use;
 };

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

@@ -524,7 +524,7 @@ spirv.func @matrix_times_scalar(%a: !matA_f32, %b: f32) "None" {
 
 spirv.func @iadd(%a: !spirv.coopmatrix<2x2xi32, Subgroup, MatrixA>,
                  %b: !spirv.coopmatrix<2x2xi32, Subgroup, MatrixB>) "None" {
-  // expected-error @+1 {{op failed to verify that cooperative matrix types match}}
+  // expected-error @+1 {{failed to verify that all of {operand1, operand2, result} have same type}}
   %q = "spirv.IAdd"(%a, %b) :
     (!spirv.coopmatrix<2x2xi32, Subgroup, MatrixA>, !spirv.coopmatrix<2x2xi32, Subgroup, MatrixB>)
     -> !spirv.coopmatrix<2x2xi32, Subgroup, MatrixA>
@@ -535,7 +535,7 @@ spirv.func @iadd(%a: !spirv.coopmatrix<2x2xi32, Subgroup, MatrixA>,
 
 spirv.func @fadd(%a: !spirv.coopmatrix<2x2xf32, Subgroup, MatrixA>,
                  %b: !spirv.coopmatrix<2x2xf32, Subgroup, MatrixAcc>) "None" {
-  // expected-error @+1 {{op failed to verify that cooperative matrix types match}}
+  // expected-error @+1 {{failed to verify that all of {operand1, operand2, result} have same type}}
   %q = "spirv.FAdd"(%a, %b) :
     (!spirv.coopmatrix<2x2xf32, Subgroup, MatrixA>, !spirv.coopmatrix<2x2xf32, Subgroup, MatrixAcc>)
     -> !spirv.coopmatrix<2x2xf32, Subgroup, MatrixAcc>