Struct MmaDefinitionExpand 

pub struct MmaDefinitionExpand<A: CubeType, B: CubeType, CD: CubeType> {
    pub shape: MatrixShape,
    pub a_type: TypeHandle,
    pub b_type: TypeHandle,
    pub cd_type: TypeHandle,
    pub scales_factor: Option<usize>,
    pub scales_type: Option<TypeHandle>,
    /* private fields */
}

Expand type of MmaDefinition.

Fields

shape: MatrixShape

a_type: TypeHandle

b_type: TypeHandle

cd_type: TypeHandle

scales_factor: Option<usize>

scales_type: Option<TypeHandle>

Implementations

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

pub fn __expand_num_elems_method( &self, scope: &Scope, ident: MatrixIdent, ) -> usize

Number of elements in the matrix

pub fn __expand_elems_per_lane_method( &self, scope: &Scope, 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 __expand_vectors_per_lane_method( &self, scope: &Scope, 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 __expand_vector_layout_method( &self, scope: &Scope, ident: MatrixIdent, ) -> MatrixLayout

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

pub fn __expand_vector_size_method( &self, scope: &Scope, 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 __expand_position_of_nth_method( &self, scope: &Scope, 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.

pub fn __expand_scales_index_method( &self, scope: &Scope, 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.

pub fn __expand_scales_count_method(&self, scope: &Scope) -> usize

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

pub fn __expand_scales_vector_size_method(&self, scope: &Scope) -> VectorSize

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

pub fn __expand_load_matrix_method<E: CubePrimitive, NO: Size>( &self, scope: &Scope, 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_method<E: Scalar, N: Size>( &self, scope: &Scope, row: &<[E] as CubeType>::ExpandType, fragment: &mut <Array<Vector<E, N>> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, )

pub fn __expand_store_matrix_method<E: CubePrimitive, N: Size>( &self, scope: &Scope, 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

pub fn __expand_execute_method<NA: Size, NB: Size, NC: Size>( &self, scope: &Scope, 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_method<NA: Size, NB: Size, NC: Size>( &self, scope: &Scope, 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, )

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

pub fn __expand_new(scope: &Scope, 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 __expand_new_scaled<S: CubePrimitive>( scope: &Scope, 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 __expand_num_elems( scope: &Scope, this: &Self, ident: MatrixIdent, ) -> usize

Number of elements in the matrix

pub fn __expand_elems_per_lane( scope: &Scope, this: &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 __expand_vectors_per_lane( scope: &Scope, this: &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 __expand_vector_layout( scope: &Scope, this: &Self, ident: MatrixIdent, ) -> MatrixLayout

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

pub fn __expand_vector_size( scope: &Scope, this: &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 __expand_position_of_nth( scope: &Scope, this: &Self, 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.

pub fn __expand_scales_index( scope: &Scope, this: &Self, 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.

pub fn __expand_scales_count(scope: &Scope, this: &Self) -> usize

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

pub fn __expand_scales_vector_size(scope: &Scope, this: &Self) -> VectorSize

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

pub fn __expand_load_matrix<E: CubePrimitive, NO: Size>( scope: &Scope, this: &Self, 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, row: &<[E] as CubeType>::ExpandType, fragment: &mut <Array<Vector<E, N>> as CubeType>::ExpandType, ident: MatrixIdent, num_matrices: usize, transpose: bool, )

pub fn __expand_store_matrix<E: CubePrimitive, N: Size>( scope: &Scope, this: &Self, 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

pub fn __expand_execute<NA: Size, NB: Size, NC: Size>( scope: &Scope, this: &Self, 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, 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, )

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

impl<A: CubeType, B: CubeType, CD: CubeType> AsMutExpand for MmaDefinitionExpand<A, B, CD>

fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self

fn __expand_as_mut_method(&mut self, scope: &Scope) -> &mut T

impl<A: CubeType, B: CubeType, CD: CubeType> AsRefExpand for MmaDefinitionExpand<A, B, CD>

fn __expand_ref_method(&self, _scope: &Scope) -> &Self

fn __expand_as_ref_method(&self, scope: &Scope) -> &T

impl<A: CubeType, B: CubeType, CD: CubeType> Clone for MmaDefinitionExpand<A, B, CD>

fn clone(&self) -> Self

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> Copy for MmaDefinitionExpand<A, B, CD>

impl<A: CubeType, B: CubeType, CD: CubeType> CubeDebug for MmaDefinitionExpand<A, B, CD>

fn set_debug_name(&self, scope: &Scope, name: &'static str)

Set the debug name of this type’s expansion. Should do nothing for types that don’t appear at runtime

impl<A: Debug + CubeType, B: Debug + CubeType, CD: Debug + CubeType> Debug for MmaDefinitionExpand<A, B, CD>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<A: CubeType, B: CubeType, CD: CubeType> ExpandTypeClone for MmaDefinitionExpand<A, B, CD>

fn clone_unchecked(&self) -> Self

Unchecked clone that only clones the conceptual runtime value. Should only be used in cases where each copy is used in a mutually exclusive branch (i.e. match, runtime enums). This is intentionally separated from Rust’s Clone semantics and should only be used for the conceptual expand values, never real data. Using two values in the same branch is undefined behaviour.

impl<A: CubeType, B: CubeType, CD: CubeType> IntoExpand for MmaDefinitionExpand<A, B, CD>

type Expand = MmaDefinitionExpand<A, B, CD>

fn into_expand(self, _: &Scope) -> Self::Expand

impl<A: CubeType, B: CubeType, CD: CubeType> IntoMut for MmaDefinitionExpand<A, B, CD>

fn into_mut(self, _scope: &Scope) -> Self

Convert the variable into a potentially new mutable variable in scope, copying if needed.