Skip to content

Commit

Permalink
[mlir][spirv] Add definition for VectorTimesMatrixOp (llvm#124571)
Browse files Browse the repository at this point in the history 
Adding op as defined in section 3.52.13. (Arithmetic Instructions) of
the SPIR-V specification.
  • Loading branch information
@IgWod-IMG
IgWod-IMG authored Jan 29, 2025
1 parent f226cab commit a01097f
Show file tree
Hide file tree
Showing 5 changed files with 144 additions and 16 deletions.
16 changes: 13 additions & 3 deletions mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4234,8 +4234,13 @@ class SPIRV_CoopMatrixOfType<list<Type> allowedTypes> :
"::llvm::cast<::mlir::spirv::CooperativeMatrixType>($_self).getElementType()",
"Cooperative Matrix">;

class SPIRV_MatrixOfType<list<Type> allowedTypes> :
ContainerType<AnyTypeOf<allowedTypes>, SPIRV_IsMatrixType,
"::llvm::cast<::mlir::spirv::MatrixType>($_self).getElementType()",
"Matrix">;

class SPIRV_VectorOf<Type type> :
VectorOfLengthAndType<[2, 3, 4, 8,16], [type]>;
VectorOfLengthAndType<[2, 3, 4, 8, 16], [type]>;

class SPIRV_ScalarOrVectorOf<Type type> :
AnyTypeOf<[type, SPIRV_VectorOf<type>]>;
Expand All @@ -4248,6 +4253,9 @@ class SPIRV_MatrixOrCoopMatrixOf<Type type> :
AnyTypeOf<[SPIRV_AnyMatrix,
SPIRV_CoopMatrixOfType<[type]>]>;

class SPIRV_MatrixOf<Type type> :
SPIRV_MatrixOfType<[type]>;

def SPIRV_ScalarOrVector : AnyTypeOf<[SPIRV_Scalar, SPIRV_Vector]>;
def SPIRV_ScalarOrVectorOrPtr : AnyTypeOf<[SPIRV_ScalarOrVector, SPIRV_AnyPtr]>;

Expand Down Expand Up @@ -4387,7 +4395,8 @@ def SPIRV_OC_OpFRem : I32EnumAttrCase<"OpFRem", 140>;
def SPIRV_OC_OpFMod : I32EnumAttrCase<"OpFMod", 141>;
def SPIRV_OC_OpVectorTimesScalar : I32EnumAttrCase<"OpVectorTimesScalar", 142>;
def SPIRV_OC_OpMatrixTimesScalar : I32EnumAttrCase<"OpMatrixTimesScalar", 143>;
def SPIRV_OC_OpMatrixTimesVector : I32EnumAttrCase<"OpMatrixTimesVector", 145>;
def SPIRV_OC_OpVectorTimesMatrix : I32EnumAttrCase<"OpVectorTimesMatrix", 144>;
def SPIRV_OC_OpMatrixTimesVector : I32EnumAttrCase<"OpMatrixTimesVector", 145>;
def SPIRV_OC_OpMatrixTimesMatrix : I32EnumAttrCase<"OpMatrixTimesMatrix", 146>;
def SPIRV_OC_OpDot : I32EnumAttrCase<"OpDot", 148>;
def SPIRV_OC_OpIAddCarry : I32EnumAttrCase<"OpIAddCarry", 149>;
Expand Down Expand Up @@ -4559,7 +4568,8 @@ def SPIRV_OpcodeAttr :
SPIRV_OC_OpFSub, SPIRV_OC_OpIMul, SPIRV_OC_OpFMul, SPIRV_OC_OpUDiv,
SPIRV_OC_OpSDiv, SPIRV_OC_OpFDiv, SPIRV_OC_OpUMod, SPIRV_OC_OpSRem,
SPIRV_OC_OpSMod, SPIRV_OC_OpFRem, SPIRV_OC_OpFMod,
SPIRV_OC_OpVectorTimesScalar, SPIRV_OC_OpMatrixTimesScalar, SPIRV_OC_OpMatrixTimesVector,
SPIRV_OC_OpVectorTimesScalar, SPIRV_OC_OpMatrixTimesScalar,
SPIRV_OC_OpVectorTimesMatrix, SPIRV_OC_OpMatrixTimesVector,
SPIRV_OC_OpMatrixTimesMatrix, SPIRV_OC_OpDot, SPIRV_OC_OpIAddCarry,
SPIRV_OC_OpISubBorrow, SPIRV_OC_OpUMulExtended, SPIRV_OC_OpSMulExtended,
SPIRV_OC_OpIsNan, SPIRV_OC_OpIsInf, SPIRV_OC_OpOrdered, SPIRV_OC_OpUnordered,
Expand Down
65 changes: 58 additions & 7 deletions mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMatrixOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,7 @@ def SPIRV_MatrixTimesMatrixOp : SPIRV_Op<"MatrixTimesMatrix", [Pure]> {

// -----

def SPIRV_MatrixTimesScalarOp : SPIRV_Op<
"MatrixTimesScalar", [Pure, AllTypesMatch<["matrix", "result"]>]> {
def SPIRV_MatrixTimesScalarOp : SPIRV_Op<"MatrixTimesScalar", [Pure, AllTypesMatch<["matrix", "result"]>]> {
let summary = "Scale a floating-point matrix.";

let description = [{
Expand Down Expand Up @@ -114,8 +113,11 @@ def SPIRV_MatrixTimesScalarOp : SPIRV_Op<

// -----

def SPIRV_MatrixTimesVectorOp : SPIRV_Op<"MatrixTimesVector", [Pure]> {
let summary = "Linear-algebraic multiply of matrix X vector.";
def SPIRV_MatrixTimesVectorOp : SPIRV_Op<"MatrixTimesVector", [
Pure,
AllElementTypesMatch<["vector", "result"]>
]> {
let summary = "Linear-algebraic Matrix X Vector.";

let description = [{
Result Type must be a vector of floating-point type.
Expand All @@ -140,12 +142,12 @@ def SPIRV_MatrixTimesVectorOp : SPIRV_Op<"MatrixTimesVector", [Pure]> {
];

let arguments = (ins
SPIRV_AnyMatrix:$matrix,
SPIRV_AnyVector:$vector
SPIRV_MatrixOf<SPIRV_Float>:$matrix,
SPIRV_VectorOf<SPIRV_Float>:$vector
);

let results = (outs
SPIRV_AnyVector:$result
SPIRV_VectorOf<SPIRV_Float>:$result
);

let assemblyFormat = [{
Expand Down Expand Up @@ -198,4 +200,53 @@ def SPIRV_TransposeOp : SPIRV_Op<"Transpose", [Pure]> {

// -----

def SPIRV_VectorTimesMatrixOp : SPIRV_Op<"VectorTimesMatrix", [
Pure,
AllElementTypesMatch<["vector", "result"]>
]> {
let summary = "Linear-algebraic Vector X Matrix.";

let description = [{
Result Type must be a vector of floating-point type.

Vector must be a vector with the same Component Type as the Component
Type in Result Type. Its number of components must equal the number of
components in each column in Matrix.

Matrix 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
components in Result Type.

<!-- End of AutoGen section -->

#### Example:

```mlir
%result = spirv.VectorTimesMatrix %vector, %matrix : vector<4xf32>, !spirv.matrix<4 x vector<4xf32>> -> vector<4xf32>
```
}];

let availability = [
MinVersion<SPIRV_V_1_0>,
MaxVersion<SPIRV_V_1_6>,
Extension<[]>,
Capability<[SPIRV_C_Matrix]>
];

let arguments = (ins
SPIRV_VectorOf<SPIRV_Float>:$vector,
SPIRV_MatrixOf<SPIRV_Float>:$matrix
);

let results = (outs
SPIRV_VectorOf<SPIRV_Float>:$result
);

let assemblyFormat = [{
operands attr-dict `:` type($vector) `,` type($matrix) `->` type($result)
}];
}

// -----

#endif // MLIR_DIALECT_SPIRV_IR_MATRIX_OPS
30 changes: 26 additions & 4 deletions mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1717,10 +1717,32 @@ LogicalResult spirv::MatrixTimesVectorOp::verify() {
<< resultType.getNumElements() << ") must match the matrix rows ("
<< matrixType.getNumRows() << ")";

auto matrixElementType = matrixType.getElementType();
if (matrixElementType != vectorType.getElementType() ||
matrixElementType != resultType.getElementType())
return emitOpError("matrix, vector, and result element types must match");
if (matrixType.getElementType() != resultType.getElementType())
return emitOpError("matrix and result element types must match");

return success();
}

//===----------------------------------------------------------------------===//
// spirv.VectorTimesMatrix
//===----------------------------------------------------------------------===//

LogicalResult spirv::VectorTimesMatrixOp::verify() {
auto vectorType = llvm::cast<VectorType>(getVector().getType());
auto matrixType = llvm::cast<spirv::MatrixType>(getMatrix().getType());
auto resultType = llvm::cast<VectorType>(getType());

if (matrixType.getNumRows() != vectorType.getNumElements())
return emitOpError("number of components in vector must equal the number "
"of components in each column in matrix");

if (resultType.getNumElements() != matrixType.getNumColumns())
return emitOpError("number of columns in matrix must equal the number of "
"components in result");

if (matrixType.getElementType() != resultType.getElementType())
return emitOpError("matrix must be a matrix with the same component type "
"as the component type in result");

return success();
}
Expand Down
42 changes: 40 additions & 2 deletions mlir/test/Dialect/SPIRV/IR/matrix-ops.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,13 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
spirv.ReturnValue %result : vector<4xf32>
}

// CHECK-LABEL: @vector_times_matrix_1
spirv.func @vector_times_matrix_1(%arg0: vector<3xf32>, %arg1: !spirv.matrix<4 x vector<3xf32>>) -> vector<4xf32> "None" {
// CHECK: {{%.*}} = spirv.VectorTimesMatrix {{%.*}}, {{%.*}} : vector<3xf32>, !spirv.matrix<4 x vector<3xf32>> -> vector<4xf32>
%result = spirv.VectorTimesMatrix %arg0, %arg1 : vector<3xf32>, !spirv.matrix<4 x vector<3xf32>> -> vector<4xf32>
spirv.ReturnValue %result : vector<4xf32>
}

// CHECK-LABEL: @matrix_times_matrix_1
spirv.func @matrix_times_matrix_1(%arg0: !spirv.matrix<3 x vector<3xf32>>, %arg1: !spirv.matrix<3 x vector<3xf32>>) -> !spirv.matrix<3 x vector<3xf32>> "None"{
// CHECK: {{%.*}} = spirv.MatrixTimesMatrix {{%.*}}, {{%.*}} : !spirv.matrix<3 x vector<3xf32>>, !spirv.matrix<3 x vector<3xf32>> -> !spirv.matrix<3 x vector<3xf32>>
Expand Down Expand Up @@ -123,7 +130,6 @@ func.func @matrix_times_matrix_component_type_mismatch_1(%arg0 : !spirv.matrix<3
return
}


// -----

func.func @matrix_times_matrix_component_type_mismatch_2(%arg0 : !spirv.matrix<3 x vector<3xf64>>, %arg1 : !spirv.matrix<3x vector<3xf32>>){
Expand All @@ -135,7 +141,7 @@ func.func @matrix_times_matrix_component_type_mismatch_2(%arg0 : !spirv.matrix<3
// -----

func.func @matrix_times_vector_element_type_mismatch(%arg0: !spirv.matrix<4 x vector<3xf32>>, %arg1: vector<4xf16>) {
// expected-error @+1 {{matrix, vector, and result element types must match}}
// expected-error @+1 {{op failed to verify that all of {vector, result} have same element type}}
%result = spirv.MatrixTimesVector %arg0, %arg1 : !spirv.matrix<4 x vector<3xf32>>, vector<4xf16> -> vector<3xf32>
return
}
Expand All @@ -155,3 +161,35 @@ func.func @matrix_times_vector_column_mismatch(%arg0: !spirv.matrix<4 x vector<3
%result = spirv.MatrixTimesVector %arg0, %arg1 : !spirv.matrix<4 x vector<3xf32>>, vector<3xf32> -> vector<3xf32>
return
}

// -----

func.func @vector_times_matrix_vector_matrix_mismatch(%arg0: vector<4xf32>, %arg1: !spirv.matrix<4 x vector<3xf32>>) {
// expected-error @+1 {{number of components in vector must equal the number of components in each column in matrix}}
%result = spirv.VectorTimesMatrix %arg0, %arg1 : vector<4xf32>, !spirv.matrix<4 x vector<3xf32>> -> vector<3xf32>
return
}

// -----

func.func @vector_times_matrix_result_matrix_mismatch(%arg0: vector<3xf32>, %arg1: !spirv.matrix<4 x vector<3xf32>>) {
// expected-error @+1 {{number of columns in matrix must equal the number of components in result}}
%result = spirv.VectorTimesMatrix %arg0, %arg1 : vector<3xf32>, !spirv.matrix<4 x vector<3xf32>> -> vector<3xf32>
return
}

// -----

func.func @vector_times_matrix_vector_type_mismatch(%arg0: vector<3xf16>, %arg1: !spirv.matrix<4 x vector<3xf32>>) {
// expected-error @+1 {{op failed to verify that all of {vector, result} have same element type}}
%result = spirv.VectorTimesMatrix %arg0, %arg1 : vector<3xf16>, !spirv.matrix<4 x vector<3xf32>> -> vector<4xf32>
return
}

// -----

func.func @vector_times_matrix_matrix_type_mismatch(%arg0: vector<3xf32>, %arg1: !spirv.matrix<4 x vector<3xf16>>) {
// expected-error @+1 {{matrix must be a matrix with the same component type as the component type in result}}
%result = spirv.VectorTimesMatrix %arg0, %arg1 : vector<3xf32>, !spirv.matrix<4 x vector<3xf16>> -> vector<4xf32>
return
}
7 changes: 7 additions & 0 deletions mlir/test/Target/SPIRV/matrix.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,13 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
%result = spirv.MatrixTimesVector %arg0, %arg1 : !spirv.matrix<3 x vector<4xf32>>, vector<3xf32> -> vector<4xf32>
spirv.ReturnValue %result : vector<4xf32>
}

// CHECK-LABEL: @vector_times_matrix_1
spirv.func @vector_times_matrix_1(%arg0: vector<3xf32>, %arg1: !spirv.matrix<4 x vector<3xf32>>) -> vector<4xf32> "None" {
// CHECK: {{%.*}} = spirv.VectorTimesMatrix {{%.*}}, {{%.*}} : vector<3xf32>, !spirv.matrix<4 x vector<3xf32>> -> vector<4xf32>
%result = spirv.VectorTimesMatrix %arg0, %arg1 : vector<3xf32>, !spirv.matrix<4 x vector<3xf32>> -> vector<4xf32>
spirv.ReturnValue %result : vector<4xf32>
}

// CHECK-LABEL: @matrix_times_matrix_1
spirv.func @matrix_times_matrix_1(%arg0: !spirv.matrix<3 x vector<3xf32>>, %arg1: !spirv.matrix<3 x vector<3xf32>>) -> !spirv.matrix<3 x vector<3xf32>> "None"{
Expand Down

0 comments on commit a01097f

Please sign in to comment.