[mlir][spirv] Add MatrixTimesMatrix operation

[Hazem]

Summary

Add MatrixTimesMatrix operation to SPIRV Dialect and add NoSideEffect trait to Matrix ops.

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2960

By construct a MatrixType's element type is VectorType. so this is gonna succeed so there is no need to wrap it in a if statement. Essentially I think writing something like the following might look nicer:

c++
auto leftM = leftMatrix.getNumElements();
auto leftVector = leftMatrix.getElementType().cast<VectorType>();
auto leftN = leftVector.getNumElements();
auto leftElemType = leftVector.getElementType();


// Similarly for rightMatrix and result Matrix

What do you think?

[Hazem]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2960

I agree that this is better. In fact, I had some changes in mind for the matrix type, mainly:

1. add getNumRows() and getNumColumns() methods
2. refactor getNumElements to return total number of elements (rows*cols)
3. refactor getElementType to return the type of a single cell (e.g, F32) instead of Vector
4. Add getColumnType to return the Vector type of the matrix columns

These are minor changes coding wise, but they will make the verification methods much nicer than they are now, and the code easier to follow. And honestly they are intuitive for a matrix type. How that sounds?

[antiagainst]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2960

Sorry for the delay. SGTM! Please go for it! :)

[Hazem]

• add helpers for the matrix type
• update verification methods for matrix operations using the new helpers

[Hazem]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2960

Done, please have a look.

[antiagainst]

Cool!

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

 def SPV_OC_OpSDiv                      : I32EnumAttrCase<"OpSDiv", 135>;
 def SPV_OC_OpFDiv                      : I32EnumAttrCase<"OpFDiv", 136>;
 def SPV_OC_OpUMod                      : I32EnumAttrCase<"OpUMod", 137>;
 def SPV_OC_OpSRem                      : I32EnumAttrCase<"OpSRem", 138>;
 def SPV_OC_OpSMod                      : I32EnumAttrCase<"OpSMod", 139>;
 def SPV_OC_OpFRem                      : I32EnumAttrCase<"OpFRem", 140>;
 def SPV_OC_OpFMod                      : I32EnumAttrCase<"OpFMod", 141>;
 def SPV_OC_OpMatrixTimesScalar         : I32EnumAttrCase<"OpMatrixTimesScalar", 143>;
+def SPV_OC_OpMatrixTimesMatrix         : I32EnumAttrCase<"OpMatrixTimesMatrix", 146>;
 def SPV_OC_OpLogicalEqual              : I32EnumAttrCase<"OpLogicalEqual", 164>;
 def SPV_OC_OpLogicalNotEqual           : I32EnumAttrCase<"OpLogicalNotEqual", 165>;
 def SPV_OC_OpLogicalOr                 : I32EnumAttrCase<"OpLogicalOr", 166>;
 def SPV_OC_OpLogicalAnd                : I32EnumAttrCase<"OpLogicalAnd", 167>;
 def SPV_OC_OpLogicalNot                : I32EnumAttrCase<"OpLogicalNot", 168>;
 def SPV_OC_OpSelect                    : I32EnumAttrCase<"OpSelect", 169>;
 def SPV_OC_OpIEqual                    : I32EnumAttrCase<"OpIEqual", 170>;
 def SPV_OC_OpINotEqual                 : I32EnumAttrCase<"OpINotEqual", 171>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpAccessChain, SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate,
       SPV_OC_OpCompositeConstruct, SPV_OC_OpCompositeExtract,
       SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose, SPV_OC_OpConvertFToU,
       SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF, SPV_OC_OpConvertUToF,
       SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert, SPV_OC_OpBitcast,
       SPV_OC_OpFNegate, SPV_OC_OpIAdd, SPV_OC_OpFAdd, SPV_OC_OpISub, SPV_OC_OpFSub,
       SPV_OC_OpIMul, SPV_OC_OpFMul, SPV_OC_OpUDiv, SPV_OC_OpSDiv, SPV_OC_OpFDiv,
       SPV_OC_OpUMod, SPV_OC_OpSRem, SPV_OC_OpSMod, SPV_OC_OpFRem, SPV_OC_OpFMod,
