cubek_convolution/routines/
simple.rs

1use cubecl::{client::Client, ir::ElemType, prelude::TensorBinding, server::LaunchError};
2use cubek_matmul::multi_level::tile::{ColMajorTilingOrder, RowMajorTilingOrder};
3use cubek_matmul::{
4    definition::AvailableVectorSizes,
5    multi_level::{
6        args::{TensorArgs, TensorMapArgs},
7        components::global::read::{
8            FullLoadingStrategy, async_full_tma::AsyncFullTmaLoading,
9            sync_full_cyclic::SyncFullCyclicLoading, sync_full_strided::SyncFullStridedLoading,
10            sync_full_tilewise::SyncFullTilewiseLoading,
11        },
12        definition::BatchMatmulBlueprint,
13        routines::batch::simple::{SimpleAlgorithm, SimpleArgs},
14    },
15};
16use std::marker::PhantomData;
17
18use crate::{
19    components::{
20        ConvolutionOperation,
21        global::{
22            args::RuntimeArgs,
23            read::strategy::{
24                async_full_cyclic::AsyncFullCyclicLoading,
25                async_full_strided::AsyncFullStridedLoading, sync_bias::SyncBiasLoading,
26            },
27        },
28    },
29    routines::{Routine, contiguous_pitched_layout, into_tensor_handle_tma},
30};
31
32/// Cmma convolution
33pub struct SimpleConv<LL: FullLoadingStrategy<RuntimeArgs>, LR: FullLoadingStrategy<RuntimeArgs>> {
34    _loader: PhantomData<(LL, LR)>,
35}
36
37pub type SimpleSyncCyclicConv = SimpleConv<
38    SyncFullCyclicLoading<RowMajorTilingOrder>,
39    SyncFullCyclicLoading<ColMajorTilingOrder>,
40>;
41pub type SimpleSyncStridedConv = SimpleConv<SyncFullStridedLoading, SyncFullStridedLoading>;
42pub type SimpleSyncTilewiseConv = SimpleConv<
43    SyncFullTilewiseLoading<RowMajorTilingOrder>,
44    SyncFullTilewiseLoading<ColMajorTilingOrder>,
45>;
46pub type SimpleAsyncCyclicConv = SimpleConv<
47    AsyncFullCyclicLoading<RowMajorTilingOrder>,
48    AsyncFullCyclicLoading<ColMajorTilingOrder>,
49>;
50pub type SimpleAsyncStridedConv = SimpleConv<AsyncFullStridedLoading, AsyncFullStridedLoading>;
51
52pub struct SimpleAsyncTmaConv;
53
54impl<
55    LL: FullLoadingStrategy<RuntimeArgs>,
56    LR: FullLoadingStrategy<RuntimeArgs, SyncStrategy = LL::SyncStrategy>,
57> Routine for SimpleConv<LL, LR>
58{
59    type Blueprint = BatchMatmulBlueprint;
60    type Strategy = SimpleArgs;
61    type MatmulRoutine = SimpleAlgorithm<LL, LR, SyncBiasLoading>;
62    type Args = TensorArgs<RuntimeArgs>;
63
64    fn correct_layout(
65        client: &Client,
66        handle: TensorBinding,
67        dtype: ElemType,
68        _operation: ConvolutionOperation,
69    ) -> Result<TensorBinding, LaunchError> {
70        contiguous_pitched_layout(client, handle, dtype)
71    }
72}
73
74impl Routine for SimpleAsyncTmaConv {
75    type Blueprint = BatchMatmulBlueprint;
76    type Strategy = SimpleArgs;
77    type MatmulRoutine = SimpleAlgorithm<AsyncFullTmaLoading, AsyncFullTmaLoading, SyncBiasLoading>;
78    type Args = TensorMapArgs<RuntimeArgs>;
79
80    fn correct_layout(
81        client: &Client,
82        handle: TensorBinding,
83        dtype: ElemType,
84        operation: ConvolutionOperation,
85    ) -> Result<TensorBinding, LaunchError> {
86        into_tensor_handle_tma(client, handle, dtype, operation)
87    }
88
89    fn filter_vector_sizes(vector_sizes: AvailableVectorSizes) -> AvailableVectorSizes {
90        AvailableVectorSizes {
91            lhs: vec![1],
92            rhs: vec![1],
93            out: vector_sizes.out,
94        }
95    }
96}
cubek_convolution/routines/simple.rs

cubek_convolution/routines/
simple.rs
