rublas 0.1.45

Linear algebra and matrix multiplication for Ruda.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use ruda_kernel::dsl as kernel_dsl;
use crate::kernel_ir::{
    components::global::read::{FullLoadingStrategy, stage::FullStageLayout},
    components::global::{GlobalReaderConfig, PlaneFlowPartition},
    components::global::{multi_stage::LoadMaxRoundPlaneCount, read::sync::Synchronous},
    components::stage::StridedStageFamily,
    components::stage::{StridedStageMemory, StridedTilingLayout},
    components::{global::memory::GlobalIterator, stage::TilingValidation},
    definition::{MatmulElems, MatmulProblem, StageIdent},
    {components::global::read::validate_swizzle_atom_size, launch::RuntimeConfig},
};
use ruda_kernel::dsl::ir::DeviceProperties;
use ruda_kernel::dsl::prelude::*;
use ruda_kernel::tiling::{InvalidConfigError, tile::Strided};

use super::{LoadingJob, LoadingValidation};

#[derive(RudaType, Clone, Copy)]
/// Loads the content of all the stage using all planes,
/// keeping the original layout, making each tile strided
pub struct SyncFullStridedLoading {}

impl LoadingValidation for SyncFullStridedLoading {
    fn validate_with_config(
        _device_props: &DeviceProperties,
        config: &GlobalReaderConfig,
    ) -> Result<(), InvalidConfigError> {
        let vector_size = config.gmem_config.vector_size;

        if !config.smem_config.elements_per_stage().is_multiple_of(vector_size as u32) {
            return Err(Box::new("Stage size must be divisible by the load vector size"));
        }

        validate_swizzle_atom_size(config.smem_config)?;
        StridedTilingLayout::check(config.smem_config)?;

        Ok(())
    }

    fn validate_with_problem(
        _problem: &MatmulProblem,
        _dtypes: &MatmulElems,
        _ident: StageIdent,
    ) -> Result<(), InvalidConfigError> {
        Ok(())
    }
}

impl LoadMaxRoundPlaneCount for SyncFullStridedLoading {
    fn max_round_plane_count(
        elements_per_tile: u32,
        tiles_per_stage: u32,
        vector_size: VectorSize,
        plane_dim: u32,
        _dtype: StorageType,
    ) -> u32 {
        let elements_per_stage = elements_per_tile * tiles_per_stage;
        let num_vectors = elements_per_stage / vector_size as u32;
        num_vectors.div_ceil(plane_dim)
    }
}

#[ruda]
impl<RC: RuntimeConfig> FullLoadingStrategy<RC> for SyncFullStridedLoading {
    type TilingLayout = StridedTilingLayout;
    type SyncStrategy = Synchronous;
    type Job<EG: Numeric, NG: Size, ES: Numeric, NS: Size> = SyncFullStridedJob;
    type Stage = StridedStageFamily;
    type TileKind = Strided;

    fn new_job<EG: Numeric, NG: Size, ES: Numeric, NS: Size>(
        _runtime_config: RC,
        #[comptime] config: GlobalReaderConfig,
    ) -> Self::Job<EG, NG, ES, NS> {
        let vector_size = NG::value().comptime() as u32;
        let num_stage_vectors = config.smem_config.elements_per_stage() / vector_size;
        let unit_count = config.loading_planes_count() * config.plane_dim;
        let num_tasks_per_unit = num_stage_vectors.div_ceil(unit_count);

        let unit_position_base = PlaneFlowPartition::new(config.plane_flow_config.partition_rule)
            .load_index(config.input_load_flow)
            * config.plane_dim
            + UNIT_POS_X;

        SyncFullStridedJob {
            unit_position_base,
            num_tasks_per_unit,
            unit_count,
        }
    }
}

#[derive(RudaType, Clone, Copy)]
pub struct SyncFullStridedJob {
    unit_position_base: u32,

    #[ruda(comptime)]
    num_tasks_per_unit: u32,
    #[ruda(comptime)]
    unit_count: u32,
}

#[ruda]
impl<EG: Numeric, NG: Size, ES: Numeric, NS: Size>
    LoadingJob<EG, NG, ES, NS, StridedTilingLayout, Synchronous> for SyncFullStridedJob
{
    type Stage = StridedStageFamily;

    fn execute_task(
        this: &mut Self,
        #[comptime] task_id: u32,
        global_iter: &GlobalIterator<Vector<EG, NG>>,
        stage: &mut StridedStageMemory<ES, NS, StridedTilingLayout>,
        _barrier: &mut (),
        #[comptime] config: GlobalReaderConfig,
    ) {
        let unit_position = this.unit_position_base + task_id * this.unit_count;
        if unit_position < config.smem_config.elements_per_stage() / NG::value() as u32 {
            let layout = FullStageLayout::new(config.smem_config);
            let view = global_iter.view().view(layout);
            let vector_read = view.read_checked(unit_position * NG::value() as u32);
            let type_size = Vector::<ES, NS>::type_size();
            let stage_offs = stage.swizzle.apply(unit_position, type_size);
            stage.as_slice_mut::<NS>()[stage_offs as usize] = Vector::cast_from(vector_read);
        }
    }

    fn task_count(this: &Self) -> comptime_type!(u32) {
        this.num_tasks_per_unit
    }
}