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

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

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