[mlir][spirv] Add MatrixTimesScalar operation

[Hazem]

Summary

• Define the MatrixTimesScalar operation and add roundtrip tests.
• Added a new base class for matrix-specific operations to avoid invalid operands type mismatch check.
• Created a separate Matrix arithmetic operations td file to add more operations in the future.
• Augmented the automatically generated verify method to print more fine-grained error messages.
• Made minor Updates to the matrix type tests.

[Hazem]

• remove unintentionally added include line

[antiagainst]

Thanks @Hazem! A few nits.

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

3001

Could you keep this alphabetically sorted?

3029

Same here.

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

2

Nit: align to 80 characters

19

Could you keep each line wrapped to 80 characters?

23

This isn't correct actually? I think we need two different operand types because for example for OpMatrixTimesScalar the second operand should be a scalar type? This is actually not needed given in the following you override with explicit arguments and results. We can just remove this and derive from SPV_Op.

47

Could you keep each line wrapped to 80 characters?

55

Same here

70

Actually we can have just one matrix type given that the input and result are the same and the scalar's type can be deduced from it. If you'd like to handle it later a TOOD here would be nice.

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

11

This should match with the function name.

18

Same here.

29

Why adding this? Seems unrelated..

[Hazem]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

23

The comment about similar operands was a copy paste error. Thanks for SPV_Op hint, I thought we need to group arithmetic ops under SPV_ArithmeticOps class by any means. Fixed it now, and all subsequent matrix ops will follow the same pattern if possible.

70

I am still trying to understand how the Op declarative form work, so let's say I change this line to only have the input Matrix type, then I get the following error:

"error: type of operand #1, named 'scalar', is not buildable and a buildable type cannot be inferred"

The same happens for result/matrix as well. I looked in the documentation and my understanding is that I need to inherit from BuildableType class for these variables types. Is that correct? Because otherwise I already set a type constraint for each of them as ins/outs.

I couldn't find SPIRV examples using assembly format definition (except cooperative matrix), and the documentation is somewhat lacking diverse examples. I could use a custom printer/parser but I wanted to leverage the auto-gen feature as it is more concise, representative, and easier to update, if needed, in the future.

Please if you have hints/docs on what needs to be done to get this right send me and I will work on it. I believe this will be beneficial as I plan to add the rest of the matrix-related Ops after this patch.

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

29

This and the next two are minor re-naming of matrix tests, should they be separated in a different patch?

[Hazem]

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

29

Removed them as unneeded.

[antiagainst]

Hey @Hazem , did you forget to push new changes? I'm still seeing the old commits.

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

Documentation for the Declarative Assembly Format is at here: https://mlir.llvm.org/docs/OpDefinitions/#declarative-assembly-format. It's a bit opaque as you said, similar like other declarative mechanisms because the goal is to be a high-level description so by nature it is inevitably terse. :) At the moment a good way to understand is to search the codebase and find other examples.

I don't think you can express the relationship that $scalar is of a type that is the same as the element type of $matrix. So pretty much it means manually defining a parser/printer.. So if you'd like to avoid that for now, it's fine for me with just a TODO here so we can remind ourselves later. :)

[rriddle]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

If the type of $scalar is the element type of something else, you can use a trait for that. TypesMatchWith in this case:

https://github.com/llvm/llvm-project/blob/13331477c0d1aeb8f3c9f24b3d0487bc6fcaa225/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td#L1028

You can find many other usages in that file.

[Hazem]

• address revision comments

[Hazem]

[Edit] I believe this patch will crash the build if D81763 is committed to master before it, due to accesschain tests format change. So once D81763 is committed, I can update this patch to reflect the format changes before committing it.

In D81677#2093864, @antiagainst wrote:

Hey @Hazem , did you forget to push new changes? I'm still seeing the old commits.

I was reluctant to add a TODO versus actually implementing the change. I went for a TODO for now. Might revisit it later. I pushed the changes now. Sorry for the confusion.

[Hazem]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

The $scalar is the element type of an element type.

For example:
Scalar is f16 when

float16 is the element type of Vector, and this Vector is the element type of a Matrix. So, one more level of nesting.

[rriddle]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

The level of nesting doesn't matter. All that matters is that you can transform one type to get the other. As long as there is a translation from A->B that doesn't require external information, you can represent it using TypesMatchWith.

[Hazem]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

I see. Is it a good practice in this case to do something like that in the transformation:

"$_self.cast<MatrixType>().getElementType().cast<VectorType>().getElementType()"

If yes, wouldn't this cause issues if the matrix element type is not a valid vector type?

[rriddle]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

It's fine as long as you ensure that the matrix type is verified to have a vector element type before hand(i.e. via the type constraint). This is good practice in general though. The type constraint for $matrix(or whatever the matrix type is) should be as complete as possible. The declarative parser will use the type constraint to verify the type before applying the transformation.

[Hazem]

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

Makes sense. Thank you!

[antiagainst]

LGTM!

