Trait TileMatmul 

pub trait TileMatmul<L: Numeric, R: Numeric, A: Numeric>:
    'static
    + Send
    + Sync {
    type Config: TileConfig;
    type LhsFragment: CubeType;
    type RhsFragment: CubeType;
    type AccFragment: CubeType;
    type LhsTile: TileKind;
    type RhsTile: TileKind;
    type AccTile: TileKind;
    type OutTile: TileKind<ReadWrite>;

    // Required methods
    fn execute(
        lhs: &Self::LhsFragment,
        rhs: &Self::RhsFragment,
        out: &mut Self::AccFragment,
        config: Self::Config,
    );
    fn allocate_lhs(config: Self::Config) -> Self::LhsFragment;
    fn load_lhs<E: Numeric>(
        tile: &Tile<Self::LhsTile, E>,
        lhs: &mut Self::LhsFragment,
        config: Self::Config,
    );
    fn allocate_rhs(config: Self::Config) -> Self::RhsFragment;
    fn load_rhs<E: Numeric>(
        tile: &Tile<Self::RhsTile, E>,
        rhs: &mut Self::RhsFragment,
        config: Self::Config,
    );
    fn allocate_acc(config: Self::Config) -> Self::AccFragment;
    fn load_acc<E: Numeric>(
        tile: &Tile<Self::AccTile, E>,
        acc: &mut Self::AccFragment,
        config: Self::Config,
    );
    fn write_results<E: Numeric>(
        tile: &mut TileMut<Self::OutTile, E>,
        out: &Self::AccFragment,
        config: Self::Config,
    );
    fn __expand_execute(
        scope: &mut Scope,
        lhs: <Self::LhsFragment as CubeType>::ExpandType,
        rhs: <Self::RhsFragment as CubeType>::ExpandType,
        out: <Self::AccFragment as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_lhs(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::LhsFragment as CubeType>::ExpandType;
    fn __expand_load_lhs<E: Numeric>(
        scope: &mut Scope,
        tile: <Tile<Self::LhsTile, E> as CubeType>::ExpandType,
        lhs: <Self::LhsFragment as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_rhs(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::RhsFragment as CubeType>::ExpandType;
    fn __expand_load_rhs<E: Numeric>(
        scope: &mut Scope,
        tile: <Tile<Self::RhsTile, E> as CubeType>::ExpandType,
        rhs: <Self::RhsFragment as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_acc(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::AccFragment as CubeType>::ExpandType;
    fn __expand_load_acc<E: Numeric>(
        scope: &mut Scope,
        tile: <Tile<Self::AccTile, E> as CubeType>::ExpandType,
        acc: <Self::AccFragment as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_write_results<E: Numeric>(
        scope: &mut Scope,
        tile: <TileMut<Self::OutTile, E> as CubeType>::ExpandType,
        out: <Self::AccFragment as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
}

Provides matrix multiplication operations at the tile level.

At the tile level,

• Dimensions M, N and K are fixed to an integer, and the matrix multiplication works only for size (M, K) · (K, N) = (M, N).

Assumptions:

• Inputs 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 units are present to perform the whole computation

Required Associated Types

type Config: TileConfig

The configuration type associated with this Matmul.

type LhsFragment: CubeType

Contains Lhs data for computation

type RhsFragment: CubeType

Contains Rhs data for computation

type AccFragment: CubeType

Contains and accumulates results of the Tile Matmul execution

type LhsTile: TileKind

Tile for the lhs data

type RhsTile: TileKind

Tile for the rhs data

type AccTile: TileKind

Tile for the accumulator data

type OutTile: TileKind<ReadWrite>

Tile for the output data

Required Methods

fn execute( lhs: &Self::LhsFragment, rhs: &Self::RhsFragment, out: &mut Self::AccFragment, config: Self::Config, )

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

fn allocate_lhs(config: Self::Config) -> Self::LhsFragment

Create the container for Lhs

Safety

This may point towards uninitialized memory. Make sure to call load_lhs prior to execute.

fn load_lhs<E: Numeric>( tile: &Tile<Self::LhsTile, E>, lhs: &mut Self::LhsFragment, config: Self::Config, )

Load the container of Lhs from tile data

fn allocate_rhs(config: Self::Config) -> Self::RhsFragment

Create the container for Rhs

Safety

This may point towards uninitialized memory. Make sure to call load_rhs prior to execute.

fn load_rhs<E: Numeric>( tile: &Tile<Self::RhsTile, E>, rhs: &mut Self::RhsFragment, config: Self::Config, )

Load the container of Rhs from tile data

fn allocate_acc(config: Self::Config) -> Self::AccFragment

Allocate the container to receive the execution output.

Safety

The output container must be initialized to some value (typically 0), because the execution adds to the already present value. Make sure to call load_acc prior to execute.

fn load_acc<E: Numeric>( tile: &Tile<Self::AccTile, E>, acc: &mut Self::AccFragment, config: Self::Config, )

Load the container of Acc from tile data

fn write_results<E: Numeric>( tile: &mut TileMut<Self::OutTile, E>, out: &Self::AccFragment, config: Self::Config, )

Write the content of the output container to the given slice

fn __expand_execute( scope: &mut Scope, lhs: <Self::LhsFragment as CubeType>::ExpandType, rhs: <Self::RhsFragment as CubeType>::ExpandType, out: <Self::AccFragment as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_load_lhs<E: Numeric>( scope: &mut Scope, tile: <Tile<Self::LhsTile, E> as CubeType>::ExpandType, lhs: <Self::LhsFragment as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_load_rhs<E: Numeric>( scope: &mut Scope, tile: <Tile<Self::RhsTile, E> as CubeType>::ExpandType, rhs: <Self::RhsFragment as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_load_acc<E: Numeric>( scope: &mut Scope, tile: <Tile<Self::AccTile, E> as CubeType>::ExpandType, acc: <Self::AccFragment as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

fn __expand_write_results<E: Numeric>( scope: &mut Scope, tile: <TileMut<Self::OutTile, E> as CubeType>::ExpandType, out: <Self::AccFragment 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<L: Numeric, R: Numeric, A: Numeric, AccTile: TileKind> TileMatmul<L, R, A> for AcceleratedMatmul<AccTile>

where CmmaStageReader<AccTile>: CmmaFragmentReader<TileKind = AccTile>,

type Config = AcceleratedConfig

type LhsFragment = Matrix<L>

type RhsFragment = Matrix<R>

type AccFragment = Matrix<A>

type LhsTile = Strided

type RhsTile = Strided

type AccTile = AccTile

type OutTile = Strided

impl<L: Numeric, R: Numeric, A: Numeric, AccTile: TileKind> TileMatmul<L, R, A> for MmaMatmul<AccTile>

where MmaStageReader<AccTile>: MmaFragmentReader<TileKind = AccTile>,

type Config = MmaMatmulConfig

type LhsFragment = Sequence<Line<L>>

type RhsFragment = Sequence<Line<R>>

type AccFragment = Sequence<Line<A>>

type LhsTile = Strided

type RhsTile = Strided

type AccTile = AccTile

type OutTile = Strided

impl<L: Numeric, R: Numeric, A: Numeric, AccTile: TileKind> TileMatmul<L, R, A> for PlaneVecMatInnerProduct<AccTile>

where MatrixStageReader<AccTile>: MatrixFragmentReader<TileKind = AccTile>,

type Config = PlaneVecMatInnerProductConfig

type LhsFragment = LineContainer<L>

type RhsFragment = Sequence<LineContainer<R>>

type AccFragment = Sequence<LineContainer<A>>

type LhsTile = Strided

type RhsTile = Strided

type AccTile = AccTile

type OutTile = Strided

impl<L: Numeric, R: Numeric, A: Numeric, AccTile: TileKind> TileMatmul<L, R, A> for RegisterMatmul<AccTile>

where RegisterStageReader<AccTile>: RegisterFragmentReader<TileKind = AccTile>,

type Config = RegisterConfig

type LhsFragment = Array<L>

type RhsFragment = Array<R>

type AccFragment = Array<A>

type LhsTile = Strided

type RhsTile = Strided

type AccTile = AccTile

type OutTile = Strided