Trait BatchMatmul

pub trait BatchMatmul<MP: MatmulPrecision>:
    'static
    + Send
    + Sync {
    type Config: BatchConfig;

    // Required methods
    fn execute(
        lhs: VirtualTensor<MP::EI>,
        rhs: VirtualTensor<MP::EI>,
        out: VirtualTensor<MP::EO, ReadWrite>,
        size_k: u32,
        quantization: CubeOption<Quantization<MP>>,
        config: Self::Config,
    );
    fn __expand_execute(
        context: &mut Scope,
        lhs: <VirtualTensor<MP::EI> as CubeType>::ExpandType,
        rhs: <VirtualTensor<MP::EI> as CubeType>::ExpandType,
        out: <VirtualTensor<MP::EO, ReadWrite> as CubeType>::ExpandType,
        size_k: <u32 as CubeType>::ExpandType,
        quantization: <CubeOption<Quantization<MP>> as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
}

Provides matrix multiplication operations at the batch level.

At the batch level,

• Inputs are whole tensors in global memory.
• All Cubes can collaborate to solve the problem
• Dimensions M, N and K can be arbitrary large, as well as the number of batches.

Assumptions

• Line sizes of the inputs evenly divide the dimension they are aligned with.
• Enough Cubes are launched to perform the whole computation.

Safety

It is not assumed that the matmul’s dimensions match its inputs dimensions perfectly. It is therefore important to use an underlying global matmul that performs check bounds, and to not launch more Cubes than necessary.

Required Associated Types

type Config: BatchConfig

Required Methods

fn execute( lhs: VirtualTensor<MP::EI>, rhs: VirtualTensor<MP::EI>, out: VirtualTensor<MP::EO, ReadWrite>, size_k: u32, quantization: CubeOption<Quantization<MP>>, config: Self::Config, )

Performs batchwise matrix multiplication over tensors.

fn __expand_execute( context: &mut Scope, lhs: <VirtualTensor<MP::EI> as CubeType>::ExpandType, rhs: <VirtualTensor<MP::EI> as CubeType>::ExpandType, out: <VirtualTensor<MP::EO, ReadWrite> as CubeType>::ExpandType, size_k: <u32 as CubeType>::ExpandType, quantization: <CubeOption<Quantization<MP>> 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, GMM: GlobalMatmul<MP>, C: CubeDispatch> BatchMatmul<MP> for OneToOneMatmul<MP, GMM, C>

type Config = Config<<GMM as GlobalMatmul<MP>>::Config, C>

impl<MP: MatmulPrecision, GMM: GlobalMatmul<MP>, S: SpanMatmul, C: CubeDispatch> BatchMatmul<MP> for OneToManyMatmul<MP, GMM, S, C>

type Config = Config<<GMM as GlobalMatmul<MP>>::Config, C>