-      SPV_OC_OpMatrixTimesScalar, SPV_OC_OpLogicalEqual, SPV_OC_OpLogicalNotEqual,
-      SPV_OC_OpLogicalOr, SPV_OC_OpLogicalAnd, SPV_OC_OpLogicalNot, SPV_OC_OpSelect,
-      SPV_OC_OpIEqual, SPV_OC_OpINotEqual, SPV_OC_OpUGreaterThan,
-      SPV_OC_OpSGreaterThan, SPV_OC_OpUGreaterThanEqual, SPV_OC_OpSGreaterThanEqual,
-      SPV_OC_OpULessThan, SPV_OC_OpSLessThan, SPV_OC_OpULessThanEqual,
-      SPV_OC_OpSLessThanEqual, SPV_OC_OpFOrdEqual, SPV_OC_OpFUnordEqual,
-      SPV_OC_OpFOrdNotEqual, SPV_OC_OpFUnordNotEqual, SPV_OC_OpFOrdLessThan,
-      SPV_OC_OpFUnordLessThan, SPV_OC_OpFOrdGreaterThan, SPV_OC_OpFUnordGreaterThan,
-      SPV_OC_OpFOrdLessThanEqual, SPV_OC_OpFUnordLessThanEqual,
-      SPV_OC_OpFOrdGreaterThanEqual, SPV_OC_OpFUnordGreaterThanEqual,
-      SPV_OC_OpShiftRightLogical, SPV_OC_OpShiftRightArithmetic,
-      SPV_OC_OpShiftLeftLogical, SPV_OC_OpBitwiseOr, SPV_OC_OpBitwiseXor,
-      SPV_OC_OpBitwiseAnd, SPV_OC_OpNot, SPV_OC_OpBitFieldInsert,
-      SPV_OC_OpBitFieldSExtract, SPV_OC_OpBitFieldUExtract, SPV_OC_OpBitReverse,
-      SPV_OC_OpBitCount, SPV_OC_OpControlBarrier, SPV_OC_OpMemoryBarrier,
-      SPV_OC_OpAtomicCompareExchangeWeak, SPV_OC_OpAtomicIIncrement,
-      SPV_OC_OpAtomicIDecrement, SPV_OC_OpAtomicIAdd, SPV_OC_OpAtomicISub,
-      SPV_OC_OpAtomicSMin, SPV_OC_OpAtomicUMin, SPV_OC_OpAtomicSMax,
-      SPV_OC_OpAtomicUMax, SPV_OC_OpAtomicAnd, SPV_OC_OpAtomicOr, SPV_OC_OpAtomicXor,
-      SPV_OC_OpPhi, SPV_OC_OpLoopMerge, SPV_OC_OpSelectionMerge, SPV_OC_OpLabel,
-      SPV_OC_OpBranch, SPV_OC_OpBranchConditional, SPV_OC_OpReturn,
-      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_OpCooperativeMatrixStoreNV, SPV_OC_OpCooperativeMatrixMulAddNV,
-      SPV_OC_OpCooperativeMatrixLengthNV
+      SPV_OC_OpMatrixTimesScalar, SPV_OC_OpMatrixTimesMatrix, SPV_OC_OpLogicalEqual,
+      SPV_OC_OpLogicalNotEqual, SPV_OC_OpLogicalOr, SPV_OC_OpLogicalAnd,
+      SPV_OC_OpLogicalNot, SPV_OC_OpSelect, SPV_OC_OpIEqual, SPV_OC_OpINotEqual,
+      SPV_OC_OpUGreaterThan, SPV_OC_OpSGreaterThan, SPV_OC_OpUGreaterThanEqual,
+      SPV_OC_OpSGreaterThanEqual, SPV_OC_OpULessThan, SPV_OC_OpSLessThan,
+      SPV_OC_OpULessThanEqual, SPV_OC_OpSLessThanEqual, SPV_OC_OpFOrdEqual,
+      SPV_OC_OpFUnordEqual, SPV_OC_OpFOrdNotEqual, SPV_OC_OpFUnordNotEqual,
+      SPV_OC_OpFOrdLessThan, SPV_OC_OpFUnordLessThan, SPV_OC_OpFOrdGreaterThan,
+      SPV_OC_OpFUnordGreaterThan, SPV_OC_OpFOrdLessThanEqual,
+      SPV_OC_OpFUnordLessThanEqual, SPV_OC_OpFOrdGreaterThanEqual,
+      SPV_OC_OpFUnordGreaterThanEqual, SPV_OC_OpShiftRightLogical,
+      SPV_OC_OpShiftRightArithmetic, SPV_OC_OpShiftLeftLogical, SPV_OC_OpBitwiseOr,
+      SPV_OC_OpBitwiseXor, SPV_OC_OpBitwiseAnd, SPV_OC_OpNot,
+      SPV_OC_OpBitFieldInsert, SPV_OC_OpBitFieldSExtract, SPV_OC_OpBitFieldUExtract,
+      SPV_OC_OpBitReverse, SPV_OC_OpBitCount, SPV_OC_OpControlBarrier,
+      SPV_OC_OpMemoryBarrier, SPV_OC_OpAtomicCompareExchangeWeak,
+      SPV_OC_OpAtomicIIncrement, SPV_OC_OpAtomicIDecrement, SPV_OC_OpAtomicIAdd,
+      SPV_OC_OpAtomicISub, SPV_OC_OpAtomicSMin, SPV_OC_OpAtomicUMin,
+      SPV_OC_OpAtomicSMax, SPV_OC_OpAtomicUMax, SPV_OC_OpAtomicAnd,
+      SPV_OC_OpAtomicOr, SPV_OC_OpAtomicXor, SPV_OC_OpPhi, SPV_OC_OpLoopMerge,
+      SPV_OC_OpSelectionMerge, SPV_OC_OpLabel, SPV_OC_OpBranch,
+      SPV_OC_OpBranchConditional, SPV_OC_OpReturn, 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_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/SPIRVMatrixOps.td

 //===-- SPIRVMatrixOps.td - MLIR SPIR-V Matrix Ops ---------*- 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 file contains matrix operations for the SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPIRV_MATRIX_OPS
 #define SPIRV_MATRIX_OPS
