Struct MmaDefinition 

pub struct MmaDefinition<A: CubeType, B: CubeType, CD: CubeType> { /* private fields */ }

Defines a matrix multiplication operation, including the input and output type, and the shape.

Implementations

impl<A: Scalar, B: Scalar, CD: Scalar> MmaDefinition<A, B, CD>

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::B Shape => (K, N)
• MatrixIdent::Accumulator Shape => (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.

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::B Shape => (K, N)
• MatrixIdent::Accumulator Shape => (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.

pub fn num_elems(&self, ident: MatrixIdent) -> usize

Number of elements in the matrix

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.

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.

pub fn vector_layout(&self, ident: MatrixIdent) -> MatrixLayout

The layout of each vector in this matrix (row major or column major)

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.

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.

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.

pub fn scales_count(&self) -> usize

Number of scales in each vector (not the vector size!). Vector size may include padding bytes.

pub fn scales_vector_size(&self) -> VectorSize

Vector size for the scale factors. May be larger than the total number of scales.

pub fn load_matrix<E: CubePrimitive, NO: Size>( &self, row: &Slice<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: &Slice<E>, fragment: &mut Array<Vector<E, N>>, ident: MatrixIdent, num_matrices: usize, transpose: bool, )

pub fn store_matrix<E: CubePrimitive, N: Size>( &self, row: &mut Slice<E, ReadWrite>, 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

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>>, )

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

pub fn __expand_new( scope: &mut Scope, m: usize, n: usize, k: usize, ) -> <Self as CubeType>::ExpandType

pub fn __expand_new_scaled<S: CubePrimitive>( scope: &mut Scope, m: usize, n: usize, k: usize, scale_factor: usize, ) -> <Self as CubeType>::ExpandType

pub fn __expand_num_elems( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> usize

pub fn __expand_elems_per_lane( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> usize

pub fn __expand_vectors_per_lane( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> usize

pub fn __expand_vector_layout( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> MatrixLayout

pub fn __expand_vector_size( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ident: MatrixIdent, ) -> VectorSize

pub fn __expand_position_of_nth( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, lane_id: <u32 as CubeType>::ExpandType, elem_idx: <u32 as CubeType>::ExpandType, ident: MatrixIdent, ) -> <(u32, u32) as CubeType>::ExpandType

pub fn __expand_scales_index( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, lane_id: <u32 as CubeType>::ExpandType, ident: MatrixIdent, ) -> <u32 as CubeType>::ExpandType

pub fn __expand_scales_count( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ) -> usize

pub fn __expand_scales_vector_size( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, ) -> VectorSize

pub fn __expand_load_matrix<E: CubePrimitive, NO: Size>( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, row: <Slice<E> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, ) -> <Array<Vector<E::Scalar, NO>> as CubeType>::ExpandType

pub fn __expand_load_matrix_inplace<E: Scalar, N: Size>( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, row: <Slice<E> as CubeType>::ExpandType, fragment: <Array<Vector<E, N>> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, ) -> <() as CubeType>::ExpandType

pub fn __expand_store_matrix<E: CubePrimitive, N: Size>( scope: &mut Scope, this: &<Self as CubeType>::ExpandType, row: <Slice<E, ReadWrite> as CubeType>::ExpandType, registers: <Array<Vector<E::Scalar, N>> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, ) -> <() as CubeType>::ExpandType

pub fn __expand_execute<NA: Size, NB: Size, NC: Size>( scope: &mut 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_inplace<NA: Size, NB: Size, NC: Size>( scope: &mut 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, ) -> <() as CubeType>::ExpandType

pub fn __expand_execute_scaled<S: Scalar, NA: Size, NB: Size, NC: Size, NS: Size>( scope: &mut 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

Trait Implementations

impl<A: Clone + CubeType, B: Clone + CubeType, CD: Clone + CubeType> Clone for MmaDefinition<A, B, CD>

fn clone(&self) -> MmaDefinition<A, B, CD>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<A: CubeType, B: CubeType, CD: CubeType> CubeType for MmaDefinition<A, B, CD>

type ExpandType = MmaDefinitionExpand<A, B, CD>

impl<A: Copy + CubeType, B: Copy + CubeType, CD: Copy + CubeType> Copy for MmaDefinition<A, B, CD>