pub struct MmaDefinition<A: CubeType, B: CubeType, CD: CubeType> { /* private fields */ }Expand description
Defines a matrix multiplication operation, including the input and output type, and the shape.
Implementations§
Source§impl<A: Scalar, B: Scalar, CD: Scalar> MmaDefinition<A, B, CD>
impl<A: Scalar, B: Scalar, CD: Scalar> MmaDefinition<A, B, CD>
Sourcepub fn new(m: usize, n: usize, k: usize) -> Self
pub fn new(m: usize, n: usize, k: usize) -> Self
Create a new matrix definition that is going to be used in the manual
matrix-multiply and accumulate execute_manual_mma() function.
You have to declare the shape used for the execution. The shape of the current matrix is determined using the MatrixIdent.
- MatrixIdent::A Shape => (M, K)
MatrixIdent::BShape => (K, N)MatrixIdent::AccumulatorShape => (M, N)
Not all shapes are supported, and the permitted shapes depend on the element type.
Layout for manual MMA is determined by the runtime and must be handled manually.
Use Self::vector_layout to check the correct data layout for each element.
Refer to nvidia documentation.
Sourcepub fn new_scaled<S: CubePrimitive>(
m: usize,
n: usize,
k: usize,
scale_factor: usize,
) -> Self
pub fn new_scaled<S: CubePrimitive>( m: usize, n: usize, k: usize, scale_factor: usize, ) -> Self
Create a new matrix definition that is going to be used in the manual
matrix-multiply and accumulate execute_manual_mma() function.
You have to declare the shape used for the execution. The shape of the current matrix is determined using the MatrixIdent.
- MatrixIdent::A Shape => (M, K)
MatrixIdent::BShape => (K, N)MatrixIdent::AccumulatorShape => (M, N)
Not all shapes are supported, and the permitted shapes depend on the element type.
Layout for manual MMA is determined by the runtime and must be handled manually.
Use Self::vector_layout to check the correct data layout for each element.
Refer to nvidia documentation.
Sourcepub fn num_elems(&self, ident: MatrixIdent) -> usize
pub fn num_elems(&self, ident: MatrixIdent) -> usize
Number of elements in the matrix
Sourcepub fn elems_per_lane(&self, ident: MatrixIdent) -> usize
pub fn elems_per_lane(&self, ident: MatrixIdent) -> usize
Returns the number of elements handled by each lane. Should be packed into Vectors of size
vector_size with Self::vector_layout.
§Note
“Lane” here refers to the unit relative to a plane, to distinguish it from a unit relative to a cube.
Sourcepub fn vectors_per_lane(&self, ident: MatrixIdent) -> usize
pub fn vectors_per_lane(&self, ident: MatrixIdent) -> usize
Returns the number of vectors of size vector_size with layout vector_layout per lane.
§Note
“Lane” here refers to the unit relative to a plane, to distinguish it from a unit relative to a cube.
Sourcepub fn vector_layout(&self, ident: MatrixIdent) -> MatrixLayout
pub fn vector_layout(&self, ident: MatrixIdent) -> MatrixLayout
The layout of each vector in this matrix (row major or column major)
Sourcepub fn vector_size(&self, ident: MatrixIdent) -> VectorSize
pub fn vector_size(&self, ident: MatrixIdent) -> VectorSize
Number of elements in each vector passed to the execute function. Represents the maximum number of contiguous elements held by the thread.
Sourcepub fn position_of_nth(
&self,
lane_id: u32,
elem_idx: u32,
ident: MatrixIdent,
) -> (u32, u32)
pub fn position_of_nth( &self, lane_id: u32, elem_idx: u32, ident: MatrixIdent, ) -> (u32, u32)
Returns the coordinates of the nth element handled by the lane_id
Each lane contains Self::elems_per_lane elements in Self::vector_size chunks.
Returns (row_idx, col_idx)
§Note
“Lane” here refers to the unit relative to a plane, to distinguish it from a unit relative to a cube.
Sourcepub fn scales_index(&self, lane_id: u32, ident: MatrixIdent) -> u32
pub fn scales_index(&self, lane_id: u32, ident: MatrixIdent) -> u32
Index of the scales for this thread, along the non-major dimension of the matrix.
Each thread loads all scales in the major direction into a single Vector.
Sourcepub fn scales_count(&self) -> usize
pub fn scales_count(&self) -> usize
Number of scales in each vector (not the vector size!). Vector size may include padding bytes.
Sourcepub fn scales_vector_size(&self) -> VectorSize
pub fn scales_vector_size(&self) -> VectorSize
Vector size for the scale factors. May be larger than the total number of scales.
Sourcepub fn load_matrix<E: CubePrimitive, NO: Size>(
&self,
row: &[E],
ident: MatrixIdent,
num_matrices: usize,
transpose: bool,
) -> Array<Vector<E::Scalar, NO>> ⓘ
pub fn load_matrix<E: CubePrimitive, NO: Size>( &self, row: &[E], ident: MatrixIdent, num_matrices: usize, transpose: bool, ) -> Array<Vector<E::Scalar, NO>> ⓘ
Load one or more matrix register using intrinsic instructions. CUDA only.
The number of matrices must be 1, 2, or 4. The rows for the nth matrix are passed by the 8
lanes starting at n * 8. All slice starts must be valid, even for non-participating lanes.
The slice determines the starting address for a 16-byte row loaded by this unit, with
the row index being UNIT_POS_PLANE % 8.
The number of elements is determined by element size.
§Constraints:
Address must be aligned to 16 bytes Address must be in shared memory
pub fn load_matrix_inplace<E: Scalar, N: Size>( &self, row: &[E], fragment: &mut Array<Vector<E, N>>, ident: MatrixIdent, num_matrices: usize, transpose: bool, )
Sourcepub fn store_matrix<E: CubePrimitive, N: Size>(
&self,
row: &mut [E],
registers: &Array<Vector<E::Scalar, N>>,
ident: MatrixIdent,
num_matrices: usize,
transpose: bool,
)
pub fn store_matrix<E: CubePrimitive, N: Size>( &self, row: &mut [E], registers: &Array<Vector<E::Scalar, N>>, ident: MatrixIdent, num_matrices: usize, transpose: bool, )
Store one or more matrix register using intrinsic instructions. CUDA only.
The number of matrices must be 1, 2, or 4. The rows for the nth matrix are passed by the 8
lanes starting at n * 8. All slice starts must be valid, even for non-participating lanes.
The slice determines the starting address for a 16-byte row loaded by this unit, with
the row index being UNIT_POS_PLANE % 8.
The number of elements is determined by element size.
§Constraints:
Address must be aligned to 16 bytes Address must be in shared memory
Sourcepub fn execute<NA: Size, NB: Size, NC: Size>(
&self,
registers_a: &Array<Vector<A, NA>>,
registers_b: &Array<Vector<B, NB>>,
registers_c: &Array<Vector<CD, NC>>,
) -> Array<Vector<CD, NC>> ⓘ
pub fn execute<NA: Size, NB: Size, NC: Size>( &self, registers_a: &Array<Vector<A, NA>>, registers_b: &Array<Vector<B, NB>>, registers_c: &Array<Vector<CD, NC>>, ) -> Array<Vector<CD, NC>> ⓘ
Execute a low level mma operation with manually managed registers. Register layout
and index mapping can be retrieved from the MmaDefinition
pub fn execute_inplace<NA: Size, NB: Size, NC: Size>( &self, registers_a: &Array<Vector<A, NA>>, registers_b: &Array<Vector<B, NB>>, registers_c: &mut Array<Vector<CD, NC>>, )
Sourcepub fn execute_scaled<S: Scalar, NA: Size, NB: Size, NC: Size, NS: Size>(
&self,
registers_a: &Array<Vector<A, NA>>,
registers_b: &Array<Vector<B, NB>>,
registers_c: &Array<Vector<CD, NC>>,
scales_a: Vector<S, NS>,
scales_b: Vector<S, NS>,
) -> Array<Vector<CD, NC>> ⓘ
pub fn execute_scaled<S: Scalar, NA: Size, NB: Size, NC: Size, NS: Size>( &self, registers_a: &Array<Vector<A, NA>>, registers_b: &Array<Vector<B, NB>>, registers_c: &Array<Vector<CD, NC>>, scales_a: Vector<S, NS>, scales_b: Vector<S, NS>, ) -> Array<Vector<CD, NC>> ⓘ
Execute a low level block scaled mma operation with manually managed registers. Register
layout and index mapping can be retrieved from the MmaDefinition
Sourcepub fn __expand_new(
scope: &Scope,
m: usize,
n: usize,
k: usize,
) -> <Self as CubeType>::ExpandType
pub fn __expand_new( scope: &Scope, m: usize, n: usize, k: usize, ) -> <Self as CubeType>::ExpandType
Create a new matrix definition that is going to be used in the manual
matrix-multiply and accumulate execute_manual_mma() function.
You have to declare the shape used for the execution. The shape of the current matrix is determined using the MatrixIdent.
- MatrixIdent::A Shape => (M, K)
MatrixIdent::BShape => (K, N)MatrixIdent::AccumulatorShape => (M, N)
Not all shapes are supported, and the permitted shapes depend on the element type.
Layout for manual MMA is determined by the runtime and must be handled manually.
Use Self::vector_layout to check the correct data layout for each element.
Refer to nvidia documentation.
Sourcepub fn __expand_new_scaled<S: CubePrimitive>(
scope: &Scope,
m: usize,
n: usize,
k: usize,
scale_factor: usize,
) -> <Self as CubeType>::ExpandType
pub fn __expand_new_scaled<S: CubePrimitive>( scope: &Scope, m: usize, n: usize, k: usize, scale_factor: usize, ) -> <Self as CubeType>::ExpandType
Create a new matrix definition that is going to be used in the manual
matrix-multiply and accumulate execute_manual_mma() function.
You have to declare the shape used for the execution. The shape of the current matrix is determined using the MatrixIdent.
- MatrixIdent::A Shape => (M, K)
MatrixIdent::BShape => (K, N)MatrixIdent::AccumulatorShape => (M, N)
Not all shapes are supported, and the permitted shapes depend on the element type.
Layout for manual MMA is determined by the runtime and must be handled manually.
Use Self::vector_layout to check the correct data layout for each element.
Refer to nvidia documentation.
Sourcepub fn __expand_num_elems(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
ident: MatrixIdent,
) -> usize
pub fn __expand_num_elems( scope: &Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> usize
Number of elements in the matrix
Sourcepub fn __expand_elems_per_lane(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
ident: MatrixIdent,
) -> usize
pub fn __expand_elems_per_lane( scope: &Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> usize
Returns the number of elements handled by each lane. Should be packed into Vectors of size
vector_size with Self::vector_layout.
§Note
“Lane” here refers to the unit relative to a plane, to distinguish it from a unit relative to a cube.
Sourcepub fn __expand_vectors_per_lane(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
ident: MatrixIdent,
) -> usize
pub fn __expand_vectors_per_lane( scope: &Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> usize
Returns the number of vectors of size vector_size with layout vector_layout per lane.
§Note
“Lane” here refers to the unit relative to a plane, to distinguish it from a unit relative to a cube.
Sourcepub fn __expand_vector_layout(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
ident: MatrixIdent,
) -> MatrixLayout
pub fn __expand_vector_layout( scope: &Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> MatrixLayout
The layout of each vector in this matrix (row major or column major)
Sourcepub fn __expand_vector_size(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
ident: MatrixIdent,
) -> VectorSize
pub fn __expand_vector_size( scope: &Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> VectorSize
Number of elements in each vector passed to the execute function. Represents the maximum number of contiguous elements held by the thread.
Sourcepub fn __expand_position_of_nth(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
lane_id: <u32 as CubeType>::ExpandType,
elem_idx: <u32 as CubeType>::ExpandType,
ident: MatrixIdent,
) -> (<u32 as CubeType>::ExpandType, <u32 as CubeType>::ExpandType)
pub fn __expand_position_of_nth( scope: &Scope, this: &<Self as CubeType>::ExpandType, lane_id: <u32 as CubeType>::ExpandType, elem_idx: <u32 as CubeType>::ExpandType, ident: MatrixIdent, ) -> (<u32 as CubeType>::ExpandType, <u32 as CubeType>::ExpandType)
Returns the coordinates of the nth element handled by the lane_id
Each lane contains Self::elems_per_lane elements in Self::vector_size chunks.
Returns (row_idx, col_idx)
§Note
“Lane” here refers to the unit relative to a plane, to distinguish it from a unit relative to a cube.
Sourcepub fn __expand_scales_index(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
lane_id: <u32 as CubeType>::ExpandType,
ident: MatrixIdent,
) -> <u32 as CubeType>::ExpandType
pub fn __expand_scales_index( scope: &Scope, this: &<Self as CubeType>::ExpandType, lane_id: <u32 as CubeType>::ExpandType, ident: MatrixIdent, ) -> <u32 as CubeType>::ExpandType
Index of the scales for this thread, along the non-major dimension of the matrix.
Each thread loads all scales in the major direction into a single Vector.
Sourcepub fn __expand_scales_count(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
) -> usize
pub fn __expand_scales_count( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> usize
Number of scales in each vector (not the vector size!). Vector size may include padding bytes.
Sourcepub fn __expand_scales_vector_size(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
) -> VectorSize
pub fn __expand_scales_vector_size( scope: &Scope, this: &<Self as CubeType>::ExpandType, ) -> VectorSize
Vector size for the scale factors. May be larger than the total number of scales.
Sourcepub fn __expand_load_matrix<E: CubePrimitive, NO: Size>(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
row: &<[E] as CubeType>::ExpandType,
ident: MatrixIdent,
num_matrices: usize,
transpose: bool,
) -> <Array<Vector<E::Scalar, NO>> as CubeType>::ExpandType ⓘ
pub fn __expand_load_matrix<E: CubePrimitive, NO: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, row: &<[E] as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, ) -> <Array<Vector<E::Scalar, NO>> as CubeType>::ExpandType ⓘ
Load one or more matrix register using intrinsic instructions. CUDA only.
The number of matrices must be 1, 2, or 4. The rows for the nth matrix are passed by the 8
lanes starting at n * 8. All slice starts must be valid, even for non-participating lanes.
The slice determines the starting address for a 16-byte row loaded by this unit, with
the row index being UNIT_POS_PLANE % 8.
The number of elements is determined by element size.
§Constraints:
Address must be aligned to 16 bytes Address must be in shared memory
pub fn __expand_load_matrix_inplace<E: Scalar, N: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, row: &<[E] as CubeType>::ExpandType, fragment: &mut <Array<Vector<E, N>> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, )
Sourcepub fn __expand_store_matrix<E: CubePrimitive, N: Size>(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
row: &mut <[E] as CubeType>::ExpandType,
registers: &<Array<Vector<E::Scalar, N>> as CubeType>::ExpandType,
ident: MatrixIdent,
num_matrices: usize,
transpose: bool,
)
pub fn __expand_store_matrix<E: CubePrimitive, N: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, row: &mut <[E] as CubeType>::ExpandType, registers: &<Array<Vector<E::Scalar, N>> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, )
Store one or more matrix register using intrinsic instructions. CUDA only.
The number of matrices must be 1, 2, or 4. The rows for the nth matrix are passed by the 8
lanes starting at n * 8. All slice starts must be valid, even for non-participating lanes.
The slice determines the starting address for a 16-byte row loaded by this unit, with
the row index being UNIT_POS_PLANE % 8.
The number of elements is determined by element size.
§Constraints:
Address must be aligned to 16 bytes Address must be in shared memory
Sourcepub fn __expand_execute<NA: Size, NB: Size, NC: Size>(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
registers_a: &<Array<Vector<A, NA>> as CubeType>::ExpandType,
registers_b: &<Array<Vector<B, NB>> as CubeType>::ExpandType,
registers_c: &<Array<Vector<CD, NC>> as CubeType>::ExpandType,
) -> <Array<Vector<CD, NC>> as CubeType>::ExpandType ⓘ
pub fn __expand_execute<NA: Size, NB: Size, NC: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, registers_a: &<Array<Vector<A, NA>> as CubeType>::ExpandType, registers_b: &<Array<Vector<B, NB>> as CubeType>::ExpandType, registers_c: &<Array<Vector<CD, NC>> as CubeType>::ExpandType, ) -> <Array<Vector<CD, NC>> as CubeType>::ExpandType ⓘ
Execute a low level mma operation with manually managed registers. Register layout
and index mapping can be retrieved from the MmaDefinition
pub fn __expand_execute_inplace<NA: Size, NB: Size, NC: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, registers_a: &<Array<Vector<A, NA>> as CubeType>::ExpandType, registers_b: &<Array<Vector<B, NB>> as CubeType>::ExpandType, registers_c: &mut <Array<Vector<CD, NC>> as CubeType>::ExpandType, )
Sourcepub fn __expand_execute_scaled<S: Scalar, NA: Size, NB: Size, NC: Size, NS: Size>(
scope: &Scope,
this: &<Self as CubeType>::ExpandType,
registers_a: &<Array<Vector<A, NA>> as CubeType>::ExpandType,
registers_b: &<Array<Vector<B, NB>> as CubeType>::ExpandType,
registers_c: &<Array<Vector<CD, NC>> as CubeType>::ExpandType,
scales_a: <Vector<S, NS> as CubeType>::ExpandType,
scales_b: <Vector<S, NS> as CubeType>::ExpandType,
) -> <Array<Vector<CD, NC>> as CubeType>::ExpandType ⓘ
pub fn __expand_execute_scaled<S: Scalar, NA: Size, NB: Size, NC: Size, NS: Size>( scope: &Scope, this: &<Self as CubeType>::ExpandType, registers_a: &<Array<Vector<A, NA>> as CubeType>::ExpandType, registers_b: &<Array<Vector<B, NB>> as CubeType>::ExpandType, registers_c: &<Array<Vector<CD, NC>> as CubeType>::ExpandType, scales_a: <Vector<S, NS> as CubeType>::ExpandType, scales_b: <Vector<S, NS> as CubeType>::ExpandType, ) -> <Array<Vector<CD, NC>> as CubeType>::ExpandType ⓘ
Execute a low level block scaled mma operation with manually managed registers. Register
layout and index mapping can be retrieved from the MmaDefinition
Trait Implementations§
Source§impl<A: Clone + CubeType, B: Clone + CubeType, CD: Clone + CubeType> Clone for MmaDefinition<A, B, CD>
impl<A: Clone + CubeType, B: Clone + CubeType, CD: Clone + CubeType> Clone for MmaDefinition<A, B, CD>
impl<A: Copy + CubeType, B: Copy + CubeType, CD: Copy + CubeType> Copy for MmaDefinition<A, B, CD>
Source§impl<A: CubeType, B: CubeType, CD: CubeType> CubeType for MmaDefinition<A, B, CD>
impl<A: CubeType, B: CubeType, CD: CubeType> CubeType for MmaDefinition<A, B, CD>
type ExpandType = MmaDefinitionExpand<A, B, CD>
Auto Trait Implementations§
impl<A, B, CD> Freeze for MmaDefinition<A, B, CD>
impl<A, B, CD> RefUnwindSafe for MmaDefinition<A, B, CD>
impl<A, B, CD> Send for MmaDefinition<A, B, CD>
impl<A, B, CD> Sync for MmaDefinition<A, B, CD>
impl<A, B, CD> Unpin for MmaDefinition<A, B, CD>
impl<A, B, CD> UnsafeUnpin for MmaDefinition<A, B, CD>
impl<A, B, CD> UnwindSafe for MmaDefinition<A, B, CD>
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
Source§impl<T> CloneExpand for Twhere
T: Clone,
impl<T> CloneExpand for Twhere
T: Clone,
fn __expand_clone_method(&self, _: &Scope) -> T
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> Downcast for Twhere
T: Any,
impl<T> Downcast for Twhere
T: Any,
Source§fn into_any(self: Box<T>) -> Box<dyn Any>
fn into_any(self: Box<T>) -> Box<dyn Any>
Box<dyn Trait> (where Trait: Downcast) to Box<dyn Any>, which can then be
downcast into Box<dyn ConcreteType> where ConcreteType implements Trait.Source§fn into_any_rc(self: Rc<T>) -> Rc<dyn Any>
fn into_any_rc(self: Rc<T>) -> Rc<dyn Any>
Rc<Trait> (where Trait: Downcast) to Rc<Any>, which can then be further
downcast into Rc<ConcreteType> where ConcreteType implements Trait.Source§fn as_any(&self) -> &(dyn Any + 'static)
fn as_any(&self) -> &(dyn Any + 'static)
&Trait (where Trait: Downcast) to &Any. This is needed since Rust cannot
generate &Any’s vtable from &Trait’s.Source§fn as_any_mut(&mut self) -> &mut (dyn Any + 'static)
fn as_any_mut(&mut self) -> &mut (dyn Any + 'static)
&mut Trait (where Trait: Downcast) to &Any. This is needed since Rust cannot
generate &mut Any’s vtable from &mut Trait’s.Source§impl<T> DowncastSend for T
impl<T> DowncastSend for T
Source§impl<T> DowncastSync for T
impl<T> DowncastSync for T
impl<T> ErasedDestructor for Twhere
T: 'static,
Source§impl<T> Instrument for T
impl<T> Instrument for T
Source§fn instrument(self, span: Span) -> Instrumented<Self> ⓘ
fn instrument(self, span: Span) -> Instrumented<Self> ⓘ
Source§fn in_current_span(self) -> Instrumented<Self> ⓘ
fn in_current_span(self) -> Instrumented<Self> ⓘ
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read more