Trait TileMatmul

pub trait TileMatmul<MP: MatmulPrecision>:
    'static
    + Send
    + Sync {
    type Config: TileConfig;
    type Lhs: CubeType;
    type Rhs: CubeType;
    type Accumulator: CubeType;

    // Required methods
    fn execute(
        lhs: &Self::Lhs,
        rhs: &Self::Rhs,
        out: &mut Self::Accumulator,
        config: Self::Config,
    );
    fn allocate_lhs(config: Self::Config) -> Self::Lhs;
    fn fill_lhs(tile: &Tile<MP::ES>, lhs: &mut Self::Lhs, config: Self::Config);
    fn allocate_rhs(config: Self::Config) -> Self::Rhs;
    fn fill_rhs(tile: &Tile<MP::ES>, rhs: &mut Self::Rhs, config: Self::Config);
    fn allocate_accumulator(config: Self::Config) -> Self::Accumulator;
    fn fill_accumulator(
        tile: &Tile<MP::EA>,
        acc: &mut Self::Accumulator,
        config: Self::Config,
    );
    fn zero_accumulator(acc: &mut Self::Accumulator, config: Self::Config);
    fn write_results(
        out: &Self::Accumulator,
        slice: &mut SliceMut<Line<MP::EO>>,
        config: Self::Config,
    );
    fn __expand_execute(
        scope: &mut Scope,
        lhs: <Self::Lhs as CubeType>::ExpandType,
        rhs: <Self::Rhs as CubeType>::ExpandType,
        out: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_lhs(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::Lhs as CubeType>::ExpandType;
    fn __expand_fill_lhs(
        scope: &mut Scope,
        tile: <Tile<MP::ES> as CubeType>::ExpandType,
        lhs: <Self::Lhs as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_rhs(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::Rhs as CubeType>::ExpandType;
    fn __expand_fill_rhs(
        scope: &mut Scope,
        tile: <Tile<MP::ES> as CubeType>::ExpandType,
        rhs: <Self::Rhs as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_accumulator(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::Accumulator as CubeType>::ExpandType;
    fn __expand_fill_accumulator(
        scope: &mut Scope,
        tile: <Tile<MP::EA> as CubeType>::ExpandType,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_zero_accumulator(
        scope: &mut Scope,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_write_results(
        scope: &mut Scope,
        out: <Self::Accumulator as CubeType>::ExpandType,
        slice: <SliceMut<Line<MP::EO>> 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 Lhs: CubeType

Contains Lhs data for computation

type Rhs: CubeType

Contains Rhs data for computation

type Accumulator: CubeType

Contains and accumulates results of the Tile Matmul execution

Required Methods

fn execute( lhs: &Self::Lhs, rhs: &Self::Rhs, out: &mut Self::Accumulator, config: Self::Config, )

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

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

Create the container for Lhs

Safety

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

fn fill_lhs(tile: &Tile<MP::ES>, lhs: &mut Self::Lhs, config: Self::Config)

Fill the container of Lhs with tile data

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

Create the container for Rhs

Safety

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

fn fill_rhs(tile: &Tile<MP::ES>, rhs: &mut Self::Rhs, config: Self::Config)

Fill the container of Rhs with tile data

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

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 either fill_accumulator or zero_accumulator.

fn fill_accumulator( tile: &Tile<MP::EA>, acc: &mut Self::Accumulator, config: Self::Config, )

Fill the accumulator with data

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

Fill the accumulator with zeros.

fn write_results( out: &Self::Accumulator, slice: &mut SliceMut<Line<MP::EO>>, config: Self::Config, )

Write the content of the output container to the given slice

fn __expand_execute( scope: &mut Scope, lhs: <Self::Lhs as CubeType>::ExpandType, rhs: <Self::Rhs as CubeType>::ExpandType, out: <Self::Accumulator as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_fill_lhs( scope: &mut Scope, tile: <Tile<MP::ES> as CubeType>::ExpandType, lhs: <Self::Lhs as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_fill_rhs( scope: &mut Scope, tile: <Tile<MP::ES> as CubeType>::ExpandType, rhs: <Self::Rhs as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_fill_accumulator( scope: &mut Scope, tile: <Tile<MP::EA> as CubeType>::ExpandType, acc: <Self::Accumulator as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

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

fn __expand_write_results( scope: &mut Scope, out: <Self::Accumulator as CubeType>::ExpandType, slice: <SliceMut<Line<MP::EO>> 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: MatmulPrecision> TileMatmul<MP> for AcceleratedMatmul

type Config = AcceleratedConfig

type Lhs = Matrix<<MP as MatmulPrecision>::ES>

type Rhs = Matrix<<MP as MatmulPrecision>::ES>

type Accumulator = Matrix<<MP as MatmulPrecision>::EA>

impl<MP: MatmulPrecision> TileMatmul<MP> for RegisterMatmul

type Config = RegisterConfig

type Lhs = Array<<MP as MatmulPrecision>::ES>

type Rhs = Array<<MP as MatmulPrecision>::ES>

type Accumulator = TileAccumulator<<MP as MatmulPrecision>::EA>