mlir/include/mlir/Dialect/SPIRV/SPIRVMatrixOps.td

70

Yes. What River said. Sorry I didn't track all the nice progress there so I wasn't aware it's possible now. :)

[Hazem]

In D81677#2099415, @antiagainst wrote:

LGTM!

Excellent, I was granted commit access today, should I go ahead and land this one?

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

 
 class SignlessOrUnsignedIntOfWidths<list<int> widths> :
     AnyTypeOf<!foreach(w, widths, IOrUI<w>),
               StrJoinInt<widths, "/">.result # "-bit signless/unsigned integer">;
 
 def SPV_IsArrayType : CPred<"$_self.isa<::mlir::spirv::ArrayType>()">;
 def SPV_IsCooperativeMatrixType :
   CPred<"$_self.isa<::mlir::spirv::CooperativeMatrixNVType>()">;
+def SPV_IsMatrixType : CPred<"$_self.isa<::mlir::spirv::MatrixType>()">;
 def SPV_IsPtrType : CPred<"$_self.isa<::mlir::spirv::PointerType>()">;
 def SPV_IsRTArrayType : CPred<"$_self.isa<::mlir::spirv::RuntimeArrayType>()">;
 def SPV_IsStructType : CPred<"$_self.isa<::mlir::spirv::StructType>()">;
 
+
 // See https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html#_types
 // for the definition of the following types and type categories.
 
 def SPV_Void : TypeAlias<NoneType, "void">;
 def SPV_Bool : TypeAlias<I1, "bool">;
 def SPV_Integer : AnyIntOfWidths<[8, 16, 32, 64]>;
 def SPV_Int32 : TypeAlias<I32, "Int32">;
 def SPV_Float : FloatOfWidths<[16, 32, 64]>;
 def SPV_Float16or32 : FloatOfWidths<[16, 32]>;
 def SPV_Vector : VectorOfLengthAndType<[2, 3, 4],
                                        [SPV_Bool, SPV_Integer, SPV_Float]>;
 // Component type check is done in the type parser for the following SPIR-V
 // dialect-specific types so we use "Any" here.
 def SPV_AnyPtr : DialectType<SPIRV_Dialect, SPV_IsPtrType,
                              "any SPIR-V pointer type">;
 def SPV_AnyArray : DialectType<SPIRV_Dialect, SPV_IsArrayType,
                                "any SPIR-V array type">;
 def SPV_AnyCooperativeMatrix : DialectType<SPIRV_Dialect,
                                SPV_IsCooperativeMatrixType,
                                "any SPIR-V cooperative matrix type">;
+def SPV_AnyMatrix : DialectType<SPIRV_Dialect, SPV_IsMatrixType,
+                                "any SPIR-V matrix type">;
 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,
-               SPV_AnyCooperativeMatrix]>;
+               SPV_AnyCooperativeMatrix, SPV_AnyMatrix]>;
 def SPV_Type : AnyTypeOf<[
     SPV_Void, SPV_Bool, SPV_Integer, SPV_Float, SPV_Vector,
     SPV_AnyPtr, SPV_AnyArray, SPV_AnyRTArray, SPV_AnyStruct,
-    SPV_AnyCooperativeMatrix
+    SPV_AnyCooperativeMatrix, SPV_AnyMatrix
   ]>;
 
 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">;

 def SPV_OC_OpUDiv                      : I32EnumAttrCase<"OpUDiv", 134>;
 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_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_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,
-      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_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,

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
+
+// -----
+
+def SPV_MatrixTimesScalarOp : SPV_Op<"MatrixTimesScalar", []> {
+  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.
+
+    Scalar must have the same type as the Component Type in Result Type.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    matrix-times-scalar-op ::= ssa-id `=` `spv.MatrixTimesScalar` ssa-use,
+    ssa-use `:` matrix-type `,` float-type `->` matrix-type
+
+    ```
+
+    #### Example:
+
+    ```mlir
+
+    %0 = spv.MatrixTimesScalar %matrix, %scalar :
+    !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<3 x vector<3xf32>>
+
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_AnyMatrix:$matrix,
+    SPV_Float:$scalar
+  );
+
+  let results = (outs
+    SPV_AnyMatrix:$result
+  );
+
+  // TODO (Hazem): we need just one matrix type given that the input and result
+  // are the same and the scalar's type can be deduced from it.
+  let assemblyFormat = [{
+    operands attr-dict `:` type($matrix) `,` type($scalar) `->` type($result)
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_Matrix]>
+  ];
+
+  let verifier = [{ return verifyMatrixTimesScalar(*this); }];
+}
+
+// -----
+
+#endif // SPIRV_MATRIX_OPS
 No newline at end of file

