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

use crate::components::{
    InvalidConfigError, MatmulIdent, MatrixLayout, MatrixPrecision,
    global::{
        CopyMechanism, GlobalConfig,
        memory::{GlobalIterator, load_window_in_stage},
        read::AsyncFullLoadingStrategy,
    },
    stage::{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 contiguous slice.
/// The goal is to reduce the total number of memcpy_async calls, though it may result in idle threads.
pub struct AsyncFullMaximizeSliceLengthLoading {}

impl LoadingValidation for AsyncFullMaximizeSliceLengthLoading {
    fn check<C: GlobalConfig>(config: &C, ident: MatmulIdent) -> Result<(), InvalidConfigError> {
        StridedTilingLayout::check(config.global_memory_config(ident))?;

        Ok(())
    }
}

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

    fn new_job<IP: MatrixPrecision, G: GlobalConfig>(
        #[comptime] ident: MatmulIdent,
        #[comptime] config: G,
    ) -> AsynFullMaximizeSliceLengthJob {
        let matrix_layout = config.matrix_layout(ident);

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

        let num_tasks_per_unit = comptime!(div_ceil(num_slices, unit_count));

        AsynFullMaximizeSliceLengthJob {
            num_tasks_per_unit,
            unit_count,
            num_slices,
            ident,
        }
    }

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

#[derive(CubeType, Clone, Copy)]
pub struct AsynFullMaximizeSliceLengthJob {
    #[cube(comptime)]
    num_tasks_per_unit: u32,
    #[cube(comptime)]
    unit_count: u32,
    #[cube(comptime)]
    num_slices: u32,
    #[cube(comptime)]
    ident: MatmulIdent,
}

#[cube]
impl<IP: MatrixPrecision> AsyncLoadingJob<IP, StridedTilingLayout>
    for AsynFullMaximizeSliceLengthJob
{
    fn execute_task<CM: CopyMechanism, G: GlobalConfig>(
        this: &mut Self,
        task_id: u32,
        tensor_reader: &GlobalIterator<Line<IP::Global>>,
        stage: &mut StridedStage<IP::Stage, StridedTilingLayout>,
        mechanism: &CM,
        #[comptime] config: G,
    ) {
        let nth_slice = this.unit_count * task_id + UNIT_POS;

        #[allow(clippy::collapsible_else_if)]
        if comptime!(this.num_slices.is_multiple_of(this.unit_count)) {
            load_nth_slice::<IP::Global, IP::Stage, CM, G>(
                nth_slice,
                tensor_reader,
                stage,
                mechanism,
                this.ident,
                config,
            );
        } else {
            if nth_slice < this.num_slices {
                load_nth_slice::<IP::Global, IP::Stage, CM, G>(
                    nth_slice,
                    tensor_reader,
                    stage,
                    mechanism,
                    this.ident,
                    config,
                );
            }
        };
    }

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

#[cube]
fn load_nth_slice<EG: Numeric, ES: Numeric, CM: CopyMechanism, G: GlobalConfig>(
    nth_slice: u32,
    global_iter: &GlobalIterator<Line<EG>>,
    stage: &mut StridedStage<ES, StridedTilingLayout>,
    mechanism: &CM,
    #[comptime] ident: MatmulIdent,
    #[comptime] config: G,
) {
    let window = load_window_in_stage(
        &global_iter.view(),
        nth_slice,
        comptime!(config.global_memory_config(ident)),
    );
    let mut destination: SliceMut<Line<ES>> = StridedTilingLayout::nth_slice::<ES>(
        stage,
        nth_slice,
        comptime!(config.stage_memory_config(ident)),
    );

    CM::memcpy_async(mechanism, &window.try_cast_unchecked(), &mut destination);
}