Trait TileMatmul

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

    // 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 allocate_rhs(config: Self::Config) -> Self::Rhs;
    fn fill_lhs(slice: &Tile<MP::ES>, lhs: &mut Self::Lhs, config: Self::Config);
    fn fill_rhs(slice: &Tile<MP::ES>, rhs: &mut Self::Rhs, config: Self::Config);
    fn fill_accumulator(
        tile: &Tile<MP::EA>,
        acc: &mut Self::Accumulator,
        config: Self::Config,
    );
    fn read_accumulator<C: Numeric>(
        out: &Self::Accumulator,
        slice: &mut SliceMut<Line<C>>,
        config: Self::Config,
    );
    fn allocate_accumulator(config: Self::Config) -> Self::Accumulator;
    fn zero_accumulator(acc: &mut Self::Accumulator, config: Self::Config);
    fn __expand_execute(
        context: &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(
        context: &mut Scope,
        config: Self::Config,
    ) -> <Self::Lhs as CubeType>::ExpandType;
    fn __expand_allocate_rhs(
        context: &mut Scope,
        config: Self::Config,
    ) -> <Self::Rhs as CubeType>::ExpandType;
    fn __expand_fill_lhs(
        context: &mut Scope,
        slice: <Tile<MP::ES> as CubeType>::ExpandType,
        lhs: <Self::Lhs as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_fill_rhs(
        context: &mut Scope,
        slice: <Tile<MP::ES> as CubeType>::ExpandType,
        rhs: <Self::Rhs as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_fill_accumulator(
        context: &mut Scope,
        tile: <Tile<MP::EA> as CubeType>::ExpandType,
        acc: <Self::Accumulator as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_read_accumulator<C: Numeric>(
        context: &mut Scope,
        out: <Self::Accumulator as CubeType>::ExpandType,
        slice: <SliceMut<Line<C>> as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_allocate_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;
}

Provides matrix multiplication operations at the tile level.

At the tile level,

• Inputs are raw slices of data, called tiles.
• units within one plane 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).

Assumptions:

• Slices given as 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

type Lhs: CubeType

Contains LHS data that can be split across the units

type Rhs: CubeType

Contains RHS data that can be split across the units

type Accumulator: CubeType + Copy + Clone

Contains output data that can be split across the units

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 output

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

Create the container for LHS data

Safety

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

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

Create the container for RHS data

Safety

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

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

Fill the container of LHS with data

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

Fill the container of RHS with data

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

Fill the accumulator with data

fn read_accumulator<C: Numeric>( out: &Self::Accumulator, slice: &mut SliceMut<Line<C>>, config: Self::Config, )

Write the content of the output container to the given slice

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 zero_accumulator(acc: &mut Self::Accumulator, config: Self::Config)

Fill the accumulator with zeros.

fn __expand_execute( context: &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( context: &mut Scope, config: Self::Config, ) -> <Self::Lhs as CubeType>::ExpandType

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

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

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

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

fn __expand_read_accumulator<C: Numeric>( context: &mut Scope, out: <Self::Accumulator as CubeType>::ExpandType, slice: <SliceMut<Line<C>> as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

fn __expand_allocate_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

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 Accelerated

type Config = Config

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

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

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