Trait StageMatmulFamily

pub trait StageMatmulFamily:
    Send
    + Sync
    + 'static {
    type Matmul<MP: MatmulPrecision, TL: TilingLayout, TR: TilingLayout>: StageMatmul<MP, Config = Self::Config, LhsReader = <Self::LhsReader as ReaderFamily>::Reader<MP::ES, TL>, RhsReader = <Self::RhsReader as ReaderFamily>::Reader<MP::ES, TR>>;
    type LhsReader: ReaderFamily;
    type RhsReader: ReaderFamily;
    type Config: StageConfig;

    // Required method
    fn setup<MP: MatmulPrecision, R: Runtime>(
        client: &ComputeClient<R::Server, R::Channel>,
        problem: &MatmulProblem,
        selection: &MatmulSelection,
        line_sizes: &MatmulLineSizes,
        num_stages: NumStages,
        max_loaders: Option<MaxLoaderPlanes>,
        ordered: bool,
    ) -> Result<Self::Config, MatmulSetupError>;

    // Provided method
    fn filter_line_sizes(
        available_line_sizes: AvailableLineSizes,
    ) -> AvailableLineSizes { ... }
}

A family of StageMatmul implementations that operate with any precision.

Required Associated Types

type Matmul<MP: MatmulPrecision, TL: TilingLayout, TR: TilingLayout>: StageMatmul<MP, Config = Self::Config, LhsReader = <Self::LhsReader as ReaderFamily>::Reader<MP::ES, TL>, RhsReader = <Self::RhsReader as ReaderFamily>::Reader<MP::ES, TR>>

The specific [TileMatmul] implementation associated with this family.

type LhsReader: ReaderFamily

Reader family for Lhs

type RhsReader: ReaderFamily

Reader family for Rhs

type Config: StageConfig

The configuration type associated with this matmul family.

Required Methods

fn setup<MP: MatmulPrecision, R: Runtime>( client: &ComputeClient<R::Server, R::Channel>, problem: &MatmulProblem, selection: &MatmulSelection, line_sizes: &MatmulLineSizes, num_stages: NumStages, max_loaders: Option<MaxLoaderPlanes>, ordered: bool, ) -> Result<Self::Config, MatmulSetupError>

Constructs the configuration based on the matmul problem, selection, line sizes, number of stages, maximum of tasks per plane, and whether the algorithm is an ordered variant

This function may return an error if the configuration cannot be supported on the current runtime.

Provided Methods

fn filter_line_sizes( available_line_sizes: AvailableLineSizes, ) -> AvailableLineSizes

Filters out line sizes that are incompatible with this matmul family.

By default, returns the input unchanged.

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<TMM: TileMatmulFamily, LRF: ReaderFamily, RRF: ReaderFamily> StageMatmulFamily for PlaneMatmulFamily<TMM, LRF, RRF>

type LhsReader = LRF

type RhsReader = RRF

type Matmul<MP: MatmulPrecision, TL: TilingLayout, TR: TilingLayout> = PartitionedStageMatmul<MP, <TMM as TileMatmulFamily>::Matmul<MP>, <LRF as ReaderFamily>::Reader<<MP as MatmulPrecision>::ES, TL>, <RRF as ReaderFamily>::Reader<<MP as MatmulPrecision>::ES, TR>, PlanePartitioner, PlanePartitionedStageConfig<<<TMM as TileMatmulFamily>::Matmul<MP> as TileMatmul<MP>>::Config>>

type Config = PlanePartitionedStageConfig<<TMM as TileMatmulFamily>::Config>

impl<TMM: TileMatmulFamily, RF: ReaderFamily> StageMatmulFamily for UnitMatmulFamily<TMM, RF>

type LhsReader = RF

type RhsReader = RF

type Matmul<MP: MatmulPrecision, TL: TilingLayout, TR: TilingLayout> = PartitionedStageMatmul<MP, <TMM as TileMatmulFamily>::Matmul<MP>, <RF as ReaderFamily>::Reader<<MP as MatmulPrecision>::ES, TL>, <RF as ReaderFamily>::Reader<<MP as MatmulPrecision>::ES, TR>, UnitPartitioner, UnitPartitionedStageConfig<<<TMM as TileMatmulFamily>::Matmul<MP> as TileMatmul<MP>>::Config>>

type Config = UnitPartitionedStageConfig<<TMM as TileMatmulFamily>::Config>