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

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