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

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