GlobalConvolution

Trait GlobalConvolution 

Source
pub trait GlobalConvolution<MP: MatmulPrecision>:
    'static
    + Send
    + Sync {
    type LhsGlobalReader: CubeType;
    type RhsGlobalReader: CubeType;
    type AccGlobalReader: CubeType;
    type Config: ConvGemmConfig;
    type GlobalWriter: GlobalWriter<MP::Acc>;
    type Accumulators: CubeType;

    // Required methods
    fn execute(
        lhs_reader: Self::LhsGlobalReader,
        rhs_reader: Self::RhsGlobalReader,
        acc_reader: Self::AccGlobalReader,
        writer: Self::GlobalWriter,
        acc: &mut Self::Accumulators,
        k_range: (u32, u32),
        config: Self::Config,
    );
    fn init_lhs_global_reader(
        lhs: View<Line<LhsG<MP>>, Coords2d>,
        offset: Coords2d,
        slice_size: Coords2d,
        runtime_args: &RuntimeArgs,
        config: Self::Config,
    ) -> Self::LhsGlobalReader;
    fn init_rhs_global_reader(
        rhs: View<Line<RhsG<MP>>, Coords2d>,
        config: Self::Config,
    ) -> Self::RhsGlobalReader;
    fn init_bias_global_reader(
        bias: CubeOption<View<Line<AccG<MP>>, Coords2d>>,
        config: Self::Config,
    ) -> Self::AccGlobalReader;
    fn init_global_writer(
        out: View<Line<AccG<MP>>, Coords2d, ReadWrite>,
        config: Self::Config,
    ) -> Self::GlobalWriter;
    fn init_accumulator(config: Self::Config) -> Self::Accumulators;
    fn __expand_execute(
        scope: &mut Scope,
        lhs_reader: <Self::LhsGlobalReader as CubeType>::ExpandType,
        rhs_reader: <Self::RhsGlobalReader as CubeType>::ExpandType,
        acc_reader: <Self::AccGlobalReader as CubeType>::ExpandType,
        writer: <Self::GlobalWriter as CubeType>::ExpandType,
        acc: <Self::Accumulators as CubeType>::ExpandType,
        k_range: <(u32, u32) as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <() as CubeType>::ExpandType;
    fn __expand_init_lhs_global_reader(
        scope: &mut Scope,
        lhs: <View<Line<LhsG<MP>>, Coords2d> as CubeType>::ExpandType,
        offset: <Coords2d as CubeType>::ExpandType,
        slice_size: <Coords2d as CubeType>::ExpandType,
        runtime_args: <RuntimeArgs as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <Self::LhsGlobalReader as CubeType>::ExpandType;
    fn __expand_init_rhs_global_reader(
        scope: &mut Scope,
        rhs: <View<Line<RhsG<MP>>, Coords2d> as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <Self::RhsGlobalReader as CubeType>::ExpandType;
    fn __expand_init_bias_global_reader(
        scope: &mut Scope,
        bias: <CubeOption<View<Line<AccG<MP>>, Coords2d>> as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <Self::AccGlobalReader as CubeType>::ExpandType;
    fn __expand_init_global_writer(
        scope: &mut Scope,
        out: <View<Line<AccG<MP>>, Coords2d, ReadWrite> as CubeType>::ExpandType,
        config: Self::Config,
    ) -> <Self::GlobalWriter as CubeType>::ExpandType;
    fn __expand_init_accumulator(
        scope: &mut Scope,
        config: Self::Config,
    ) -> <Self::Accumulators as CubeType>::ExpandType;
}

Required Associated Types§

Source

type LhsGlobalReader: CubeType

The global reader for the Lhs (input feature map) tensor

Source

type RhsGlobalReader: CubeType

The global reader for the Rhs (weight) tensor

Source

type AccGlobalReader: CubeType

The global reader for the accumulator (bias) tensor

Source

type Config: ConvGemmConfig

The config type of the convolution

Source

type GlobalWriter: GlobalWriter<MP::Acc>

The writer used to write the results to the output feature map

Source

type Accumulators: CubeType

The type of the tile matmul accumulator

Required Methods§

Source

fn execute( lhs_reader: Self::LhsGlobalReader, rhs_reader: Self::RhsGlobalReader, acc_reader: Self::AccGlobalReader, writer: Self::GlobalWriter, acc: &mut Self::Accumulators, k_range: (u32, u32), config: Self::Config, )

Performs the convolution over data loaded by the LHS and RHS readers, over the range given for K, and stores with using the output writer.

To compute the whole range of k values, use k_range=(0, K) where K is the K dimension of LHS and RHS.

Source

fn init_lhs_global_reader( lhs: View<Line<LhsG<MP>>, Coords2d>, offset: Coords2d, slice_size: Coords2d, runtime_args: &RuntimeArgs, config: Self::Config, ) -> Self::LhsGlobalReader

Initializes the global reader for the input feature map with an appropriate layout

Source

fn init_rhs_global_reader( rhs: View<Line<RhsG<MP>>, Coords2d>, config: Self::Config, ) -> Self::RhsGlobalReader

Initializes the global reader for the weights with an appropriate layout

Source

fn init_bias_global_reader( bias: CubeOption<View<Line<AccG<MP>>, Coords2d>>, config: Self::Config, ) -> Self::AccGlobalReader

Initializes the global reader for the bias with an appropriate layout

Source

fn init_global_writer( out: View<Line<AccG<MP>>, Coords2d, ReadWrite>, config: Self::Config, ) -> Self::GlobalWriter

Initializes the output feature map global writer with an appropriate layout

Source

fn init_accumulator(config: Self::Config) -> Self::Accumulators

Initializes a new accumulator for the tile matmul

Source

fn __expand_execute( scope: &mut Scope, lhs_reader: <Self::LhsGlobalReader as CubeType>::ExpandType, rhs_reader: <Self::RhsGlobalReader as CubeType>::ExpandType, acc_reader: <Self::AccGlobalReader as CubeType>::ExpandType, writer: <Self::GlobalWriter as CubeType>::ExpandType, acc: <Self::Accumulators as CubeType>::ExpandType, k_range: <(u32, u32) as CubeType>::ExpandType, config: Self::Config, ) -> <() as CubeType>::ExpandType

Source

fn __expand_init_lhs_global_reader( scope: &mut Scope, lhs: <View<Line<LhsG<MP>>, Coords2d> as CubeType>::ExpandType, offset: <Coords2d as CubeType>::ExpandType, slice_size: <Coords2d as CubeType>::ExpandType, runtime_args: <RuntimeArgs as CubeType>::ExpandType, config: Self::Config, ) -> <Self::LhsGlobalReader as CubeType>::ExpandType

Source

fn __expand_init_rhs_global_reader( scope: &mut Scope, rhs: <View<Line<RhsG<MP>>, Coords2d> as CubeType>::ExpandType, config: Self::Config, ) -> <Self::RhsGlobalReader as CubeType>::ExpandType

Source

fn __expand_init_bias_global_reader( scope: &mut Scope, bias: <CubeOption<View<Line<AccG<MP>>, Coords2d>> as CubeType>::ExpandType, config: Self::Config, ) -> <Self::AccGlobalReader as CubeType>::ExpandType

Source

fn __expand_init_global_writer( scope: &mut Scope, out: <View<Line<AccG<MP>>, Coords2d, ReadWrite> as CubeType>::ExpandType, config: Self::Config, ) -> <Self::GlobalWriter as CubeType>::ExpandType

Source

fn __expand_init_accumulator( scope: &mut Scope, config: Self::Config, ) -> <Self::Accumulators 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§

Source§

impl<MP: MatmulPrecision, SMM> GlobalConvolution<MP> for MultiStageTmaConvolution<MP, SMM>
where SMM: StageMatmul<MP, LhsStage = StridedStage<LhsS<MP>, TmaIm2colTiling>, RhsStage = StridedStage<RhsS<MP>, TmaWeightTiling>, AccStage = BiasStage<AccS<MP>>, OutStage = PartitionedStage<AccS<MP>>>,

Source§

impl<MP: MatmulPrecision, SMM> GlobalConvolution<MP> for SimpleConvolution<MP, SMM>
where SMM: StageMatmul<MP, LhsStage = StridedStage<LhsS<MP>, ConvTilingLayout>, RhsStage = StridedStage<RhsS<MP>, ConvTilingLayout>, AccStage = BiasStage<AccS<MP>>, OutStage = PartitionedStage<AccS<MP>>>,

Source§

impl<MP: MatmulPrecision, SMM> GlobalConvolution<MP> for SimpleTmaConvolution<MP, SMM>
where SMM: StageMatmul<MP, LhsStage = StridedStage<LhsS<MP>, TmaIm2colTiling>, RhsStage = StridedStage<RhsS<MP>, TmaWeightTiling>, AccStage = BiasStage<AccS<MP>>, OutStage = PartitionedStage<AccS<MP>>>,