+include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // -----
 
-def SPV_MatrixTimesScalarOp : SPV_Op<"MatrixTimesScalar", []> {
+def SPV_MatrixTimesMatrixOp : SPV_Op<"MatrixTimesMatrix", [NoSideEffect]> {
+  let summary = "Linear-algebraic multiply of LeftMatrix X RightMatrix.";
+
+  let description = [{
+    Result Type must be an OpTypeMatrix whose Column Type is a vector of
+    floating-point type.
+
+    LeftMatrix must be a matrix whose Column Type is the same as the Column
+    Type in Result Type.
+
+    RightMatrix must be a matrix with the same Component Type as the
+    Component Type in Result Type. Its number of columns must equal the
+    number of columns in Result Type. Its columns must have the same number
+    of components as the number of columns in LeftMatrix.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    matrix-times-matrix-op ::= ssa-id `=` `spv.MatrixTimesMatrix` ssa-use,
+    ssa-use `:` matrix-type `,` matrix-type `->` matrix-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %0 = spv.MatrixTimesMatrix %matrix_1, %matrix_2 :
+        !spv.matrix<4 x vector<3xf32>>, !spv.matrix<3 x vector<4xf32>> ->
+        !spv.matrix<4 x vector<4xf32>>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_Matrix]>
+  ];
+
+  let arguments = (ins
+    SPV_AnyMatrix:$leftmatrix,
+    SPV_AnyMatrix:$rightmatrix
+  );
+
+  let results = (outs
+    SPV_AnyMatrix:$result
+  );
+  let assemblyFormat = [{
+    operands attr-dict `:` type($leftmatrix) `,` type($rightmatrix) `->` type($result)
+  }];
+  let verifier = [{ return verifyMatrixTimesMatrix(*this); }];
+}
+
+// -----
+
+def SPV_MatrixTimesScalarOp : SPV_Op<"MatrixTimesScalar", [NoSideEffect]> {
   let summary = "Scale a floating-point matrix.";
 
   let description = [{
     Result Type must be an OpTypeMatrix whose Column Type is a vector of
     floating-point type.
 
      The type of Matrix must be the same as Result Type. Each component in
     each column in Matrix is multiplied by Scalar.
let availability = [
     Capability<[SPV_C_Matrix]>
   ];
 
   let verifier = [{ return verifyMatrixTimesScalar(*this); }];
 }
 
 // -----
 
-def SPV_TransposeOp : SPV_Op<"Transpose", []> {
+def SPV_TransposeOp : SPV_Op<"Transpose", [NoSideEffect]> {
   let summary = "Transpose a matrix.";
 
   let description = [{
     Result Type must be an OpTypeMatrix.
 
     Matrix must be an object of type OpTypeMatrix. The number of columns and
     the column size of Matrix must be the reverse of those in Result Type.
     The types of the scalar components in Matrix and Result Type must be the

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

public:
 
   static MatrixType getChecked(Type columnType, uint32_t columnCount,
                                Location location);
 
   static LogicalResult verifyConstructionInvariants(Location loc,
                                                     Type columnType,
                                                     uint32_t columnCount);
 
-  /// Returns true if the matrix elements are vectors of float elements
+  /// Returns true if the matrix elements are vectors of float elements.
   static bool isValidColumnType(Type columnType);
 
-  Type getElementType() const;
+  Type getColumnType() const;
+
+  /// Returns the number of rows.
+  unsigned getNumRows() const;
+
+  /// Returns the number of columns.
+  unsigned getNumColumns() const;
 
+  /// Returns total number of elements (rows*columns).
   unsigned getNumElements() const;
 
+  /// Returns the elements' type (i.e, single element type).
+  Type getElementType() const;
+
   void getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                      Optional<spirv::StorageClass> storage = llvm::None);
   void getCapabilities(SPIRVType::CapabilityArrayRefVector &capabilities,
                        Optional<spirv::StorageClass> storage = llvm::None);
 };
 
 } // end namespace spirv
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVTYPES_H_

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

 
 static void print(CooperativeMatrixNVType type, DialectAsmPrinter &os) {
   os << "coopmatrix<" << type.getRows() << "x" << type.getColumns() << "x";
   os << type.getElementType() << ", " << stringifyScope(type.getScope());
   os << ">";
 }
 
 static void print(MatrixType type, DialectAsmPrinter &os) {
-  os << "matrix<" << type.getNumElements() << " x " << type.getElementType();
+  os << "matrix<" << type.getNumColumns() << " x " << type.getColumnType();
   os << ">";
 }
 
 void SPIRVDialect::printType(Type type, DialectAsmPrinter &os) const {
   switch (type.getKind()) {
   case TypeKind::Array:
     print(type.cast<ArrayType>(), os);
     return;

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 // spv.MatrixTimesScalar
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verifyMatrixTimesScalar(spirv::MatrixTimesScalarOp op) {
   // We already checked that result and matrix are both of matrix type in the
   // auto-generated verify method.
 
   auto inputMatrix = op.matrix().getType().cast<spirv::MatrixType>();
-  // Check that the scalar type is the same as the matrix components type.
-  if (auto inputMatrixColumns =
-          inputMatrix.getElementType().dyn_cast<VectorType>()) {
-    if (op.scalar().getType() != inputMatrixColumns.getElementType())
-      return op.emitError("input matrix components' type and scaling "
-                          "value must have the same type");
+  auto resultMatrix = op.result().getType().cast<spirv::MatrixType>();
+
+  // Check that the scalar type is the same as the matrix element type.
+  if (op.scalar().getType() != inputMatrix.getElementType())
+    return op.emitError("input matrix components' type and scaling value must "
+                        "have the same type");
 
-    // Note that the next three checks could be done using the AllTypesMatch
-    // trait in the Op definition file but it generates a vague error message.
 
-    // Check that the input and result matrices have the same size
-    auto resultMatrix = op.result().getType().cast<spirv::MatrixType>();
-    if (inputMatrix.getNumElements() != resultMatrix.getNumElements())
-      return op.emitError("input and result matrices must have "
-                          "the same number of columns");
-
-    if (auto resultMatrixColumns =
-            resultMatrix.getElementType().dyn_cast<VectorType>()) {
-      // Check that the input and result matrices' columns have the same type
-      if (inputMatrixColumns.getElementType() !=
-          resultMatrixColumns.getElementType())
-        return op.emitError("input and result matrices' columns must "
-                            "have the same component type");
-
-      // Check that the input and result matrices' columns have the same size
-      if (inputMatrixColumns.getNumElements() !=
-          resultMatrixColumns.getNumElements())
-        return op.emitError("input and result matrices' columns must "
-                            "have the same size");
-    }
-  }
+  // Check that the input and result matrices have the same columns' count
+  if (inputMatrix.getNumColumns() != resultMatrix.getNumColumns())
+    return op.emitError("input and result matrices must have the same "
+                        "number of columns");
+
+  // Check that the input and result matrices' have the same rows count
+  if (inputMatrix.getNumRows() != resultMatrix.getNumRows())
+    return op.emitError("input and result matrices' columns must have "
+                        "the same size");
+
+  // Check that the input and result matrices' have the same component type
+  if (inputMatrix.getElementType() != resultMatrix.getElementType())
+    return op.emitError("input and result matrices' columns must have "
+                        "the same component type");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.CopyMemory
 //===----------------------------------------------------------------------===//
 

 // spv.Transpose
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verifyTranspose(spirv::TransposeOp op) {
   auto inputMatrix = op.matrix().getType().cast<spirv::MatrixType>();
   auto resultMatrix = op.result().getType().cast<spirv::MatrixType>();
 
   // Verify that the input and output matrices have correct shapes.
-  if (auto inputMatrixColumns =
-          inputMatrix.getElementType().dyn_cast<VectorType>()) {
-    if (inputMatrixColumns.getNumElements() != resultMatrix.getNumElements())
-      return op.emitError("input matrix rows count must be equal to "
-                          "output matrix columns count");
-    if (auto resultMatrixColumns =
-            resultMatrix.getElementType().dyn_cast<VectorType>()) {
-      if (resultMatrixColumns.getNumElements() != inputMatrix.getNumElements())
-        return op.emitError("input matrix columns count must be equal "
-                            "to output matrix rows count");
+  if (inputMatrix.getNumRows() != resultMatrix.getNumColumns())
+
+  if (inputMatrix.getNumColumns() != resultMatrix.getNumRows())
+    return op.emitError("input matrix columns count must be equal to "
+                        "output matrix rows count");
 
-      // Verify that the input and output matrices have the same component type
-      if (inputMatrixColumns.getElementType() !=
-          resultMatrixColumns.getElementType())
-        return op.emitError("input and output matrices must have the "
-                            "same component type");
-    }
-  }
+  if (inputMatrix.getElementType() != resultMatrix.getElementType())
+    return op.emitError("input and output matrices must have the same "
+                        "component type");
+
+  return success();
+
+//===----------------------------------------------------------------------===//
+// spv.MatrixTimesMatrix
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verifyMatrixTimesMatrix(spirv::MatrixTimesMatrixOp op) {
+  auto leftMatrix = op.leftmatrix().getType().cast<spirv::MatrixType>();
+  auto rightMatrix = op.rightmatrix().getType().cast<spirv::MatrixType>();
+  auto resultMatrix = op.result().getType().cast<spirv::MatrixType>();
+
+  // left matrix columns' count and right matrix rows' count must be equal
+  if (leftMatrix.getNumColumns() != rightMatrix.getNumRows())
+    return op.emitError("left matrix columns' count must be equal to "
+                        "the right matrix rows' count");
+
+  // right and result matrices columns' count must be the same
+  if (rightMatrix.getNumColumns() != resultMatrix.getNumColumns())
+    return op.emitError(
+        "right and result matrices must have equal columns' count");
+
+  // right and result matrices component type must be the same
+  if (rightMatrix.getElementType() != resultMatrix.getElementType())
+    return op.emitError("right and result matrices' component type must"
+                        " be the same");
+
+  // left and result matrices component type must be the same
+  if (leftMatrix.getElementType() != resultMatrix.getElementType())
+    return op.emitError("left and result matrices' component type"
+                        " must be the same");
+
+  // left and result matrices rows count must be the same
+  if (leftMatrix.getNumRows() != resultMatrix.getNumRows())
+    return op.emitError("left and result matrices must have equal rows'"
+                        " count");
+
   return success();
 }
 
 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/lib/Dialect/SPIRV/SPIRVTypes.cpp

 
 Type CompositeType::getElementType(unsigned index) const {
   switch (getKind()) {
   case spirv::TypeKind::Array:
     return cast<ArrayType>().getElementType();
   case spirv::TypeKind::CooperativeMatrix:
     return cast<CooperativeMatrixNVType>().getElementType();
   case spirv::TypeKind::Matrix:
-    return cast<MatrixType>().getElementType();
+    return cast<MatrixType>().getColumnType();
   case spirv::TypeKind::RuntimeArray:
     return cast<RuntimeArrayType>().getElementType();
   case spirv::TypeKind::Struct:
     return cast<StructType>().getElementType(index);
   case StandardTypes::Vector:
     return cast<VectorType>().getElementType();
   default:
     llvm_unreachable("invalid composite type");
   }
 }
 
 unsigned CompositeType::getNumElements() const {
   switch (getKind()) {
   case spirv::TypeKind::Array:
     return cast<ArrayType>().getNumElements();
   case spirv::TypeKind::CooperativeMatrix:
     llvm_unreachable(
         "invalid to query number of elements of spirv::CooperativeMatrix type");
   case spirv::TypeKind::Matrix:
-    return cast<MatrixType>().getNumElements();
+    return cast<MatrixType>().getNumColumns();
   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:

 bool MatrixType::isValidColumnType(Type columnType) {
   if (auto vectorType = columnType.dyn_cast<VectorType>()) {
     if (vectorType.getElementType().isa<FloatType>())
       return true;
   }
   return false;
 }
 
-Type MatrixType::getElementType() const { return getImpl()->columnType; }
+Type MatrixType::getColumnType() const { return getImpl()->columnType; }
 
-unsigned MatrixType::getNumElements() const { return getImpl()->columnCount; }
+Type MatrixType::getElementType() const {
+  return getImpl()->columnType.cast<VectorType>().getElementType();
+}
+
+unsigned MatrixType::getNumColumns() const { return getImpl()->columnCount; }
+
+unsigned MatrixType::getNumRows() const {
+  return getImpl()->columnType.cast<VectorType>().getShape()[0];
+}
+
+unsigned MatrixType::getNumElements() const {
+  return (getImpl()->columnCount) * getNumRows();
+}
 
 void MatrixType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
                                Optional<StorageClass> storage) {
-  getElementType().cast<SPIRVType>().getExtensions(extensions, storage);
+  getColumnType().cast<SPIRVType>().getExtensions(extensions, storage);
 }
 
 void MatrixType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
   {
     static const Capability caps[] = {Capability::Matrix};
     ArrayRef<Capability> ref(caps, llvm::array_lengthof(caps));
     capabilities.push_back(ref);
   }
   // Add any capabilities associated with the underlying vectors (i.e., columns)
-  getElementType().cast<SPIRVType>().getCapabilities(capabilities, storage);
+  getColumnType().cast<SPIRVType>().getCapabilities(capabilities, storage);
 }

mlir/lib/Dialect/SPIRV/Serialization/Serializer.cpp

operands.push_back(
         getConstantOp(static_cast<uint32_t>(cooperativeMatrixType.getScope())));
     operands.push_back(getConstantOp(cooperativeMatrixType.getRows()));
     operands.push_back(getConstantOp(cooperativeMatrixType.getColumns()));
     return success();
   }
 
   if (auto matrixType = type.dyn_cast<spirv::MatrixType>()) {
     uint32_t elementTypeID = 0;
-    if (failed(processType(loc, matrixType.getElementType(), elementTypeID))) {
+    if (failed(processType(loc, matrixType.getColumnType(), elementTypeID))) {
       return failure();
     }
     typeEnum = spirv::Opcode::OpTypeMatrix;
     operands.push_back(elementTypeID);
-    operands.push_back(matrixType.getNumElements());
+    operands.push_back(matrixType.getNumColumns());
     return success();
   }
 
   // TODO: Handle other types.
   return emitError(loc, "unhandled type in serialization: ") << type;
 }
 
 LogicalResult

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

spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   }
 
   // CHECK-LABEL: @matrix_transpose_1
   spv.func @matrix_transpose_1(%arg0 : !spv.matrix<3 x vector<2xf32>>) -> !spv.matrix<2 x vector<3xf32>> "None" {
     // CHECK: {{%.*}} = spv.Transpose {{%.*}} : !spv.matrix<3 x vector<2xf32>> -> !spv.matrix<2 x vector<3xf32>>
     %result = spv.Transpose %arg0 : !spv.matrix<3 x vector<2xf32>> -> !spv.matrix<2 x vector<3xf32>>
     spv.ReturnValue %result : !spv.matrix<2 x vector<3xf32>>
   }
+
+  // CHECK-LABEL: @matrix_times_matrix_1
+  spv.func @matrix_times_matrix_1(%arg0: !spv.matrix<3 x vector<3xf32>>, %arg1: !spv.matrix<3 x vector<3xf32>>) -> !spv.matrix<3 x vector<3xf32>> "None"{
+    // CHECK: {{%.*}} = spv.MatrixTimesMatrix {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<3xf32>>, !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
+    %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
+    spv.ReturnValue %result : !spv.matrix<3 x vector<3xf32>>
+  }
+
+  // CHECK-LABEL: @matrix_times_matrix_2
+  spv.func @matrix_times_matrix_2(%arg0: !spv.matrix<3 x vector<2xf32>>, %arg1: !spv.matrix<2 x vector<3xf32>>) -> !spv.matrix<2 x vector<2xf32>> "None"{
+    // CHECK: {{%.*}} = spv.MatrixTimesMatrix {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<3xf32>> -> !spv.matrix<2 x vector<2xf32>>
+    %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<3xf32>> -> !spv.matrix<2 x vector<2xf32>>
+    spv.ReturnValue %result : !spv.matrix<2 x vector<2xf32>>
+  }
 }
 
 // -----
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   // CHECK: spv.globalVariable {{@.*}} : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, StorageBuffer>
   spv.globalVariable @var0 : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, StorageBuffer>
 
   // CHECK: spv.globalVariable {{@.*}} : !spv.ptr<!spv.matrix<2 x vector<3xf32>>, StorageBuffer>
   spv.globalVariable @var1 : !spv.ptr<!spv.matrix<2 x vector<3xf32>>, StorageBuffer>
 
   // CHECK: spv.globalVariable {{@.*}} : !spv.ptr<!spv.matrix<4 x vector<4xf16>>, StorageBuffer>
   spv.globalVariable @var2 : !spv.ptr<!spv.matrix<4 x vector<4xf16>>, StorageBuffer>
 }

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

spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   }
 
   // CHECK-LABEL: @matrix_transpose_2
   spv.func @matrix_transpose_2(%arg0 : !spv.matrix<3 x vector<3xf32>>) -> !spv.matrix<3 x vector<3xf32>> "None" {
     // CHECK: {{%.*}} = spv.Transpose {{%.*}} : !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
     %result = spv.Transpose %arg0 : !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
     spv.ReturnValue %result : !spv.matrix<3 x vector<3xf32>>
   }
+
+  // CHECK-LABEL: @matrix_times_matrix_1
+  spv.func @matrix_times_matrix_1(%arg0: !spv.matrix<3 x vector<3xf32>>, %arg1: !spv.matrix<3 x vector<3xf32>>) -> !spv.matrix<3 x vector<3xf32>> "None"{
+    // CHECK: {{%.*}} = spv.MatrixTimesMatrix {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<3xf32>>, !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
+    %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
+    spv.ReturnValue %result : !spv.matrix<3 x vector<3xf32>>
+  }
+
+  // CHECK-LABEL: @matrix_times_matrix_2
+  spv.func @matrix_times_matrix_2(%arg0: !spv.matrix<3 x vector<2xf32>>, %arg1: !spv.matrix<2 x vector<3xf32>>) -> !spv.matrix<2 x vector<2xf32>> "None"{
+    // CHECK: {{%.*}} = spv.MatrixTimesMatrix {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<3xf32>> -> !spv.matrix<2 x vector<2xf32>>
+    %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<3xf32>> -> !spv.matrix<2 x vector<2xf32>>
+    spv.ReturnValue %result : !spv.matrix<2 x vector<2xf32>>
+  }
 }
 
 // -----
 
 func @input_type_mismatch(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : f16) -> () {
   // expected-error @+1 {{input matrix components' type and scaling value must have the same type}}
   %result = spv.MatrixTimesScalar %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, f16 -> !spv.matrix<3 x vector<3xf32>>
 }

 
 // -----
 
 func @transpose_op_type_mismatch(%arg0 : !spv.matrix<3 x vector<4xf32>>) -> () {
    // expected-error @+1 {{input and output matrices must have the same component type}}
    %result = spv.Transpose %arg0 : !spv.matrix<3 x vector<4xf32>> -> !spv.matrix<4 x vector<3xf16>>
    spv.Return
 }
+
+// -----
+
+func @matrix_times_matrix_invalid_input_shape_1(%arg0 : !spv.matrix<3 x vector<2xf32>>, %arg1 : !spv.matrix<2 x vector<3xf32>>){
+   // expected-error @+1 {{right and result matrices must have equal columns' count}}
+   %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<3xf32>> -> !spv.matrix<3 x vector<2xf32>>
+}
+
+// -----
+
+func @matrix_times_matrix_invalid_input_shape_2(%arg0 : !spv.matrix<3 x vector<2xf32>>, %arg1 : !spv.matrix<2 x vector<3xf32>>){
+   // expected-error @+1 {{left and result matrices must have equal rows' count}}
+   %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<3xf32>> -> !spv.matrix<2 x vector<3xf32>>
+}
+
+// -----
+
+func @matrix_times_matrix_inputs_shape_mismatch(%arg0 : !spv.matrix<3 x vector<2xf32>>, %arg1 : !spv.matrix<2 x vector<2xf32>>){
+   // expected-error @+1 {{left matrix columns' count must be equal to the right matrix rows' count}}
+   %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<2xf32>>, !spv.matrix<2 x vector<2xf32>> -> !spv.matrix<2 x vector<2xf32>>
+}
+
+// -----
+
+func @matrix_times_matrix_component_type_mismatch_1(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : !spv.matrix<3x vector<3xf32>>){
+   // expected-error @+1 {{right and result matrices' component type must be the same}}
+   %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf64>>
+}
+
+
+// -----
+
+func @matrix_times_matrix_component_type_mismatch_2(%arg0 : !spv.matrix<3 x vector<3xf64>>, %arg1 : !spv.matrix<3x vector<3xf32>>){
+   // expected-error @+1 {{left and result matrices' component type must be the same}}
+   %result = spv.MatrixTimesMatrix %arg0, %arg1 : !spv.matrix<3 x vector<3xf64>>, !spv.matrix<3 x vector<3xf32>> -> !spv.matrix<3 x vector<3xf32>>
+}