add support for Cooperative matrix#589
Conversation
56b4228 to
5df8670
Compare
|
@Firestar99 test is broken not sure why |
Firestar99
left a comment
There was a problem hiding this comment.
I always thought it requires SPV_EXT_physical_storage_buffer but apparently it only interacts with it, so we can actually implement this without needing it.
| #[spirv_std_macros::gpu_only] | ||
| #[doc(alias = "OpCooperativeMatrixLoadKHR")] | ||
| #[inline] | ||
| pub unsafe fn load(ptr: *const T, layout: u32, stride: u32) -> Self { |
There was a problem hiding this comment.
Why have ptr: *const T be a pointer when spec says it must be a "pointer to an array" aka a slice?
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.
| /// | ||
| /// # Safety | ||
| /// `ptr` must be valid for `ROWS * stride` (row-major) or `COLS * stride` | ||
| /// (column-major) element reads within the subgroup's access pattern. |
There was a problem hiding this comment.
This should probably be in the safety contract of every function:
All the operands to this instruction must be dynamically uniform within every instance of the Scope of the cooperative matrix.
To be honest, I would just copy the description of the spec and adjust it to be more rustdoc-like, as it's the most accurate description you'll get for intrinsics. Maybe also mention that scope is always Subgroup.
| /// Matrix role: input operand A in D = A × B + C. | ||
| pub const MATRIX_A: u32 = 0; | ||
| /// Matrix role: input operand B in D = A × B + C. | ||
| pub const MATRIX_B: u32 = 1; | ||
| /// Matrix role: accumulator / result in D = A × B + C. | ||
| pub const MATRIX_ACCUMULATOR: u32 = 2; | ||
|
|
||
| /// Memory layout: rows are stored contiguously. | ||
| pub const ROW_MAJOR: u32 = 0; | ||
| /// Memory layout: columns are stored contiguously. | ||
| pub const COLUMN_MAJOR: u32 = 1; |
There was a problem hiding this comment.
Why not an enum? Especially for RAW_MAJOR, since it isn't used in generics but only in function calls. For MATRIX_A I could see u32 aliases, eg.
#[repr(u32)]
enum MatrixIDK {
MatrixA,
MatrixB,
MatrixAccumulator,
}
pub const MATRIX_A: u32 = MatrixIDK::MatrixA as u32;
| //! Requires the `CooperativeMatrixKHR` capability and `SPV_KHR_cooperative_matrix` extension: | ||
| //! ```text | ||
| //! -C target-feature=+CooperativeMatrixKHR,+ext:SPV_KHR_cooperative_matrix | ||
| //! ``` |
There was a problem hiding this comment.
A link to the spec is always nice: https://github.khronos.org/SPIRV-Registry/extensions/KHR/SPV_KHR_cooperative_matrix.html
| "NonPrivatePointerKHR", | ||
| MemoryAccess::NON_PRIVATE_POINTER, | ||
| ), | ||
| ("NonPrivatePointerKHR", MemoryAccess::NON_PRIVATE_POINTER), |
There was a problem hiding this comment.
Have you cargo fmt-ed the first commit on it's own? These changes look like some of that ended up in the second commit
| let component_type = match args.type_at(0).kind() { | ||
| TyKind::Float(FloatTy::F32) => SpirvType::Float(32).def(span, cx), | ||
| TyKind::Float(FloatTy::F64) => SpirvType::Float(64).def(span, cx), | ||
| TyKind::Int(IntTy::I8) => SpirvType::Integer(8, true).def(span, cx), | ||
| TyKind::Int(IntTy::I16) => SpirvType::Integer(16, true).def(span, cx), | ||
| TyKind::Int(IntTy::I32) => SpirvType::Integer(32, true).def(span, cx), | ||
| TyKind::Int(IntTy::I64) => SpirvType::Integer(64, true).def(span, cx), | ||
| TyKind::Uint(UintTy::U8) => SpirvType::Integer(8, false).def(span, cx), | ||
| TyKind::Uint(UintTy::U16) => SpirvType::Integer(16, false).def(span, cx), | ||
| TyKind::Uint(UintTy::U32) => SpirvType::Integer(32, false).def(span, cx), | ||
| TyKind::Uint(UintTy::U64) => SpirvType::Integer(64, false).def(span, cx), |
There was a problem hiding this comment.
Recursively calling trans_intrinsic_type?
| /// Memory layout: rows are stored contiguously. | ||
| pub const ROW_MAJOR: MatrixLayout = MatrixLayout::RowMajor; | ||
| /// Memory layout: columns are stored contiguously. | ||
| pub const COLUMN_MAJOR: MatrixLayout = MatrixLayout::ColumnMajor; |
There was a problem hiding this comment.
we could eliminate these but we'll need #![feature(adt_const_params)]
|
@Firestar99 i modified the code according to the suggestions |
this adds support for Cooperative matrix, i had to upgrade
rspirvfrom 0.12->0.13 to have access to coop-mat features.I made it two separate commits, so it'll be easier to review
this only supports
mul_addfor now