[MLIR][SPIRV] Adding new data type and decorators to SPIRV Dialect

[Hazem]

Summary

Add a new data type to the SPIRV dialect: matrix type.
Add several decorators needed for successful deserialization of graphics shaders

[antiagainst]

Awesome, thanks a lot @Hazem for the contribution! I have a few nits inlined. One major thing: could you add (de)serialization round-trip tests for matrix under https://github.com/llvm/llvm-project/tree/master/mlir/test/Dialect/SPIRV/Serialization? The tests at the moment only covers IR parsing/printing round-tripping. :)

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

64

Could you keep this alphabetically sorted?

75

Same here. Please keep it alphabetically sorted. :)

381

Maybe renaming the paramters as columnType and numColumns to be consistent with the spec? Similarly for the rest.

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

120

Same here for alphabetical order.

208

I think we can use CompositeType::isValid(t) to check the vector rank and num elements.

221

Generally error messages should start with lower case; this composes well in sentences.

408

Lower-case to start and no need for ending period.

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

164

Alphabetical order.

1027

Keep this together with the other field?

1043

No need for this getters given this is a struct anyway.

1080

if (auto vecType = elementType.dy_cast<VectorType>(...)) 
  if (vecType.getElementType().isa...)

[Hazem]

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

208

CompositeType::isValid(t) will return true for IntegerTypes as well, Matrix need to have Float types, unless you meant to use it then re-check that the vector elements are Float and not Integer in the parseAndVerifyMatrixType. Please let me know which one you had in mind. Thanks.

[mravishankar]

Did a glance through of this. Mostly looks fine to me. Only question is the decorations added here dont seem to be related to the matrix type as such. If so can we split those into two separate CLs. And also I dont see any test for the roundtrip through SPIR-V binary for those decorations....

[Hazem]

In D80594#2057716, @mravishankar wrote:

Did a glance through of this. Mostly looks fine to me. Only question is the decorations added here dont seem to be related to the matrix type as such. If so can we split those into two separate CLs. And also I dont see any test for the roundtrip through SPIR-V binary for those decorations....

Makes sense, and just to be sure I understood correctly, I will submit the decorators along with their tests in a separate patch and eliminate them from this matrix type patch?

[Hazem]

Updating D80594: [MLIR][SPIRV] Addressed revision changes requested, and removed decorators to another CL as requested.

[antiagainst]

Awesome, thanks for the contribution!

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

208

Actually both are fine. The current way makes sure we can have more meaningful error message, at the expense of duplicating the logic a bit. It's not too much duplication so it's fine. Thanks! :)

[antiagainst]

Do you have commit access to llvm-project mono repo? Let me know if not then I can land this for you.

[Hazem]

In D80594#2061575, @antiagainst wrote:

Do you have commit access to llvm-project mono repo? Let me know if not then I can land this for you.

Based on the documentation here https://mlir.llvm.org/getting_started/Contributing/#first-time-contributors
and since I am a new contributor, I don't have direct access to commit and I have to request from you if you please to land this for me.
Thanks!

[antiagainst]

@Hazem: sure no problem! sorry for the delay; was OOO for a few days. :)

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

 def SPV_OC_OpEntryPoint                : I32EnumAttrCase<"OpEntryPoint", 15>;
 def SPV_OC_OpExecutionMode             : I32EnumAttrCase<"OpExecutionMode", 16>;
 def SPV_OC_OpCapability                : I32EnumAttrCase<"OpCapability", 17>;
 def SPV_OC_OpTypeVoid                  : I32EnumAttrCase<"OpTypeVoid", 19>;
 def SPV_OC_OpTypeBool                  : I32EnumAttrCase<"OpTypeBool", 20>;
 def SPV_OC_OpTypeInt                   : I32EnumAttrCase<"OpTypeInt", 21>;
 def SPV_OC_OpTypeFloat                 : I32EnumAttrCase<"OpTypeFloat", 22>;
 def SPV_OC_OpTypeVector                : I32EnumAttrCase<"OpTypeVector", 23>;
