Trait StageMatmul 

pub trait StageMatmul<MP: MatmulPrecision>:
    'static
    + Send
    + Sync {
    type Config: StageConfig;
    type Accumulators: CubeType;
    type LhsStage: CubeType;
    type RhsStage: CubeType;
    type AccStage: CubeType;
    type OutStage: CubeType;
    type LhsTile: CubeType;
    type RhsTile: CubeType;

    // Required methods
    fn execute(
        lhs: &Self::LhsStage,
        rhs: &Self::RhsStage,
        instruction_lhs: &mut Self::LhsTile,
        instruction_rhs: &mut Self::RhsTile,
        acc: &mut Self::Accumulators,
        config: Self::Config,
        partition_scheduler: &PartitionScheduler,
    );
    fn execute_with_listener<SEL: StageEventListener<Self::Config>>(
        lhs: &Self::LhsStage,
        rhs: &Self::RhsStage,
        instruction_lhs: &mut Self::LhsTile,
        instruction_rhs: &mut Self::RhsTile,
        acc: &mut Self::Accumulators,
        config: Self::Config,
        listener: SEL,
        partition_scheduler: &PartitionScheduler,
    );
    fn init_tile_inputs(config: Self::Config) -> (Self::LhsTile, Self::RhsTile);
    fn init_accumulators(config: Self::Config) -> Self::Accumulators;
    fn load_accumulators(
        reader: &Self::AccStage,
        acc: &mut Self::Accumulators,
        config: Self::Config,
    );
    fn write_results<W: WriteEventListener, G: GlobalConfig>(
        acc: &Self::Accumulators,
        stage: &mut Self::OutStage,
        listener: &mut W,
        partition_scheduler: &PartitionScheduler,
        stage_config: Self::Config,
        global_config: G,
    );
    fn init_scheduler(config: Self::Config) -> PartitionScheduler;
    fn __expand_execute(
        scope: &mut Scope,
        lhs: <Self::LhsStage as CubeType>::ExpandType,
        rhs: <Self::RhsStage as CubeType>::ExpandType,
        instruction_lhs: <Self::LhsTile as CubeType>::ExpandType,
        instruction_rhs: <Self::RhsTile as CubeType>::ExpandType,
        acc: <Self::Accumulators as CubeType>::ExpandType,
        config: Self::Config,
        partition_scheduler: <PartitionScheduler as CubeType>::ExpandType,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_execute_with_listener<SEL: StageEventListener<Self::Config>>(
        scope: &mut Scope,
        lhs: <Self::LhsStage as CubeType>::ExpandType,
        rhs: <Self::RhsStage as CubeType>::ExpandType,
        instruction_lhs: <Self::LhsTile as CubeType>::ExpandType,
        instruction_rhs: <Self::RhsTile as CubeType>::ExpandType,
        acc: <Self::Accumulators as CubeType>::ExpandType,
        config: Self::Config,
        listener: <SEL as CubeType>::ExpandType,
        partition_scheduler: <PartitionScheduler as CubeType>::ExpandType,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_init_tile_inputs(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <(Self::LhsTile, Self::RhsTile) as CubeType>::ExpandType;
    fn __expand_init_accumulators(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::Accumulators as CubeType>::ExpandType;
    fn __expand_load_accumulators(
        scope: &mut Scope,
        reader: <Self::AccStage as CubeType>::ExpandType,
        acc: <Self::Accumulators as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_write_results<W: WriteEventListener, G: GlobalConfig>(
        scope: &mut Scope,
        acc: <Self::Accumulators as CubeType>::ExpandType,
        stage: <Self::OutStage as CubeType>::ExpandType,
        listener: <W as CubeType>::ExpandType,
        partition_scheduler: <PartitionScheduler as CubeType>::ExpandType,
        stage_config: Self::Config,
        global_config: G,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_init_scheduler(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <PartitionScheduler as CubeType>::ExpandType;
}

Provides matrix multiplication operations at the stage level.

At the stage level,

• Inputs are assumed to be already staged into a shared memory.
• All main flow planes within a Cube are used to solve the problem
• Dimensions M, N and K are fixed to an integer, and the matrix multiplication works only for size (M, K) · (K, N) = (M, N). These integers are multiples of the underlying Tile matmul, corresponding to the number of tiles in each dimension.

Assumptions:

• Data given as inputs by stage readers must always be valid. If the actual matrix multiplication should be done on smaller sizes than M, N and K, padding with zeros must be done beforehand.
• Enough planes/units are launched to perform the whole computation

Required Associated Types

type Config: StageConfig

The configuration type associated with this Matmul.

type Accumulators: CubeType

Contains the matrix multiplication output, that can be shared across the different planes of the cube. The same Accumulator will be added to across multiple executions of the Stage Matmul.

type LhsStage: CubeType

Stage for Lhs

type RhsStage: CubeType

Stage for Rhs

type AccStage: CubeType

Stage for Accumulator

type OutStage: CubeType

Stage for Out

type LhsTile: CubeType

Lhs input of the underlying Tile Matmul

type RhsTile: CubeType

Rhs input of the underlying Tile Matmul

Required Methods

fn execute( lhs: &Self::LhsStage, rhs: &Self::RhsStage, instruction_lhs: &mut Self::LhsTile, instruction_rhs: &mut Self::RhsTile, acc: &mut Self::Accumulators, config: Self::Config, partition_scheduler: &PartitionScheduler, )

Executes the matrix multiplication of Lhs and Rhs, adding the result to the accumulator

Equivalent to execute_with_listener with SEL:=NoEvent

fn execute_with_listener<SEL: StageEventListener<Self::Config>>( lhs: &Self::LhsStage, rhs: &Self::RhsStage, instruction_lhs: &mut Self::LhsTile, instruction_rhs: &mut Self::RhsTile, acc: &mut Self::Accumulators, config: Self::Config, listener: SEL, partition_scheduler: &PartitionScheduler, )

Executes the matrix multiplication of Lhs and Rhs, with the addition of injected event listener.

fn init_tile_inputs(config: Self::Config) -> (Self::LhsTile, Self::RhsTile)

Inits inputs of the underlying Tile Matmul

fn init_accumulators(config: Self::Config) -> Self::Accumulators

Create an instance of the accumulators, without data

fn load_accumulators( reader: &Self::AccStage, acc: &mut Self::Accumulators, config: Self::Config, )

Load all accumulators in the stage from data

fn write_results<W: WriteEventListener, G: GlobalConfig>( acc: &Self::Accumulators, stage: &mut Self::OutStage, listener: &mut W, partition_scheduler: &PartitionScheduler, stage_config: Self::Config, global_config: G, )

Reads the result of the accumulator and hands it to the stage writer

fn init_scheduler(config: Self::Config) -> PartitionScheduler

fn __expand_execute( scope: &mut Scope, lhs: <Self::LhsStage as CubeType>::ExpandType, rhs: <Self::RhsStage as CubeType>::ExpandType, instruction_lhs: <Self::LhsTile as CubeType>::ExpandType, instruction_rhs: <Self::RhsTile as CubeType>::ExpandType, acc: <Self::Accumulators as CubeType>::ExpandType, config: Self::Config, partition_scheduler: <PartitionScheduler as CubeType>::ExpandType, ) -> <() as CubeType>::ExpandType

fn __expand_execute_with_listener<SEL: StageEventListener<Self::Config>>( scope: &mut Scope, lhs: <Self::LhsStage as CubeType>::ExpandType, rhs: <Self::RhsStage as CubeType>::ExpandType, instruction_lhs: <Self::LhsTile as CubeType>::ExpandType, instruction_rhs: <Self::RhsTile as CubeType>::ExpandType, acc: <Self::Accumulators as CubeType>::ExpandType, config: Self::Config, listener: <SEL as CubeType>::ExpandType, partition_scheduler: <PartitionScheduler as CubeType>::ExpandType, ) -> <() as CubeType>::ExpandType

fn __expand_init_tile_inputs( scope: &mut Scope, config: Self::Config, ) -> <(Self::LhsTile, Self::RhsTile) as CubeType>::ExpandType

fn __expand_init_accumulators( scope: &mut Scope, config: Self::Config, ) -> <Self::Accumulators as CubeType>::ExpandType

fn __expand_load_accumulators( scope: &mut Scope, reader: <Self::AccStage as CubeType>::ExpandType, acc: <Self::Accumulators as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

fn __expand_write_results<W: WriteEventListener, G: GlobalConfig>( scope: &mut Scope, acc: <Self::Accumulators as CubeType>::ExpandType, stage: <Self::OutStage as CubeType>::ExpandType, listener: <W as CubeType>::ExpandType, partition_scheduler: <PartitionScheduler as CubeType>::ExpandType, stage_config: Self::Config, global_config: G, ) -> <() as CubeType>::ExpandType

fn __expand_init_scheduler( scope: &mut Scope, config: Self::Config, ) -> <PartitionScheduler as CubeType>::ExpandType

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety", so this trait is not object safe.

Implementors