Trait StageMatmul

pub trait StageMatmul<MP: MatmulPrecision>:
    'static
    + Send
    + Sync {
    type Config: StageConfig;
    type Accumulator: CubeType;
    type LhsReader: CubeType;
    type RhsReader: CubeType;
    type LhsTile: CubeType;
    type RhsTile: CubeType;

    // Required methods
    fn execute(
        lhs: &Self::LhsReader,
        rhs: &Self::RhsReader,
        instruction_lhs: &mut Self::LhsTile,
        instruction_rhs: &mut Self::RhsTile,
        acc: &mut Self::Accumulator,
        config: Self::Config,
    );
    fn execute_with_listener<SEL: StageEventListener>(
        lhs: &Self::LhsReader,
        rhs: &Self::RhsReader,
        instruction_lhs: &mut Self::LhsTile,
        instruction_rhs: &mut Self::RhsTile,
        acc: &mut Self::Accumulator,
        config: Self::Config,
        listener: SEL,
    );
    fn init_tile_inputs(config: Self::Config) -> (Self::LhsTile, Self::RhsTile);
    fn read_accumulator<Out: StageWriter<MP::EO>, G: GlobalConfig>(
        acc: &Self::Accumulator,
        out: &mut Out,
        stage_config: Self::Config,
        global_config: G,
    );
    fn init_accumulator(config: Self::Config) -> Self::Accumulator;
    fn zero_accumulator(acc: &mut Self::Accumulator, config: Self::Config);
    fn fill_accumulator<L: AccumulatorLoader<MP>>(
        loader: &mut L,
        acc: &mut Self::Accumulator,
        config: Self::Config,
    );
    fn __expand_execute(
        context: &mut Scope,
        lhs: <Self::LhsReader as CubeType>::ExpandType,
        rhs: <Self::RhsReader as CubeType>::ExpandType,
        instruction_lhs: <Self::LhsTile as CubeType>::ExpandType,
        instruction_rhs: <Self::RhsTile as CubeType>::ExpandType,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_execute_with_listener<SEL: StageEventListener>(
        context: &mut Scope,
        lhs: <Self::LhsReader as CubeType>::ExpandType,
        rhs: <Self::RhsReader as CubeType>::ExpandType,
        instruction_lhs: <Self::LhsTile as CubeType>::ExpandType,
        instruction_rhs: <Self::RhsTile as CubeType>::ExpandType,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
        listener: <SEL as CubeType>::ExpandType,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_init_tile_inputs(
        context: &mut Scope,
        config: Self::Config,
    ) -> <(Self::LhsTile, Self::RhsTile) as CubeType>::ExpandType;
    fn __expand_read_accumulator<Out: StageWriter<MP::EO>, G: GlobalConfig>(
        context: &mut Scope,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        out: <Out as CubeType>::ExpandType,
        stage_config: Self::Config,
        global_config: G,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_init_accumulator(
        context: &mut Scope,
        config: Self::Config,
    ) -> <Self::Accumulator as CubeType>::ExpandType;
    fn __expand_zero_accumulator(
        context: &mut Scope,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_fill_accumulator<L: AccumulatorLoader<MP>>(
        context: &mut Scope,
        loader: <L as CubeType>::ExpandType,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
}

Provides matrix multiplication operations at the stage level.

At the stage level,

• Inputs are staged into an intermediate memory called stage (typically a shared memory).
• All planes within a Cube can collaborate 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 are launched to perform the whole computation

Required Associated Types

type Config: StageConfig

type Accumulator: 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 LhsReader: CubeType

type RhsReader: CubeType

type LhsTile: CubeType

type RhsTile: CubeType

Required Methods

fn execute( lhs: &Self::LhsReader, rhs: &Self::RhsReader, instruction_lhs: &mut Self::LhsTile, instruction_rhs: &mut Self::RhsTile, acc: &mut Self::Accumulator, config: Self::Config, )

Executes the matrix multiplication of LHS and RHS, adding the result to the accumulator

Equivalent to execute_with_listener with SEL:=NoEvent

Quantization

If scaling is provided, the matmul will be performed in a quantized version. This assumes that [read_accumulator] is called with some quantization provided.

fn execute_with_listener<SEL: StageEventListener>( lhs: &Self::LhsReader, rhs: &Self::RhsReader, instruction_lhs: &mut Self::LhsTile, instruction_rhs: &mut Self::RhsTile, acc: &mut Self::Accumulator, config: Self::Config, listener: SEL, )

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)

fn read_accumulator<Out: StageWriter<MP::EO>, G: GlobalConfig>( acc: &Self::Accumulator, out: &mut Out, stage_config: Self::Config, global_config: G, )

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

Quantization

If some quantization is provided, the read will also requantize the stage in the output and update the scaling of the output tensor. This assumes that [execute] is called with some scaling provided.

fn init_accumulator(config: Self::Config) -> Self::Accumulator

Create an instance of the accumulator, without data

fn zero_accumulator(acc: &mut Self::Accumulator, config: Self::Config)

Fill the accumulator with zeros

fn fill_accumulator<L: AccumulatorLoader<MP>>( loader: &mut L, acc: &mut Self::Accumulator, config: Self::Config, )

Fill the accumulator with data

fn __expand_execute( context: &mut Scope, lhs: <Self::LhsReader as CubeType>::ExpandType, rhs: <Self::RhsReader as CubeType>::ExpandType, instruction_lhs: <Self::LhsTile as CubeType>::ExpandType, instruction_rhs: <Self::RhsTile as CubeType>::ExpandType, acc: <Self::Accumulator as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

fn __expand_execute_with_listener<SEL: StageEventListener>( context: &mut Scope, lhs: <Self::LhsReader as CubeType>::ExpandType, rhs: <Self::RhsReader as CubeType>::ExpandType, instruction_lhs: <Self::LhsTile as CubeType>::ExpandType, instruction_rhs: <Self::RhsTile as CubeType>::ExpandType, acc: <Self::Accumulator as CubeType>::ExpandType, config: Self::Config, listener: <SEL as CubeType>::ExpandType, ) -> <() as CubeType>::ExpandType

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

fn __expand_read_accumulator<Out: StageWriter<MP::EO>, G: GlobalConfig>( context: &mut Scope, acc: <Self::Accumulator as CubeType>::ExpandType, out: <Out as CubeType>::ExpandType, stage_config: Self::Config, global_config: G, ) -> <() as CubeType>::ExpandType

fn __expand_init_accumulator( context: &mut Scope, config: Self::Config, ) -> <Self::Accumulator as CubeType>::ExpandType

fn __expand_zero_accumulator( context: &mut Scope, acc: <Self::Accumulator as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

fn __expand_fill_accumulator<L: AccumulatorLoader<MP>>( context: &mut Scope, loader: <L as CubeType>::ExpandType, acc: <Self::Accumulator as CubeType>::ExpandType, config: Self::Config, ) -> <() 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

impl<MP, TMM, RL, RR> StageMatmul<MP> for PlaneMatmul<MP, TMM, RL, RR>

where MP: MatmulPrecision, TMM: TileMatmul<MP>, RL: Reader<MP::ES>, RR: Reader<MP::ES>,

type Config = CommonStageConfig<<TMM as TileMatmul<MP>>::Config>

type LhsReader = RL

type RhsReader = RR

type Accumulator = Accumulators<MP, TMM>

type LhsTile = Sequence<<TMM as TileMatmul<MP>>::Lhs>

type RhsTile = RhsTile<<TMM as TileMatmul<MP>>::Rhs>