cubek_convolution/components/global/read/strategy/
async_full_cyclic.rs

1use std::marker::PhantomData;
2
3use cubecl::{
4    prelude::*,
5    std::tensor::layout::{Layout, LayoutExpand},
6    {ir::DeviceProperties, prelude::barrier::Barrier},
7};
8use cubek_matmul::multi_level::tile::{ContiguousTilingLayout, TilingOrder};
9use cubek_matmul::{
10    definition::{MatmulElems, MatmulProblem},
11    multi_level::components::{
12        global::{
13            GlobalReaderConfig, PlaneFlowPartition,
14            memory::GlobalIterator,
15            multi_stage::LoadMaxRoundPlaneCount,
16            read::{
17                FullLoadingStrategy, LoadingJob, LoadingValidation, ReaderMode,
18                async_barrier::AsyncCopy,
19                async_full_cyclic::AsyncFullCyclicLoading as MatmulCyclicLoading,
20                tiled::TiledLayout,
21            },
22        },
23        stage::{StridedStageFamily, StridedStageMemory},
24    },
25};
26use cubek_std::{InvalidConfigError, StageIdent};
27
28use crate::components::global::{
29    args::RuntimeArgs,
30    read::strategy::async_copy::{ASYNC_COPY_WIDTH, async_copy_from},
31};
32
33#[derive(CubeType, Clone, Copy)]
34/// Loads the content of all tiles in the stage using all planes.
35/// Unit with pos X loads vectors with indices X, X + NUM_UNITS, X + 2 * NUM_UNITS, ...
36pub struct AsyncFullCyclicLoading<T: TilingOrder> {
37    #[cube(comptime)]
38    _t: PhantomData<T>,
39}
40
41impl<TO: TilingOrder> LoadingValidation for AsyncFullCyclicLoading<TO> {
42    fn validate_with_config(
43        device_props: &DeviceProperties,
44        config: &GlobalReaderConfig,
45    ) -> Result<(), InvalidConfigError> {
46        MatmulCyclicLoading::<TO>::validate_with_config(device_props, config)
47    }
48
49    fn validate_with_problem(
50        problem: &MatmulProblem,
51        dtypes: &MatmulElems,
52        ident: StageIdent,
53    ) -> Result<(), InvalidConfigError> {
54        MatmulCyclicLoading::<TO>::validate_with_problem(problem, dtypes, ident)
55    }
56}
57
58impl<TO: TilingOrder> LoadMaxRoundPlaneCount for AsyncFullCyclicLoading<TO> {
59    fn max_round_plane_count(
60        elements_per_tile: u32,
61        tiles_per_stage: u32,
62        vector_size: VectorSize,
63        plane_dim: u32,
64        dtype: ElemType,
65    ) -> u32 {
66        MatmulCyclicLoading::<TO>::max_round_plane_count(
67            elements_per_tile,
68            tiles_per_stage,
69            vector_size,
70            plane_dim,
71            dtype,
72        )
73    }
74}
75
76#[cube]
77impl<TO: TilingOrder> FullLoadingStrategy<RuntimeArgs> for AsyncFullCyclicLoading<TO> {
78    type TilingLayout = ContiguousTilingLayout<TO>;
79    type SyncStrategy = AsyncCopy;
80    type Job<EG: Numeric, NG: Size, ES: Numeric, NS: Size> = AsyncFullCyclicJob;
81    type Stage = StridedStageFamily;
82    fn new_job<EG: Numeric, NG: Size, ES: Numeric, NS: Size>(
83        runtime_args: RuntimeArgs,
84        #[comptime] config: GlobalReaderConfig,
85    ) -> Self::Job<EG, NG, ES, NS> {
86        let type_size = ES::size_bits().comptime();
87        let vector_size = ASYNC_COPY_WIDTH / type_size as u32;
88        let tile_num_elements = config.smem_config.elements_per_tile();
89        let num_stage_elements = config.smem_config.elements_per_stage();
90
91        let num_stage_vectors = num_stage_elements.div_ceil(vector_size);
92        let total_units = config.loading_units_count();
93        let num_tasks_per_unit = num_stage_vectors.div_ceil(total_units);
94        let balanced_workload = num_stage_vectors.is_multiple_of(total_units);
95        let jump_length = total_units * vector_size;
96
97        let unit_id = PlaneFlowPartition::new(config.plane_flow_config.partition_rule)
98            .load_index(config.input_load_flow)
99            * config.plane_dim
100            + UNIT_POS_X;
101        let unit_position_base = unit_id * vector_size;
102
103        AsyncFullCyclicJob {
104            unit_position_base,
105            runtime_args,
106            num_tasks_per_unit,
107            tile_num_elements,
108            jump_length,
109            copy_vector_size: vector_size,
110            balanced_workload,
111            num_stage_elements,
112            reader_mode: config.reader_mode,
113        }
114    }
115}
116
117#[derive(CubeType, Clone)]
118#[expand(derive(Clone))]
119pub struct AsyncFullCyclicJob {
120    unit_position_base: u32,
121    runtime_args: RuntimeArgs,
122
123    #[cube(comptime)]
124    num_tasks_per_unit: u32,
125    #[cube(comptime)]
126    tile_num_elements: u32,
127    #[cube(comptime)]
128    jump_length: u32,
129    #[cube(comptime)]
130    copy_vector_size: u32,
131    #[cube(comptime)]
132    balanced_workload: bool,
133    #[cube(comptime)]
134    num_stage_elements: u32,
135    #[cube(comptime)]
136    reader_mode: ReaderMode,
137}
138
139#[cube]
140impl<EG: Numeric, NG: Size, ES: Numeric, NS: Size, TO: TilingOrder>
141    LoadingJob<EG, NG, ES, NS, ContiguousTilingLayout<TO>, AsyncCopy> for AsyncFullCyclicJob
142{
143    type Stage = StridedStageFamily;
144
145    fn execute_task(
146        this: &mut Self,
147        #[comptime] task_id: u32,
148        global_iter: &GlobalIterator<Vector<EG, NG>>,
149        stage: &mut StridedStageMemory<ES, NS, ContiguousTilingLayout<TO>>,
150        _barrier: &Shared<Barrier>,
151        #[comptime] config: GlobalReaderConfig,
152    ) {
153        let unit_position = this.unit_position_base + task_id * this.jump_length;
154
155        #[allow(clippy::collapsible_else_if)]
156        if comptime!(this.reader_mode == ReaderMode::Strict || this.balanced_workload) {
157            copy_vector::<EG, NG, ES, NS, TO>(
158                &*this,
159                unit_position,
160                global_iter,
161                stage,
162                &this.runtime_args,
163                config,
164            );
165        } else {
166            if unit_position < this.num_stage_elements {
167                copy_vector::<EG, NG, ES, NS, TO>(
168                    &*this,
169                    unit_position,
170                    global_iter,
171                    stage,
172                    &this.runtime_args,
173                    config,
174                );
175            }
176        }
177    }
178
179    fn task_count(this: &Self) -> comptime_type!(u32) {
180        this.num_tasks_per_unit
181    }
182}
183
184#[cube]
185pub(crate) fn copy_vector<EG: Numeric, NG: Size, ES: Numeric, NS: Size, TO: TilingOrder>(
186    job: &AsyncFullCyclicJob,
187    unit_position: u32,
188    global_iter: &GlobalIterator<Vector<EG, NG>>,
189    stage: &mut StridedStageMemory<ES, NS, ContiguousTilingLayout<TO>>,
190    runtime_args: &RuntimeArgs,
191    #[comptime] config: GlobalReaderConfig,
192) {
193    let nth_tile = unit_position / job.tile_num_elements;
194    let pos_within_tile = unit_position % job.tile_num_elements;
195
196    let layout = TiledLayout::new(config.stage_ident, config.smem_config);
197    let view = global_iter.view();
198
199    let tile = ContiguousTilingLayout::<TO>::to_x_y(nth_tile, config.smem_config);
200
201    let pos = layout.to_source_pos((tile, pos_within_tile));
202    let stage_offset = unit_position / stage.smem.vector_size() as u32;
203
204    async_copy_from(
205        view,
206        pos,
207        stage,
208        stage_offset,
209        runtime_args,
210        global_iter.offset(),
211        config,
212        job.copy_vector_size,
213    );
214}
cubek_convolution/components/global/read/strategy/async_full_cyclic.rs

cubek_convolution/components/global/read/strategy/
async_full_cyclic.rs
