rudnn/convolution/components/global/read/strategy/
async_full_cyclic.rs

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

rudnn/convolution/components/global/read/strategy/
async_full_cyclic.rs
