cubek_convolution/routines/
specialized.rs

1use cubecl::{client::Client, ir::ElemType, prelude::TensorBinding, server::LaunchError};
2use cubek_matmul::multi_level::tile::ColMajorTilingOrder;
3use cubek_matmul::{
4    definition::AvailableVectorSizes,
5    multi_level::{
6        args::{TensorArgs, TensorMapArgs},
7        components::global::read::{
8            AsyncPartialLoadingStrategy, async_partial_cyclic::AsyncPartialCyclicLoading,
9            async_partial_strided::AsyncPartialStridedLoading,
10            async_partial_tma::AsyncPartialTmaLoading,
11        },
12        definition::BatchMatmulBlueprint,
13        routines::batch::specialized::{SpecializedAlgorithm, SpecializedStrategy},
14    },
15};
16use std::marker::PhantomData;
17
18use crate::{
19    components::{
20        ConvolutionOperation,
21        global::{args::RuntimeArgs, read::strategy::sync_bias::SyncBiasLoading},
22    },
23    routines::{Routine, contiguous_pitched_layout, into_tensor_handle_tma},
24};
25
26/// Cmma convolution with a partial async loading strategy.
27pub struct SpecializedConv<L: AsyncPartialLoadingStrategy<RuntimeArgs>> {
28    _loader: PhantomData<L>,
29}
30
31pub type SpecializedAsyncCyclicConv =
32    SpecializedConv<AsyncPartialCyclicLoading<ColMajorTilingOrder>>;
33pub type SpecializedAsyncStridedConv = SpecializedConv<AsyncPartialStridedLoading>;
34
35pub struct SpecializedTmaConv;
36
37impl<L: AsyncPartialLoadingStrategy<RuntimeArgs>> Routine for SpecializedConv<L> {
38    type Blueprint = BatchMatmulBlueprint;
39    type Strategy = SpecializedStrategy;
40    type MatmulRoutine = SpecializedAlgorithm<L, SyncBiasLoading>;
41    type Args = TensorArgs<RuntimeArgs>;
42    const IS_SPECIALIZED: bool = true;
43
44    fn correct_layout(
45        client: &Client,
46        handle: TensorBinding,
47        dtype: ElemType,
48        _operation: ConvolutionOperation,
49    ) -> Result<TensorBinding, LaunchError> {
50        contiguous_pitched_layout(client, handle, dtype)
51    }
52}
53
54impl Routine for SpecializedTmaConv {
55    type Blueprint = BatchMatmulBlueprint;
56    type Strategy = SpecializedStrategy;
57    type MatmulRoutine = SpecializedAlgorithm<AsyncPartialTmaLoading, SyncBiasLoading>;
58    type Args = TensorMapArgs<RuntimeArgs>;
59    const IS_SPECIALIZED: bool = true;
60
61    fn correct_layout(
62        client: &Client,
63        handle: TensorBinding,
64        dtype: ElemType,
65        operation: ConvolutionOperation,
66    ) -> Result<TensorBinding, LaunchError> {
67        into_tensor_handle_tma(client, handle, dtype, operation)
68    }
69
70    fn filter_vector_sizes(vector_sizes: AvailableVectorSizes) -> AvailableVectorSizes {
71        AvailableVectorSizes {
72            lhs: vec![1],
73            rhs: vec![1],
74            out: vector_sizes.out,
75        }
76    }
77}
cubek_convolution/routines/specialized.rs

cubek_convolution/routines/
specialized.rs
