cubek_convolution/components/global/read/reader/strategy/
base.rs

1use cubecl::prelude::*;
2use cubek_matmul::components::{
3    global::{
4        GlobalReaderConfig,
5        read::{
6            async_full_tma::AsyncFullTmaLoading, sync_full_cyclic::SyncFullCyclicLoading,
7            sync_full_ordered::SyncFullOrderedLoading, sync_full_strided::SyncFullStridedLoading,
8            sync_full_tilewise::SyncFullTilewiseLoading,
9        },
10    },
11    stage::TilingOrder,
12};
13
14use crate::components::global::{
15    args::{RuntimeArgs, RuntimeArgsExpand},
16    read::full_reader::FullLoadingStrategy,
17};
18use cubek_matmul::components::global::read::FullLoadingStrategy as MatmulFullLoadingStrategy;
19
20macro_rules! impl_full_load_strategy {
21    ($( $($ty: ident)::* $(<$($l: lifetime,)* $($T: ident $(: $($(+)? $B: ident)*)?),+>)?,)*) => {
22        $(
23            impl$(<$($l,)* $($T: CubeType $($(+ $B)*)? ),*>)? FullLoadingStrategy for $($ty)::* $(<$($l,)* $($T),+>)? {
24                type TilingLayout = <Self as MatmulFullLoadingStrategy>::TilingLayout;
25                type SyncStrategy = <Self as MatmulFullLoadingStrategy>::SyncStrategy;
26
27                type Job<EG: Numeric, ES: Numeric> =
28                    <Self as MatmulFullLoadingStrategy>::Job<EG, ES>;
29
30                fn new_job<EG: Numeric, ES: Numeric>(
31                    _runtime_args: RuntimeArgs,
32                    line_size: u32,
33                    config: GlobalReaderConfig,
34                ) -> Self::Job<EG, ES> {
35                    <Self as MatmulFullLoadingStrategy>::new_job::<EG, ES>(
36                        line_size, config,
37                    )
38                }
39
40                fn __expand_new_job<EG: Numeric, ES: Numeric>(
41                    scope: &mut Scope,
42                    _runtime_args: RuntimeArgsExpand,
43                    line_size: u32,
44                    config: GlobalReaderConfig,
45                ) -> <Self::Job<EG, ES> as CubeType>::ExpandType {
46                    <Self as MatmulFullLoadingStrategy>::__expand_new_job::<EG, ES>(
47                        scope, line_size, config,
48                    )
49                }
50            }
51        )*
52    };
53}
54
55// These work as is with the correct layout. They don't need the runtime args to function properly.
56impl_full_load_strategy!(
57    SyncFullCyclicLoading<TO: TilingOrder>,
58    SyncFullStridedLoading,
59    SyncFullOrderedLoading,
60    SyncFullTilewiseLoading<TO: TilingOrder>,
61    AsyncFullTmaLoading,
62);
cubek_convolution/components/global/read/reader/strategy/base.rs

cubek_convolution/components/global/read/reader/strategy/
base.rs