mlir/include/mlir/Dialect/SPIRV/SPIRVOps.td

 include "mlir/Dialect/SPIRV/SPIRVBitOps.td"
 include "mlir/Dialect/SPIRV/SPIRVCastOps.td"
 include "mlir/Dialect/SPIRV/SPIRVCompositeOps.td"
 include "mlir/Dialect/SPIRV/SPIRVControlFlowOps.td"
 include "mlir/Dialect/SPIRV/SPIRVCooperativeMatrixOps.td"
 include "mlir/Dialect/SPIRV/SPIRVGLSLOps.td"
 include "mlir/Dialect/SPIRV/SPIRVGroupOps.td"
 include "mlir/Dialect/SPIRV/SPIRVLogicalOps.td"
+include "mlir/Dialect/SPIRV/SPIRVMatrixOps.td"
 include "mlir/Dialect/SPIRV/SPIRVNonUniformOps.td"
 include "mlir/Dialect/SPIRV/SPIRVStructureOps.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // -----
 
 def SPV_AccessChainOp : SPV_Op<"AccessChain", [NoSideEffect]> {
   let summary = [{

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

if (typeR.getScope() != typeA.getScope() ||
     return op.emitOpError("matrix scope must match");
   if (typeR.getElementType() != typeA.getElementType() ||
       typeR.getElementType() != typeB.getElementType() ||
       typeR.getElementType() != typeC.getElementType())
     return op.emitOpError("matrix element type must match");
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// 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");
+
+    // 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");
+    }
+  }
+  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/test/Dialect/SPIRV/Serialization/matrix.mlir

 // RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
-  spv.func @matrix_type(%arg0 : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, StorageBuffer>, %arg1 : i32) "None" {
-    // CHECK: {{%.*}} = spv.AccessChain {{%.*}}[{{%.*}}] : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, StorageBuffer>
-    %2 = spv.AccessChain %arg0[%arg1] : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, StorageBuffer>
-    spv.Return
+  // CHECK-LABEL: @matrix_access_chain
+  spv.func @matrix_access_chain(%arg0 : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, Function>, %arg1 : i32) -> !spv.ptr<vector<3xf32>, Function> "None" {
+    // CHECK: {{%.*}} = spv.AccessChain {{%.*}}[{{%.*}}] : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, Function>
+    %0 = spv.AccessChain %arg0[%arg1] : !spv.ptr<!spv.matrix<3 x vector<3xf32>>, Function>
+    spv.ReturnValue %0 : !spv.ptr<vector<3xf32>, Function>
+  }
+
+  // CHECK-LABEL: @matrix_times_scalar_1
+  spv.func @matrix_times_scalar_1(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : f32) -> !spv.matrix<3 x vector<3xf32>> "None" {
+    // CHECK: {{%.*}} = spv.MatrixTimesScalar {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<3 x vector<3xf32>>
+    %result = spv.MatrixTimesScalar %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<3 x vector<3xf32>>
+    spv.ReturnValue %result : !spv.matrix<3 x vector<3xf32>>
+  }
+
+  // CHECK-LABEL: @matrix_times_scalar_2
+  spv.func @matrix_times_scalar_2(%arg0 : !spv.matrix<3 x vector<3xf16>>, %arg1 : f16) -> !spv.matrix<3 x vector<3xf16>> "None" {
+    // CHECK: {{%.*}} = spv.MatrixTimesScalar {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<3xf16>>, f16 -> !spv.matrix<3 x vector<3xf16>>
+    %result = spv.MatrixTimesScalar %arg0, %arg1 : !spv.matrix<3 x vector<3xf16>>, f16 -> !spv.matrix<3 x vector<3xf16>>
+    spv.ReturnValue %result : !spv.matrix<3 x vector<3xf16>>
   }
 }
 
 // -----
 
 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

+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -verify-diagnostics %s | FileCheck %s
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  // CHECK-LABEL: @matrix_times_scalar
+  spv.func @matrix_times_scalar_1(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : f32) -> !spv.matrix<3 x vector<3xf32>> "None" {
+    // CHECK: {{%.*}} = spv.MatrixTimesScalar {{%.*}}, {{%.*}} : !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<3 x vector<3xf32>>
+    %result = spv.MatrixTimesScalar %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<3 x vector<3xf32>>
+    spv.ReturnValue %result : !spv.matrix<3 x vector<3xf32>>
+  }
+}
+
+// -----
+
+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 @input_type_mismatch(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : f64) -> () {
+  // 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>>, f64 -> !spv.matrix<3 x vector<3xf32>>
+}
+
+// -----
+
+func @input_output_component_type_mismatch(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : f32) -> () {
+   // expected-error @+1 {{input and result matrices' columns must have the same component type}}
+   %result = spv.MatrixTimesScalar %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<3 x vector<3xf64>>
+}
+
+// -----
+
+func @input_output_size_mismatch(%arg0 : !spv.matrix<3 x vector<3xf32>>, %arg1 : f32) -> () {
+   // expected-error @+1 {{input and result matrices must have the same number of columns}}
+   %result = spv.MatrixTimesScalar %arg0, %arg1 : !spv.matrix<3 x vector<3xf32>>, f32 -> !spv.matrix<4 x vector<3xf32>>
+}
+
+
+