350 lines
13 KiB
TableGen
350 lines
13 KiB
TableGen
//===- SPIRVCooperativeMatrixOps.td - cooperative matmul ---*- 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 is the op definition spec of cooperative matrix multiply extension ops.
|
|
// We support both cooperative matrix extensions:
|
|
// - SPV_NV_cooperative_matrix
|
|
// - SPV_KHR_cooperative_matrix
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
|
|
#define MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
|
|
|
|
include "mlir/Interfaces/AlignmentAttrInterface.td"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// SPV_KHR_cooperative_matrix extension ops.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// -----
|
|
|
|
def SPIRV_KHRCooperativeMatrixLengthOp :
|
|
SPIRV_KhrVendorOp<"CooperativeMatrixLength", [Pure]> {
|
|
let summary = "Queries the number of cooperative matrix components";
|
|
|
|
let description = [{
|
|
Number of components of a cooperative matrix type accessible to each
|
|
invocation when treated as a composite.
|
|
|
|
The type attribute must be a cooperative matrix type.
|
|
|
|
#### Example:
|
|
|
|
```
|
|
%0 = spirv.KHR.CooperativeMatrixLength :
|
|
!spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>
|
|
```
|
|
}];
|
|
|
|
let assemblyFormat = "attr-dict `:` $cooperative_matrix_type";
|
|
|
|
let availability = [
|
|
MinVersion<SPIRV_V_1_6>,
|
|
MaxVersion<SPIRV_V_1_6>,
|
|
Extension<[SPV_KHR_cooperative_matrix]>,
|
|
Capability<[SPIRV_C_CooperativeMatrixKHR]>
|
|
];
|
|
|
|
let arguments = (ins
|
|
TypeAttrOf<SPIRV_AnyCooperativeMatrix>:$cooperative_matrix_type
|
|
);
|
|
|
|
let results = (outs
|
|
SPIRV_Int32:$result
|
|
);
|
|
|
|
let hasVerifier = false;
|
|
}
|
|
|
|
// -----
|
|
|
|
def SPIRV_KHRCooperativeMatrixLoadOp : SPIRV_KhrVendorOp<"CooperativeMatrixLoad", [DeclareOpInterfaceMethods<AlignmentAttrOpInterface>]> {
|
|
let summary = "Loads a cooperative matrix through a pointer";
|
|
|
|
let description = [{
|
|
Load a cooperative matrix through a pointer.
|
|
|
|
Result Type is the type of the loaded object. It must be a cooperative
|
|
matrix type.
|
|
|
|
Pointer is a pointer. Its type must be an OpTypePointer whose Type operand is
|
|
a scalar or vector type. If the Shader capability was declared, Pointer must
|
|
point into an array and any ArrayStride decoration on Pointer is ignored.
|
|
|
|
MemoryLayout specifies how matrix elements are laid out in memory. It must
|
|
come from a 32-bit integer constant instruction whose value corresponds to a
|
|
Cooperative Matrix Layout. See the Cooperative Matrix Layout table for a
|
|
description of the layouts and detailed layout-specific rules.
|
|
|
|
Stride further qualifies how matrix elements are laid out in memory. It must
|
|
be a scalar integer type and its exact semantics depend on MemoryLayout.
|
|
|
|
Memory Operand must be a Memory Operand literal. If not present, it is the
|
|
same as specifying None.
|
|
|
|
NOTE: In earlier versions of the SPIR-V spec, 'Memory Operand' was known
|
|
as 'Memory Access'.
|
|
|
|
For a given dynamic instance of this instruction, all operands of this
|
|
instruction must be the same for all invocations in a given scope instance
|
|
(where the scope is the scope the cooperative matrix type was created with).
|
|
All invocations in a given scope instance must be active or all must be
|
|
inactive.
|
|
|
|
TODO: In the SPIR-V spec, `stride` is an optional argument. We should also
|
|
support this optionality in the SPIR-V dialect.
|
|
|
|
#### Example:
|
|
|
|
```
|
|
%0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor>
|
|
: !spirv.ptr<i32, StorageBuffer>, i32
|
|
-> !spirv.KHR.coopmatrix<16x8xi32, Workgroup, MatrixA>
|
|
|
|
%1 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor>, <Volatile>
|
|
: !spirv.ptr<f32, StorageBuffer>, i64
|
|
-> !spirv.KHR.coopmatrix<8x8xf32, Subgroup, MatrixAcc>
|
|
```
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$pointer `,` $stride `,` $matrix_layout ( `,` $memory_operand^ )? ( `,` $alignment^ )? attr-dict `:`
|
|
type(operands) `->` type($result)
|
|
}];
|
|
|
|
let availability = [
|
|
MinVersion<SPIRV_V_1_6>,
|
|
MaxVersion<SPIRV_V_1_6>,
|
|
Extension<[SPV_KHR_cooperative_matrix]>,
|
|
Capability<[SPIRV_C_CooperativeMatrixKHR]>
|
|
];
|
|
|
|
// TODO: Add scope operand for MakePointer*. See #145485.
|
|
let arguments = (ins
|
|
SPIRV_AnyPtr:$pointer,
|
|
SPIRV_KHR_CooperativeMatrixLayoutAttr:$matrix_layout,
|
|
SPIRV_Integer:$stride,
|
|
OptionalAttr<SPIRV_MemoryAccessAttr>:$memory_operand,
|
|
OptionalAttr<IntValidAlignment<I32Attr>>:$alignment
|
|
);
|
|
|
|
let results = (outs
|
|
SPIRV_AnyCooperativeMatrix:$result
|
|
);
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Type":$result, "Value":$pointer,
|
|
"spirv::ConstantOp":$stride,
|
|
"spirv::CooperativeMatrixLayoutKHR":$layout), [{
|
|
build($_builder, $_state, result, pointer, layout, stride,
|
|
spirv::MemoryAccessAttr{}, IntegerAttr{});
|
|
}]>
|
|
];
|
|
}
|
|
|
|
// -----
|
|
|
|
def SPIRV_KHRCooperativeMatrixStoreOp : SPIRV_KhrVendorOp<"CooperativeMatrixStore", [DeclareOpInterfaceMethods<AlignmentAttrOpInterface>]> {
|
|
let summary = "Stores a cooperative matrix through a pointer";
|
|
|
|
let description = [{
|
|
Store a cooperative matrix through a pointer.
|
|
Pointer is a pointer. Its type must be an OpTypePointer whose Type operand
|
|
is a scalar or vector type. If the Shader capability was declared, Pointer
|
|
must point into an array and any ArrayStride decoration on Pointer is
|
|
ignored.
|
|
|
|
Object is the object to store. Its type must be an
|
|
OpTypeCooperativeMatrixKHR.
|
|
|
|
MemoryLayout specifies how matrix elements are laid out in memory. It must
|
|
come from a 32-bit integer constant instruction whose value corresponds to a
|
|
Cooperative Matrix Layout. See the Cooperative Matrix Layout table for a
|
|
description of the layouts and detailed layout-specific rules.
|
|
|
|
Stride further qualifies how matrix elements are laid out in memory. It must
|
|
be a scalar integer type and its exact semantics depend on MemoryLayout.
|
|
|
|
Memory Operand must be a Memory Operand literal. If not present, it is the
|
|
same as specifying None.
|
|
|
|
NOTE: In earlier versions of the SPIR-V spec, 'Memory Operand' was known
|
|
as 'Memory Access'.
|
|
|
|
For a given dynamic instance of this instruction, all operands of this
|
|
instruction must be the same for all invocations in a given scope instance
|
|
(where the scope is the scope the cooperative matrix type was created with).
|
|
All invocations in a given scope instance must be active or all must be
|
|
inactive.
|
|
|
|
TODO: In the SPIR-V spec, `stride` is an optional argument. We should also
|
|
support this optionality in the SPIR-V dialect.
|
|
|
|
#### Example:
|
|
|
|
```
|
|
spirv.KHR.CooperativeMatrixStore %ptr, %obj, %stride, <RowMajor> :
|
|
!spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i32
|
|
|
|
spirv.KHR.CooperativeMatrixStore %ptr, %obj, %stride, <ColumnMajor>, <Volatile> :
|
|
!spirv.ptr<f32, StorageBuffer>, !spirv.coopmatrix<8x8xf32, Subgroup, MatrixAcc>, i64
|
|
```
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$pointer `,` $object `,` $stride `,` $matrix_layout ( `,` $memory_operand^ )? ( `,` $alignment^ )? attr-dict `:`
|
|
type(operands)
|
|
}];
|
|
|
|
let availability = [
|
|
MinVersion<SPIRV_V_1_6>,
|
|
MaxVersion<SPIRV_V_1_6>,
|
|
Extension<[SPV_KHR_cooperative_matrix]>,
|
|
Capability<[SPIRV_C_CooperativeMatrixKHR]>
|
|
];
|
|
|
|
// TODO: Add scope operand for MakePointer*. See #145485.
|
|
let arguments = (ins
|
|
SPIRV_AnyPtr:$pointer,
|
|
SPIRV_AnyCooperativeMatrix:$object,
|
|
SPIRV_KHR_CooperativeMatrixLayoutAttr:$matrix_layout,
|
|
SPIRV_Integer:$stride,
|
|
OptionalAttr<SPIRV_MemoryAccessAttr>:$memory_operand,
|
|
OptionalAttr<IntValidAlignment<I32Attr>>:$alignment
|
|
);
|
|
|
|
let results = (outs);
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Value":$pointer, "Value":$object,
|
|
"spirv::ConstantOp":$stride,
|
|
"spirv::CooperativeMatrixLayoutKHR":$layout), [{
|
|
build($_builder, $_state, pointer, object, layout, stride,
|
|
spirv::MemoryAccessAttr{}, IntegerAttr{});
|
|
}]>
|
|
];
|
|
}
|
|
|
|
// -----
|
|
|
|
def SPIRV_KHRCooperativeMatrixMulAddOp : SPIRV_KhrVendorOp<"CooperativeMatrixMulAdd",
|
|
[Pure, AllTypesMatch<["c", "result"]>]> {
|
|
let summary = "Returns the result of `(A x B) + C` of matrices A, B, and C";
|
|
|
|
let description = [{
|
|
Linear-algebraic matrix multiply of A by B and then component-wise add C.
|
|
The order of the operations is implementation-dependent. The internal
|
|
precision of floating-point operations is defined by the client API. Integer
|
|
operations used in the multiplication of A by B are performed at the
|
|
precision of the Result Type and the resulting value will equal the
|
|
low-order N bits of the correct result R, where N is the result width and R
|
|
is computed with enough precision to avoid overflow and underflow if the
|
|
SaturatingAccumulation Cooperative Matrix Operand is not present. If the
|
|
SaturatingAccumulation Cooperative Matrix Operand is present and overflow or
|
|
underflow occurs as part of calculating that intermediate result, the result
|
|
of the instruction is undefined. Integer additions of the elements of that
|
|
intermediate result with those of C are performed at the precision of Result
|
|
Type, are exact, and are saturating if the SaturatingAccumulation
|
|
Cooperative Matrix Operand is present, with the signedness of the saturation
|
|
being that of the components of Result Type. If the SaturatingAccumulation
|
|
Cooperative Matrix Operand is not present then the resulting value will
|
|
equal the low-order N bits of the correct result R, where N is the result
|
|
width and R is computed with enough precision to avoid overflow and
|
|
underflow.
|
|
|
|
Result Type must be a cooperative matrix type with M rows and N columns
|
|
whose Use must be MatrixAccumulatorKHR.
|
|
|
|
A is a cooperative matrix with M rows and K columns whose Use must be
|
|
MatrixAKHR.
|
|
|
|
B is a cooperative matrix with K rows and N columns whose Use must be
|
|
MatrixBKHR.
|
|
|
|
C is a cooperative matrix with M rows and N columns whose Use must be
|
|
MatrixAccumulatorKHR.
|
|
|
|
The values of M, N, and K must be consistent across the result and operands.
|
|
This is referred to as an MxNxK matrix multiply.
|
|
|
|
A, B, C, and Result Type must have the same scope, and this defines the
|
|
scope of the operation. A, B, C, and Result Type need not necessarily have
|
|
the same component type, this is defined by the client API.
|
|
|
|
If the Component Type of any matrix operand is an integer type, then its
|
|
components are treated as signed if the Matrix{A,B,C,Result}SignedComponents
|
|
Cooperative Matrix Operand is present and are treated as unsigned otherwise.
|
|
|
|
Cooperative Matrix Operands is an optional Cooperative Matrix Operand
|
|
literal. If not present, it is the same as specifying the Cooperative Matrix
|
|
Operand None.
|
|
|
|
For a given dynamic instance of this instruction, all invocations in a given
|
|
scope instance must be active or all must be inactive (where the scope is
|
|
the scope of the operation).
|
|
|
|
``` {.ebnf}
|
|
cooperative-matrixmuladd-op ::= ssa-id `=` `spirv.KHR.CooperativeMatrixMulAdd`
|
|
ssa-use `,` ssa-use `,` ssa-use
|
|
(`<` matrix-operands `>`)? `:`
|
|
a-cooperative-matrix-type `,`
|
|
b-cooperative-matrix-type `->`
|
|
result-cooperative-matrix-type
|
|
```
|
|
|
|
#### Example:
|
|
|
|
```
|
|
%0 = spirv.KHR.CooperativeMatrixMulAdd %matA, %matB, %matC :
|
|
!spirv.coopmatrix<4x4xf32, Subgroup, MatrixA>,
|
|
!spirv.coopmatrix<4x4xf32, Subgroup, MatrixB> ->
|
|
!spirv.coopmatrix<4x4xf32, Subgroup, MatrixAcc>
|
|
|
|
%1 = spirv.KHR.CooperativeMatrixMulAdd %matA, %matB, %matC, <ASigned | AccSat> :
|
|
!spirv.coopmatrix<8x16xi32, Subgroup, MatrixA>,
|
|
!spirv.coopmatrix<16x4xi32, Subgroup, MatrixB> ->
|
|
!spirv.coopmatrix<8x4xi32, Subgroup, MatrixAcc>
|
|
```
|
|
}];
|
|
|
|
let assemblyFormat = [{
|
|
$a `,` $b `,` $c ( `,` $matrix_operands^ )? attr-dict `:`
|
|
type($a) `,` type($b) `->` type($c)
|
|
}];
|
|
|
|
let availability = [
|
|
MinVersion<SPIRV_V_1_6>,
|
|
MaxVersion<SPIRV_V_1_6>,
|
|
Extension<[SPV_KHR_cooperative_matrix]>,
|
|
Capability<[SPIRV_C_CooperativeMatrixKHR]>
|
|
];
|
|
|
|
let arguments = (ins
|
|
SPIRV_AnyCooperativeMatrix:$a,
|
|
SPIRV_AnyCooperativeMatrix:$b,
|
|
SPIRV_AnyCooperativeMatrix:$c,
|
|
OptionalAttr<SPIRV_KHR_CooperativeMatrixOperandsAttr>:$matrix_operands
|
|
);
|
|
|
|
let results = (outs
|
|
SPIRV_AnyCooperativeMatrix:$result
|
|
);
|
|
|
|
let builders = [
|
|
OpBuilder<(ins "Value":$a, "Value":$b, "Value":$c), [{
|
|
build($_builder, $_state, a, b, c,
|
|
spirv::CooperativeMatrixOperandsKHRAttr{});
|
|
}]>
|
|
];
|
|
}
|
|
|
|
// -----
|
|
|
|
#endif // MLIR_DIALECT_SPIRV_IR_COOPERATIVE_MATRIX_OPS
|