Enum CubeCountInput

pub enum CubeCountInput {
    FromProblem,
    SmFirst {
        m_cubes: u32,
        n_cubes: u32,
        batch_cubes: u32,
    },
    CubeFirst {
        m_cubes: u32,
        n_cubes: u32,
        batch_cubes: u32,
    },
    Flattened {
        m_cubes: u32,
        n_cubes: u32,
    },
    Spread {
        m_cubes: u32,
        n_cubes: u32,
        batch_cubes: u32,
    },
}

CubeCountPlan stripped of non-essential runtime information

This enum is given as runtime input to the matmul

Variants

FromProblem

SmFirst

Fields

m_cubes: u32

n_cubes: u32

batch_cubes: u32

CubeFirst

Fields

m_cubes: u32

n_cubes: u32

batch_cubes: u32

Flattened

Fields

m_cubes: u32

n_cubes: u32

Spread

Fields

m_cubes: u32

n_cubes: u32

batch_cubes: u32

Implementations

impl CubeCountInput

pub fn new_FromProblem() -> Self

pub fn __expand_new_FromProblem(_: &mut Scope) -> CubeCountInputExpand

pub fn new_SmFirst(m_cubes: u32, n_cubes: u32, batch_cubes: u32) -> Self

pub fn __expand_new_SmFirst( _: &mut Scope, m_cubes: <u32 as CubeType>::ExpandType, n_cubes: <u32 as CubeType>::ExpandType, batch_cubes: <u32 as CubeType>::ExpandType, ) -> CubeCountInputExpand

pub fn new_CubeFirst(m_cubes: u32, n_cubes: u32, batch_cubes: u32) -> Self

pub fn __expand_new_CubeFirst( _: &mut Scope, m_cubes: <u32 as CubeType>::ExpandType, n_cubes: <u32 as CubeType>::ExpandType, batch_cubes: <u32 as CubeType>::ExpandType, ) -> CubeCountInputExpand

pub fn new_Flattened(m_cubes: u32, n_cubes: u32) -> Self

pub fn __expand_new_Flattened( _: &mut Scope, m_cubes: <u32 as CubeType>::ExpandType, n_cubes: <u32 as CubeType>::ExpandType, ) -> CubeCountInputExpand

pub fn new_Spread(m_cubes: u32, n_cubes: u32, batch_cubes: u32) -> Self

pub fn __expand_new_Spread( _: &mut Scope, m_cubes: <u32 as CubeType>::ExpandType, n_cubes: <u32 as CubeType>::ExpandType, batch_cubes: <u32 as CubeType>::ExpandType, ) -> CubeCountInputExpand

impl CubeCountInput

pub fn num_valid_cubes(&self) -> u32

Returns the number of valid cubes

pub fn cube_pos_to_tensor_pos( &self, global_order: GlobalOrder, ) -> (u32, u32, u32)

Given a cube position (SM ID, local index), returns the tensor coordinates (m, n, batch).

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

pub fn __expand_cube_pos_to_tensor_pos( scope: &mut Scope, this: <Self as CubeType>::ExpandType, global_order: GlobalOrder, ) -> <(u32, u32, u32) as CubeType>::ExpandType

Trait Implementations

impl CubeDebug for CubeCountInput

fn set_debug_name(&self, scope: &mut 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 CubeType for CubeCountInput

type ExpandType = CubeCountInputExpand

fn into_mut(scope: &mut Scope, expand: Self::ExpandType) -> Self::ExpandType

Wrapper around the init method, necessary to type inference.

impl LaunchArg for CubeCountInput

type RuntimeArg<'a, R: Runtime> = CubeCountInputArgs<'a, R>

The runtime argument for the kernel.

fn compilation_arg<'a, R: Runtime>( runtime_arg: &Self::RuntimeArg<'a, R>, ) -> Self::CompilationArg

impl LaunchArgExpand for CubeCountInput

type CompilationArg = CubeCountInputCompilationArg

Compilation argument.

fn expand( arg: &Self::CompilationArg, builder: &mut KernelBuilder, ) -> <Self as CubeType>::ExpandType

Register an input variable during compilation that fill the KernelBuilder.

fn expand_output( arg: &Self::CompilationArg, builder: &mut KernelBuilder, ) -> <Self as CubeType>::ExpandType

Register an output variable during compilation that fill the KernelBuilder.