Skip to content

[mlir][spirv] Add definition for VectorTimesMatrixOp #124571

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 29, 2025
+144 −16
Merged

[mlir][spirv] Add definition for VectorTimesMatrixOp #124571

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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