cubecl-matmul 0.8.1

CubeCL Matrix Multiplication Kernels Engine
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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use crate::components::{
    InvalidConfigError, MatmulIdent, MatrixLayout, MatrixPrecision,
    global::{
        CopyMechanism, GlobalConfig,
        memory::{GlobalIterator, load_window_in_stage},
        read::AsyncFullLoadingStrategy,
    },
    stage::{StageConfig, StridedStage, StridedTilingLayout, TilingValidation},
};
use cubecl_core::prelude::*;
use cubecl_core::{self as cubecl, prelude::barrier::BarrierLevel};

use super::{AsyncLoadingJob, LoadingValidation};

#[derive(CubeType, Clone, Copy)]
/// Executes one memcpy_async call per unit.
/// The objective is to reduce branching, prioritizing this over maximizing memory slice length.
pub struct AsyncFullMaximizeUnitCountLoading {}

impl LoadingValidation for AsyncFullMaximizeUnitCountLoading {
    fn check<C: GlobalConfig>(config: &C, ident: MatmulIdent) -> Result<(), InvalidConfigError> {
        let matrix_layout = config.matrix_layout(ident);
        let line_size = config.global_line_size(ident);

        let (num_slices, slice_length) = match matrix_layout {
            MatrixLayout::RowMajor => (
                config.tiling_scheme().elements_in_stage_row(ident),
                config.tiling_scheme().elements_in_stage_col(ident) / line_size,
            ),
            MatrixLayout::ColMajor => (
                config.tiling_scheme().elements_in_stage_col(ident),
                config.tiling_scheme().elements_in_stage_row(ident) / line_size,
            ),
        };
        let unit_count = config.plane_dim() * config.num_loading_planes(ident);

        if !unit_count.is_multiple_of(num_slices) {
            return Err(Box::new(
                "Number of slices must divide number of units evenly",
            ));
        }
        if slice_length % (unit_count / num_slices) != 0 {
            return Err(Box::new(
                "Number of units per slice must divide slice length evenly",
            ));
        }

        StridedTilingLayout::check(config.global_memory_config(ident))?;

        Ok(())
    }
}

#[cube]
impl AsyncFullLoadingStrategy for AsyncFullMaximizeUnitCountLoading {
    type TilingLayout = StridedTilingLayout;
    type Job<IP: MatrixPrecision> = AsyncFullMaximizeUnitCountJob;

    fn new_job<IP: MatrixPrecision, G: GlobalConfig>(
        #[comptime] ident: MatmulIdent,
        #[comptime] config: G,
    ) -> AsyncFullMaximizeUnitCountJob {
        let matrix_layout = config.matrix_layout(ident);
        let line_size = config
            .stage_config()
            .stage_line_size(comptime!(ident.into_stage()));

        let (num_slices, slice_length) = match matrix_layout {
            MatrixLayout::RowMajor => (
                config.tiling_scheme().elements_in_stage_row(ident),
                config.tiling_scheme().elements_in_stage_col(ident) / line_size,
            ),
            MatrixLayout::ColMajor => (
                config.tiling_scheme().elements_in_stage_col(ident),
                config.tiling_scheme().elements_in_stage_row(ident) / line_size,
            ),
        };

        let unit_count = config.plane_dim() * config.num_loading_planes(ident);

        let units_per_slice = comptime!(unit_count / num_slices);
        let nth_slice = UNIT_POS / units_per_slice;

        let segment_length = comptime!(slice_length / units_per_slice);
        let nth_segment = UNIT_POS % units_per_slice;

        AsyncFullMaximizeUnitCountJob {
            nth_slice,
            nth_segment,
            segment_length,
            ident,
        }
    }

    fn barrier_level() -> BarrierLevel {
        BarrierLevel::cube_manual(0u32)
    }
}

#[derive(CubeType, Clone, Copy)]
pub struct AsyncFullMaximizeUnitCountJob {
    nth_slice: u32,
    nth_segment: u32,
    #[cube(comptime)]
    segment_length: u32,
    #[cube(comptime)]
    ident: MatmulIdent,
}

#[cube]
impl<IP: MatrixPrecision> AsyncLoadingJob<IP, StridedTilingLayout>
    for AsyncFullMaximizeUnitCountJob
{
    fn execute_task<CM: CopyMechanism, G: GlobalConfig>(
        this: &mut Self,
        _task_id: u32,
        global_iter: &GlobalIterator<Line<IP::Global>>,
        stage: &mut StridedStage<IP::Stage, StridedTilingLayout>,
        mechanism: &CM,
        #[comptime] config: G,
    ) {
        let mut destination: SliceMut<Line<IP::Stage>> = StridedTilingLayout::nth_slice::<IP::Stage>(
            stage,
            this.nth_slice,
            comptime!(config.stage_memory_config(this.ident)),
        );

        let window = load_window_in_stage(
            &global_iter.view(),
            this.nth_slice,
            comptime!(config.global_memory_config(this.ident)),
        );
        let seg_start = Min::min(this.nth_segment * this.segment_length, window.len());
        let seg_end = Min::min((this.nth_segment + 1) * this.segment_length, window.len());

        let src_segment = window.slice(seg_start, seg_end);
        let mut dest_segment = destination.slice_mut(seg_start, seg_end);

        CM::memcpy_async(
            mechanism,
            &src_segment.try_cast_unchecked(),
            &mut dest_segment,
        );
    }

    fn task_count(_this: &Self) -> comptime_type!(u32) {
        1
    }
}