+def SPV_OC_OpTypeMatrix                : I32EnumAttrCase<"OpTypeMatrix", 24>;
 def SPV_OC_OpTypeArray                 : I32EnumAttrCase<"OpTypeArray", 28>;
 def SPV_OC_OpTypeRuntimeArray          : I32EnumAttrCase<"OpTypeRuntimeArray", 29>;
 def SPV_OC_OpTypeStruct                : I32EnumAttrCase<"OpTypeStruct", 30>;
 def SPV_OC_OpTypePointer               : I32EnumAttrCase<"OpTypePointer", 32>;
 def SPV_OC_OpTypeFunction              : I32EnumAttrCase<"OpTypeFunction", 33>;
 def SPV_OC_OpConstantTrue              : I32EnumAttrCase<"OpConstantTrue", 41>;
 def SPV_OC_OpConstantFalse             : I32EnumAttrCase<"OpConstantFalse", 42>;
 def SPV_OC_OpConstant                  : I32EnumAttrCase<"OpConstant", 43>;

 
 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_OC_OpTypeFloat, SPV_OC_OpTypeVector, SPV_OC_OpTypeArray,
-      SPV_OC_OpTypeRuntimeArray, SPV_OC_OpTypeStruct, SPV_OC_OpTypePointer,
-      SPV_OC_OpTypeFunction, SPV_OC_OpConstantTrue, SPV_OC_OpConstantFalse,
-      SPV_OC_OpConstant, SPV_OC_OpConstantComposite, SPV_OC_OpConstantNull,
-      SPV_OC_OpSpecConstantTrue, SPV_OC_OpSpecConstantFalse, SPV_OC_OpSpecConstant,
-      SPV_OC_OpSpecConstantComposite, SPV_OC_OpFunction, SPV_OC_OpFunctionParameter,
-      SPV_OC_OpFunctionEnd, SPV_OC_OpFunctionCall, SPV_OC_OpVariable, SPV_OC_OpLoad,
-      SPV_OC_OpStore, SPV_OC_OpAccessChain, SPV_OC_OpDecorate,
-      SPV_OC_OpMemberDecorate, SPV_OC_OpCompositeConstruct,
+      SPV_OC_OpTypeFloat, SPV_OC_OpTypeVector, SPV_OC_OpTypeMatrix,
+      SPV_OC_OpTypeArray, SPV_OC_OpTypeRuntimeArray, SPV_OC_OpTypeStruct,
+      SPV_OC_OpTypePointer, SPV_OC_OpTypeFunction, SPV_OC_OpConstantTrue,
+      SPV_OC_OpConstantFalse, SPV_OC_OpConstant, SPV_OC_OpConstantComposite,
+      SPV_OC_OpConstantNull, SPV_OC_OpSpecConstantTrue, SPV_OC_OpSpecConstantFalse,
+      SPV_OC_OpSpecConstant, SPV_OC_OpSpecConstantComposite, SPV_OC_OpFunction,
+      SPV_OC_OpFunctionParameter, SPV_OC_OpFunctionEnd, SPV_OC_OpFunctionCall,
+      SPV_OC_OpVariable, SPV_OC_OpLoad, SPV_OC_OpStore, SPV_OC_OpAccessChain,
+      SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate, SPV_OC_OpCompositeConstruct,
       SPV_OC_OpCompositeExtract, SPV_OC_OpCompositeInsert, 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_OpLogicalEqual, SPV_OC_OpLogicalNotEqual, SPV_OC_OpLogicalOr,
       SPV_OC_OpLogicalAnd, SPV_OC_OpLogicalNot, SPV_OC_OpSelect, SPV_OC_OpIEqual,

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

 //===- SPIRVTypes.h - MLIR SPIR-V Types -------------------------*- C++ -*-===//
 //
 // 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 declares the types in the SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_SPIRVTYPES_H_
 #define MLIR_DIALECT_SPIRV_SPIRVTYPES_H_
 
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Location.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 
 #include <tuple>
 
 // Forward declare enum classes related to op availability. Their definitions
 // are in the TableGen'erated SPIRVEnums.h.inc and can be referenced by other

 
 namespace detail {
 struct ArrayTypeStorage;
 struct CooperativeMatrixTypeStorage;
 struct ImageTypeStorage;
 struct PointerTypeStorage;
 struct RuntimeArrayTypeStorage;
 struct StructTypeStorage;
+struct MatrixTypeStorage;
 } // namespace detail
 
 namespace TypeKind {
 enum Kind {
   Array = Type::FIRST_SPIRV_TYPE,
   CooperativeMatrix,
   Image,
   Pointer,
   RuntimeArray,
   Struct,
-  LAST_SPIRV_TYPE = Struct,
+  Matrix,
+  LAST_SPIRV_TYPE = Matrix,
 };
 }
 
 // Base SPIR-V type for providing availability queries.
 class SPIRVType : public Type {
 public:
   using Type::Type;
 
public:
   unsigned getColumns() 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 matrix type
+class MatrixType : public Type::TypeBase<MatrixType, CompositeType,
+                                         detail::MatrixTypeStorage> {
+public:
+  using Base::Base;
+
+  static bool kindof(unsigned kind) { return kind == TypeKind::Matrix; }
+
+  static MatrixType get(Type elementType, uint32_t numCols);
+
+  static MatrixType getChecked(Type elementType, uint32_t numCols,
+                               Location location);
+
+  static LogicalResult verifyConstructionInvariants(Location loc,
+                                                    Type elementType,
+                                                    uint32_t numCols);
+
+  /// Returns true if the matrix elements are vectors of float elements
+  static bool isValidElementType(Type elementType);
+
+  Type getElementType() const;
+
+  unsigned getNumElements() 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

 
 //===----------------------------------------------------------------------===//
 // SPIR-V Dialect
 //===----------------------------------------------------------------------===//
 
 SPIRVDialect::SPIRVDialect(MLIRContext *context)
     : Dialect(getDialectNamespace(), context) {
   addTypes<ArrayType, CooperativeMatrixNVType, ImageType, PointerType,
-           RuntimeArrayType, StructType>();
+           RuntimeArrayType, StructType, MatrixType>();
 
   addAttributes<InterfaceVarABIAttr, TargetEnvAttr, VerCapExtAttr>();
 
   // Add SPIR-V ops.
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/SPIRV/SPIRVOps.cpp.inc"
       >();
if (auto t = type.dyn_cast<FloatType>()) {
     parser.emitError(typeLoc, "cannot use ")
         << type << " to compose SPIR-V types";
     return Type();
   }
 
   return type;
 }
 
+static Type parseAndVerifyMatrixType(SPIRVDialect const &dialect,
+                                     DialectAsmParser &parser) {
+  Type type;
+  llvm::SMLoc typeLoc = parser.getCurrentLocation();
+  if (parser.parseType(type))
+    return Type();
+
+  if (auto t = type.dyn_cast<VectorType>()) {
+    if (t.getRank() != 1) {
+      parser.emitError(typeLoc, "only 1-D vector allowed but found ") << t;
+      return Type();
+    }
+    if (t.getNumElements() > 4 || t.getNumElements() < 2) {
+      parser.emitError(typeLoc,
+                       "Matrix columns size has to be less than or equal "
+                       "to 4 and greater than or equal 2, but found ")
+          << t.getNumElements();
+      return Type();
+    }
+
+    if (!t.getElementType().isa<FloatType>()) {
+      parser.emitError(typeLoc, "Matrix columns' elements must be of "
+                                "Float type, got ")
+          << t.getElementType();
+      return Type();
+    }
+  } else {
+    parser.emitError(typeLoc, "Matrix must be composed using vector "
+                              "type, got ")
+        << type;
+    return Type();
+  }
+
+  return type;
+}
+
 /// Parses an optional `, stride = N` assembly segment. If no parsing failure
 /// occurs, writes `N` to `stride` if existing and writes 0 to `stride` if
 /// missing.
 static LogicalResult parseOptionalArrayStride(const SPIRVDialect &dialect,
                                               DialectAsmParser &parser,
                                               unsigned &stride) {
   if (failed(parser.parseOptionalComma())) {
     stride = 0;
static Type parseRuntimeArrayType(SPIRVDialect const &dialect,
   if (failed(parseOptionalArrayStride(dialect, parser, stride)))
     return Type();
 
   if (parser.parseGreater())
     return Type();
   return RuntimeArrayType::get(elementType, stride);
 }
 
+// matrix-type ::= `!spv.matrix` `<` integer-literal `x` element-type `>`
+static Type parseMatrixType(SPIRVDialect const &dialect,
+                            DialectAsmParser &parser) {
+  if (parser.parseLess())
+    return Type();
+
+  SmallVector<int64_t, 1> countDims;
+  llvm::SMLoc countLoc = parser.getCurrentLocation();
+  if (parser.parseDimensionList(countDims, /*allowDynamic=*/false))
+    return Type();
+  if (countDims.size() != 1) {
+    parser.emitError(countLoc, "expected single unsigned "
+                               "integer for number of columns");
+    return Type();
+  }
+
+  int64_t numCols = countDims[0];
+  // According to the specification, Matrices can have 2, 3, or 4 columns
+  if (numCols < 2 || numCols > 4) {
+    parser.emitError(countLoc, "Matrix is expected to have 2, 3, or 4 "
+                               "columns.");
+    return Type();
+  }
+
+  Type elementType = parseAndVerifyMatrixType(dialect, parser);
+  if (!elementType)
+    return Type();
+
+  if (parser.parseGreater())
+    return Type();
+
+  return MatrixType::get(elementType, numCols);
+}
+
 // Specialize this function to parse each of the parameters that define an
 // ImageType. By default it assumes this is an enum type.
 template <typename ValTy>
 static Optional<ValTy> parseAndVerify(SPIRVDialect const &dialect,
                                       DialectAsmParser &parser) {
   StringRef enumSpec;
   llvm::SMLoc enumLoc = parser.getCurrentLocation();
   if (parser.parseKeyword(&enumSpec)) {
Type SPIRVDialect::parseType(DialectAsmParser &parser) const {
   if (keyword == "image")
     return parseImageType(*this, parser);
   if (keyword == "ptr")
     return parsePointerType(*this, parser);
   if (keyword == "rtarray")
     return parseRuntimeArrayType(*this, parser);
   if (keyword == "struct")
     return parseStructType(*this, parser);
-
+  if (keyword == "matrix")
+    return parseMatrixType(*this, parser);
   parser.emitError(parser.getNameLoc(), "unknown SPIR-V type: ") << keyword;
   return Type();
 }
 
 //===----------------------------------------------------------------------===//
 // Type Printing
 //===----------------------------------------------------------------------===//
 

 }
 
 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 << ">";
+}
+
 void SPIRVDialect::printType(Type type, DialectAsmPrinter &os) const {
   switch (type.getKind()) {
   case TypeKind::Array:
     print(type.cast<ArrayType>(), os);
     return;
   case TypeKind::CooperativeMatrix:
     print(type.cast<CooperativeMatrixNVType>(), os);
     return;
   case TypeKind::Pointer:
     print(type.cast<PointerType>(), os);
     return;
   case TypeKind::RuntimeArray:
     print(type.cast<RuntimeArrayType>(), os);
     return;
   case TypeKind::Image:
     print(type.cast<ImageType>(), os);
     return;
   case TypeKind::Struct:
     print(type.cast<StructType>(), os);
     return;
+  case TypeKind::Matrix:
+    print(type.cast<MatrixType>(), os);
+    return;
   default:
     llvm_unreachable("unhandled SPIR-V type");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Attribute Parsing
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

 //===----------------------------------------------------------------------===//
 
 bool CompositeType::classof(Type type) {
   switch (type.getKind()) {
   case TypeKind::Array:
   case TypeKind::CooperativeMatrix:
   case TypeKind::RuntimeArray:
   case TypeKind::Struct:
+  case TypeKind::Matrix:
     return true;
   case StandardTypes::Vector:
     return isValid(type.cast<VectorType>());
   default:
     return false;
   }
 }
 
Type CompositeType::getElementType(unsigned index) const {
   case spirv::TypeKind::CooperativeMatrix:
     return cast<CooperativeMatrixNVType>().getElementType();
   case spirv::TypeKind::RuntimeArray:
     return cast<RuntimeArrayType>().getElementType();
   case spirv::TypeKind::Struct:
     return cast<StructType>().getElementType(index);
   case StandardTypes::Vector:
     return cast<VectorType>().getElementType();
+  case spirv::TypeKind::Matrix:
+    return cast<MatrixType>().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::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();
+  case spirv::TypeKind::Matrix:
+    return cast<MatrixType>().getNumElements();
   default:
     llvm_unreachable("invalid composite type");
   }
 }
 
 bool CompositeType::hasCompileTimeKnownNumElements() const {
   switch (getKind()) {
   case TypeKind::CooperativeMatrix:
case spirv::TypeKind::RuntimeArray:
     break;
   case spirv::TypeKind::Struct:
     cast<StructType>().getExtensions(extensions, storage);
     break;
   case StandardTypes::Vector:
     cast<VectorType>().getElementType().cast<ScalarType>().getExtensions(
         extensions, storage);
     break;
+  case spirv::TypeKind::Matrix:
+    cast<MatrixType>().getExtensions(extensions, storage);
+    break;
   default:
     llvm_unreachable("invalid composite type");
   }
 }
 
 void CompositeType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
case spirv::TypeKind::RuntimeArray:
     break;
   case spirv::TypeKind::Struct:
     cast<StructType>().getCapabilities(capabilities, storage);
     break;
   case StandardTypes::Vector:
     cast<VectorType>().getElementType().cast<ScalarType>().getCapabilities(
         capabilities, storage);
     break;
+  case spirv::TypeKind::Matrix:
+    cast<MatrixType>().getCapabilities(capabilities, storage);
+    break;
   default:
     llvm_unreachable("invalid composite type");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // CooperativeMatrixType
 //===----------------------------------------------------------------------===//
void SPIRVType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
   if (auto scalarType = dyn_cast<ScalarType>()) {
     scalarType.getExtensions(extensions, storage);
   } else if (auto compositeType = dyn_cast<CompositeType>()) {
     compositeType.getExtensions(extensions, storage);
   } else if (auto ptrType = dyn_cast<PointerType>()) {
     ptrType.getExtensions(extensions, storage);
   } else if (auto imageType = dyn_cast<ImageType>()) {
     imageType.getExtensions(extensions, storage);
+  } else if (auto matrixType = dyn_cast<MatrixType>()) {
+    matrixType.getExtensions(extensions, storage);
   } else {
     llvm_unreachable("invalid SPIR-V Type to getExtensions");
   }
 }
 
 void SPIRVType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
   if (auto scalarType = dyn_cast<ScalarType>()) {
     scalarType.getCapabilities(capabilities, storage);
   } else if (auto compositeType = dyn_cast<CompositeType>()) {
     compositeType.getCapabilities(capabilities, storage);
   } else if (auto ptrType = dyn_cast<PointerType>()) {
     ptrType.getCapabilities(capabilities, storage);
   } else if (auto imageType = dyn_cast<ImageType>()) {
     imageType.getCapabilities(capabilities, storage);
+  } else if (auto matrixType = dyn_cast<MatrixType>()) {
+    matrixType.getCapabilities(capabilities, storage);
   } else {
     llvm_unreachable("invalid SPIR-V Type to getCapabilities");
   }
 }
 
 //===----------------------------------------------------------------------===//
 // StructType
 //===----------------------------------------------------------------------===//

 }
 
 void StructType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
   for (Type elementType : getElementTypes())
     elementType.cast<SPIRVType>().getCapabilities(capabilities, storage);
 }
+
+//===----------------------------------------------------------------------===//
+// MatrixType
+//===----------------------------------------------------------------------===//
+
+struct spirv::detail::MatrixTypeStorage : public TypeStorage {
+  MatrixTypeStorage(Type elementTy, uint32_t numCols)
+      : TypeStorage(), elementType(elementTy), numCols(numCols) {}
+
+  Type elementType;
+
+  using KeyTy = std::tuple<Type, uint32_t>;
+
+  static MatrixTypeStorage *construct(TypeStorageAllocator &allocator,
+                                      const KeyTy &key) {
+
+    // Initialize the memory using placement new.
+    return new (allocator.allocate<MatrixTypeStorage>())
+        MatrixTypeStorage(std::get<0>(key), std::get<1>(key));
+  }
+
+  bool operator==(const KeyTy &key) const {
+    return key == KeyTy(elementType, getNumColumns());
+  }
+
+  uint32_t getNumColumns() const { return numCols; }
+
+  const uint32_t numCols;
+};
+
+MatrixType MatrixType::get(Type elementType, uint32_t numCols) {
+  return Base::get(elementType.getContext(), TypeKind::Matrix, elementType,
+                   numCols);
+}
+
+MatrixType MatrixType::getChecked(Type elementType, uint32_t numCols,
+                                  Location location) {
+  return Base::getChecked(location, TypeKind::Matrix, elementType, numCols);
+}
+
+LogicalResult MatrixType::verifyConstructionInvariants(Location loc,
+                                                       Type elementType,
+                                                       uint32_t numCols) {
+  if (numCols < 2 || numCols > 4)
+    return emitError(loc, "Matrix can have 2, 3, or 4 columns only.");
+
+  if (!isValidElementType(elementType))
+    return emitError(loc, "Matrix columns must be vectors of floats.");
+
+  /// The underlying vectors (columns) must be of size 2, 3, or 4
+  ArrayRef<int64_t> colShape = elementType.cast<VectorType>().getShape();
+  if (colShape.size() != 1)
+    return emitError(loc, "Matrix columns must be 1D vectors.");
+
+  if (colShape[0] < 2 || colShape[0] > 4)
+    return emitError(loc, "Matrix columns must be of size 2, 3, or 4.");
+
+  return success();
+}
+
+/// Returns true if the matrix elements are vectors of float elements
+bool MatrixType::isValidElementType(Type elementType) {
+  if (elementType.isa<VectorType>()) {
+    Type vectorElementsType =
+        elementType.dyn_cast<VectorType>().getElementType();
+    if (vectorElementsType.isa<FloatType>())
+      return true;
+  }
+  return false;
+}
+
+Type MatrixType::getElementType() const { return getImpl()->elementType; }
+
+unsigned MatrixType::getNumElements() const {
+  return getImpl()->getNumColumns();
+}
+
+void MatrixType::getExtensions(SPIRVType::ExtensionArrayRefVector &extensions,
+                               Optional<StorageClass> storage) {
+  getElementType().cast<SPIRVType>().getExtensions(extensions, storage);
+}
+
+void MatrixType::getCapabilities(
+    SPIRVType::CapabilityArrayRefVector &capabilities,
+    Optional<StorageClass> storage) {
+  // Add any capabilities associated with the underlying vectors (i.e., columns)
+  getElementType().cast<SPIRVType>().getCapabilities(capabilities, storage);
+
+  // Add the Matrix capability
+  capabilities.push_back(Capability::Matrix);
+}

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

private:
   LogicalResult processCooperativeMatrixType(ArrayRef<uint32_t> operands);
 
   LogicalResult processFunctionType(ArrayRef<uint32_t> operands);
 
   LogicalResult processRuntimeArrayType(ArrayRef<uint32_t> operands);
 
   LogicalResult processStructType(ArrayRef<uint32_t> operands);
 
+  LogicalResult processMatrixType(ArrayRef<uint32_t> operands);
+
   //===--------------------------------------------------------------------===//
   // Constant
   //===--------------------------------------------------------------------===//
 
   /// Processes a SPIR-V Op{|Spec}Constant instruction with the given
   /// `operands`. `isSpec` indicates whether this is a specialization constant.
   LogicalResult processConstant(ArrayRef<uint32_t> operands, bool isSpec);
 
case spirv::Decoration::BuiltIn:
     break;
   case spirv::Decoration::ArrayStride:
     if (words.size() != 3) {
       return emitError(unknownLoc, "OpDecorate with ")
              << decorationName << " needs a single integer literal";
     }
     typeDecorations[words[0]] = words[2];
     break;
+  case spirv::Decoration::NoPerspective:
+  case spirv::Decoration::Flat:
+  case spirv::Decoration::NonWritable:
+  case spirv::Decoration::NonReadable:
   case spirv::Decoration::Block:
   case spirv::Decoration::BufferBlock:
     if (words.size() != 2) {
       return emitError(unknownLoc, "OpDecoration with ")
              << decorationName << "needs a single target <id>";
     }
     // Block decoration does not affect spv.struct type, but is still stored for
     // verification.
     // TODO: Update StructType to contain this information since
     // it is needed for many validation rules.
     decorations[words[0]].set(symbol, opBuilder.getUnitAttr());
     break;
+  case spirv::Decoration::Location:
   case spirv::Decoration::SpecId:
     if (words.size() != 3) {
       return emitError(unknownLoc, "OpDecoration with ")
              << decorationName << "needs a single integer literal";
     }
     decorations[words[0]].set(
         symbol, opBuilder.getI32IntegerAttr(static_cast<int32_t>(words[2])));
     break;
+
   default:
     return emitError(unknownLoc, "unhandled Decoration : '") << decorationName;
   }
   return success();
 }
 
 LogicalResult Deserializer::processMemberDecoration(ArrayRef<uint32_t> words) {
   // The binary layout of OpMemberDecorate is different comparing to OpDecorate
LogicalResult Deserializer::processType(spirv::Opcode opcode,
   case spirv::Opcode::OpTypeCooperativeMatrixNV:
     return processCooperativeMatrixType(operands);
   case spirv::Opcode::OpTypeFunction:
     return processFunctionType(operands);
   case spirv::Opcode::OpTypeRuntimeArray:
     return processRuntimeArrayType(operands);
   case spirv::Opcode::OpTypeStruct:
     return processStructType(operands);
+  case spirv::Opcode::OpTypeMatrix:
+    return processMatrixType(operands);
   default:
     return emitError(unknownLoc, "unhandled type instruction");
   }
   return success();
 }
 
 LogicalResult Deserializer::processArrayType(ArrayRef<uint32_t> operands) {
   if (operands.size() != 3) {
LogicalResult Deserializer::processStructType(ArrayRef<uint32_t> operands) {
   }
   typeMap[operands[0]] =
       spirv::StructType::get(memberTypes, layoutInfo, memberDecorationsInfo);
   // TODO(ravishankarm): Update StructType to have member name as attribute as
   // well.
   return success();
 }
 
+LogicalResult Deserializer::processMatrixType(ArrayRef<uint32_t> operands) {
+  if (operands.size() != 3) {
+    // Three operands are needed: result_id, column_type, and column_count
+    return emitError(unknownLoc, "OpTypeMatrix must have 3 operands"
+                                 " (result_id, column_type, and column_count)");
+  }
+  // Matrix columns must be of vector type
+  Type elementTy = getType(operands[1]);
+  if (!elementTy) {
+    return emitError(unknownLoc,
+                     "OpTypeMatrix references undefined column type.")
+           << operands[1];
+  }
+
+  uint32_t colsCount = operands[2];
+  typeMap[operands[0]] = spirv::MatrixType::get(elementTy, colsCount);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Constant
 //===----------------------------------------------------------------------===//
 
 LogicalResult Deserializer::processConstant(ArrayRef<uint32_t> operands,
                                             bool isSpec) {
   StringRef opname = isSpec ? "OpSpecConstant" : "OpConstant";
 
case spirv::Opcode::OpSourceExtension:
     // TODO: This is debug information embedded in the binary which should be
     // translated into the spv.module.
     return success();
   case spirv::Opcode::OpTypeVoid:
   case spirv::Opcode::OpTypeBool:
   case spirv::Opcode::OpTypeInt:
   case spirv::Opcode::OpTypeFloat:
   case spirv::Opcode::OpTypeVector:
+  case spirv::Opcode::OpTypeMatrix:
   case spirv::Opcode::OpTypeArray:
   case spirv::Opcode::OpTypeFunction:
   case spirv::Opcode::OpTypeRuntimeArray:
   case spirv::Opcode::OpTypeStruct:
   case spirv::Opcode::OpTypePointer:
   case spirv::Opcode::OpTypeCooperativeMatrixNV:
     return processType(opcode, operands);
   case spirv::Opcode::OpConstant:

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

if (!decoration) {
     return emitError(
                loc, "non-argument attributes expected to have snake-case-ified "
                     "decoration name, unhandled attribute with name : ")
            << attrName;
   }
   SmallVector<uint32_t, 1> args;
   switch (decoration.getValue()) {
   case spirv::Decoration::DescriptorSet:
+  case spirv::Decoration::Location:
   case spirv::Decoration::Binding:
     if (auto intAttr = attr.second.dyn_cast<IntegerAttr>()) {
       args.push_back(intAttr.getValue().getZExtValue());
       break;
     }
     return emitError(loc, "expected integer attribute for ") << attrName;
   case spirv::Decoration::BuiltIn:
     if (auto strAttr = attr.second.dyn_cast<StringAttr>()) {
       auto enumVal = spirv::symbolizeBuiltIn(strAttr.getValue());
       if (enumVal) {
         args.push_back(static_cast<uint32_t>(enumVal.getValue()));
         break;
       }
       return emitError(loc, "invalid ")
              << attrName << " attribute " << strAttr.getValue();
     }
     return emitError(loc, "expected string attribute for ") << attrName;
+  case spirv::Decoration::NoPerspective:
+  case spirv::Decoration::Flat:
+  case spirv::Decoration::NonWritable:
+  case spirv::Decoration::NonReadable:
+  case spirv::Decoration::Block:
+  case spirv::Decoration::BufferBlock:
+    if (auto unitAttr = attr.second.dyn_cast<UnitAttr>()) {
+      // For unit attributes, the args list has no values so we do nothing
+      break;
+    }
+    return emitError(loc, "expected unit attribute for ") << attrName;
   default:
     return emitError(loc, "unhandled decoration ") << decorationName;
   }
   return emitDecoration(resultID, decoration.getValue(), args);
 }
 
 LogicalResult Serializer::processName(uint32_t resultID, StringRef name) {
   assert(!name.empty() && "unexpected empty string for OpName");
if (auto cooperativeMatrixType =
     typeEnum = spirv::Opcode::OpTypeCooperativeMatrixNV;
     operands.push_back(elementTypeID);
     operands.push_back(static_cast<uint32_t>(cooperativeMatrixType.getScope()));
     operands.push_back(cooperativeMatrixType.getRows());
     operands.push_back(cooperativeMatrixType.getColumns());
     return success();
   }
 
+  if (auto matrixType = type.dyn_cast<spirv::MatrixType>()) {
+    uint32_t elementTypeID = 0;
+    if (failed(processType(loc, matrixType.getElementType(), elementTypeID))) {
+      return failure();
+    }
+    typeEnum = spirv::Opcode::OpTypeMatrix;
+    operands.push_back(elementTypeID);
+    operands.push_back(matrixType.getNumElements());
+    return success();
+  }
+
   // TODO(ravishankarm) : Handle other types.
   return emitError(loc, "unhandled type in serialization: ") << type;
 }
 
 LogicalResult
 Serializer::prepareFunctionType(Location loc, FunctionType type,
                                 spirv::Opcode &typeEnum,
                                 SmallVectorImpl<uint32_t> &operands) {

mlir/test/Dialect/SPIRV/types.mlir

 // expected-error @+1 {{expected ','}}
 func @missing_scope(!spv.coopmatrix<8x16xi32>) -> ()
 
 // -----
 
 // expected-error @+1 {{expected rows and columns size}}
 func @missing_count(!spv.coopmatrix<8xi32, Subgroup>) -> ()
 
+// -----
+
+//===----------------------------------------------------------------------===//
+// Matrix
+//===----------------------------------------------------------------------===//
+// CHECK: func @matrix_type(!spv.matrix<2 x vector<2xf16>>)
+func @matrix_type(!spv.matrix<2 x vector<2xf16>>) -> ()
+
+// -----
+
+// CHECK: func @matrix_type(!spv.matrix<3 x vector<3xf32>>)
+func @matrix_type(!spv.matrix<3 x vector<3xf32>>) -> ()
+
+// -----
+
+// CHECK: func @matrix_type(!spv.matrix<4 x vector<4xf16>>)
+func @matrix_type(!spv.matrix<4 x vector<4xf16>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix is expected to have 2, 3, or 4 columns.}}
+func @matrix_invalid_size(!spv.matrix<5 x vector<3xf32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix is expected to have 2, 3, or 4 columns.}}
+func @matrix_invalid_size(!spv.matrix<1 x vector<3xf32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix columns size has to be less than or equal to 4 and greater than or equal 2, but found 5}}
+func @matrix_invalid_columns_size(!spv.matrix<3 x vector<5xf32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix columns size has to be less than or equal to 4 and greater than or equal 2, but found 1}}
+func @matrix_invalid_columns_size(!spv.matrix<3 x vector<1xf32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected '<'}}
+func @matrix_invalid_format(!spv.matrix 3 x vector<3xf32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{unbalanced ')' character in pretty dialect name}}
+func @matrix_invalid_format(!spv.matrix< 3 x vector<3xf32>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected 'x' in dimension list}}
+func @matrix_invalid_format(!spv.matrix<2 vector<3xi32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix must be composed using vector type, got 'i32'}}
+func @matrix_invalid_type(!spv.matrix< 3 x i32>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix must be composed using vector type, got '!spv.array<16 x f32>'}}
+func @matrix_invalid_type(!spv.matrix< 3 x !spv.array<16 x f32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix must be composed using vector type, got '!spv.rtarray<i32>'}}
+func @matrix_invalid_type(!spv.matrix< 3 x !spv.rtarray<i32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{Matrix columns' elements must be of Float type, got 'i32'}}
+func @matrix_invalid_type(!spv.matrix<2 x vector<3xi32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected single unsigned integer for number of columns}}
+func @matrix_size_type(!spv.matrix< x vector<3xi32>>) -> ()
+
+// -----
+
+// expected-error @+1 {{expected single unsigned integer for number of columns}}
+func @matrix_size_type(!spv.matrix<2.0 x vector<3xi32>>) -> ()
+
+// -----
 No newline at end of file