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

1use cubecl::{
2    prelude::*,
3    std::tensor::layout::{Layout, LayoutExpand},
4    {ir::DeviceProperties, prelude::barrier::Barrier},
5};
6use cubek_matmul::multi_level::tile::StridedTilingLayout;
7use cubek_matmul::{
8    definition::{MatmulElems, MatmulProblem},
9    multi_level::components::{
10        global::{
11            GlobalReaderConfig, PlaneFlowPartition,
12            memory::GlobalIterator,
13            multi_stage::LoadMaxRoundPlaneCount,
14            read::{
15                FullLoadingStrategy, LoadingJob, LoadingValidation, async_barrier::AsyncCopy,
16                async_full_strided::AsyncFullStridedLoading as MatmulStridedLoading,
17                stage::FullStageLayout,
18            },
19        },
20        stage::{StridedStageFamily, StridedStageMemory},
21    },
22};
23use cubek_std::{InvalidConfigError, StageIdent};
24
25use crate::components::global::{
26    args::RuntimeArgs,
27    read::strategy::async_copy::{ASYNC_COPY_WIDTH, async_copy_from},
28};
29
30#[derive(CubeType, Clone, Copy)]
31/// Loads the content of all the stage using all planes,
32/// keeping the original layout, making each tile strided
33pub struct AsyncFullStridedLoading {}
34
35impl LoadingValidation for AsyncFullStridedLoading {
36    fn validate_with_config(
37        device_props: &DeviceProperties,
38        config: &GlobalReaderConfig,
39    ) -> Result<(), InvalidConfigError> {
40        MatmulStridedLoading::validate_with_config(device_props, config)
41    }
42
43    fn validate_with_problem(
44        problem: &MatmulProblem,
45        dtypes: &MatmulElems,
46        ident: StageIdent,
47    ) -> Result<(), InvalidConfigError> {
48        MatmulStridedLoading::validate_with_problem(problem, dtypes, ident)
49    }
50}
51
52impl LoadMaxRoundPlaneCount for AsyncFullStridedLoading {
53    fn max_round_plane_count(
54        elements_per_tile: u32,
55        tiles_per_stage: u32,
56        vector_size: VectorSize,
57        plane_dim: u32,
58        dtype: ElemType,
59    ) -> u32 {
60        MatmulStridedLoading::max_round_plane_count(
61            elements_per_tile,
62            tiles_per_stage,
63            vector_size,
64            plane_dim,
65            dtype,
66        )
67    }
68}
69
70#[cube]
71impl FullLoadingStrategy<RuntimeArgs> for AsyncFullStridedLoading {
72    type TilingLayout = StridedTilingLayout;
73    type SyncStrategy = AsyncCopy;
74    type Job<EG: Numeric, NG: Size, ES: Numeric, NS: Size> = AsyncFullStridedJob;
75    type Stage = StridedStageFamily;
76    fn new_job<EG: Numeric, NG: Size, ES: Numeric, NS: Size>(
77        runtime_args: RuntimeArgs,
78        #[comptime] config: GlobalReaderConfig,
79    ) -> Self::Job<EG, NG, ES, NS> {
80        let type_size = ES::size_bits().comptime();
81        let vector_size = ASYNC_COPY_WIDTH / type_size as u32;
82        let num_stage_vectors = config.smem_config.elements_per_stage() / vector_size;
83        let unit_count = config.loading_planes_count() * config.plane_dim;
84        let num_tasks_per_unit = num_stage_vectors / unit_count;
85
86        let unit_position_base = PlaneFlowPartition::new(config.plane_flow_config.partition_rule)
87            .load_index(config.input_load_flow)
88            * config.plane_dim
89            + UNIT_POS_X;
90
91        AsyncFullStridedJob {
92            unit_position_base,
93            runtime_args,
94            num_tasks_per_unit,
95            unit_count,
96            copy_vector_size: vector_size,
97        }
98    }
99}
100
101#[derive(CubeType, Clone)]
102#[expand(derive(Clone))]
103pub struct AsyncFullStridedJob {
104    unit_position_base: u32,
105    runtime_args: RuntimeArgs,
106
107    #[cube(comptime)]
108    num_tasks_per_unit: u32,
109    #[cube(comptime)]
110    unit_count: u32,
111    #[cube(comptime)]
112    copy_vector_size: u32,
113}
114
115#[cube]
116impl<EG: Numeric, NG: Size, ES: Numeric, NS: Size>
117    LoadingJob<EG, NG, ES, NS, StridedTilingLayout, AsyncCopy> for AsyncFullStridedJob
118{
119    type Stage = StridedStageFamily;
120
121    fn execute_task(
122        this: &mut Self,
123        #[comptime] task_id: u32,
124        global_iter: &GlobalIterator<Vector<EG, NG>>,
125        stage: &mut StridedStageMemory<ES, NS, StridedTilingLayout>,
126        _barrier: &Shared<Barrier>,
127        #[comptime] config: GlobalReaderConfig,
128    ) {
129        let unit_position = this.unit_position_base + task_id * this.unit_count;
130        let unit_position_abs = unit_position * this.copy_vector_size;
131
132        let layout = FullStageLayout::new(config.smem_config);
133        let view = global_iter.view();
134
135        let pos = layout.to_source_pos(unit_position_abs);
136        let stage_offset = unit_position_abs / stage.smem.vector_size() as u32;
137
138        async_copy_from(
139            view,
140            pos,
141            stage,
142            stage_offset,
143            &this.runtime_args,
144            global_iter.offset(),
145            config,
146            this.copy_vector_size,
147        );
148    }
149
150    fn task_count(this: &Self) -> comptime_type!(u32) {
151        this.num_tasks_per_unit
152    }
153}
cubek_convolution/components/global/read/strategy/async_full_strided.rs

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