cubecl_matmul/
tune_key.rs

1use cubecl_core::client::ComputeClient;
2use cubecl_core::{self as cubecl, Runtime};
3
4use cubecl_core::{AutotuneKey, ir::ElemType};
5use serde::{Deserialize, Serialize};
6
7use cubecl_std::tensor::{MatrixBatchLayout, matrix_batch_layout};
8
9use super::components::{MatmulKind, MatmulProblemSize};
10
11#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)]
12/// Autotune key representative of matmul versions
13pub struct MatmulAutotuneKey {
14    pub definition: MatmulProblemDefinition,
15    pub analysis: MatmulAutotuneAnalysis,
16}
17
18#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)]
19pub struct MatmulProblemDefinition {
20    #[autotune(anchor)]
21    pub m: usize,
22    #[autotune(anchor)]
23    pub n: usize,
24    #[autotune(anchor)]
25    pub k: usize,
26    pub lhs_pow2_factor: u8,
27    pub rhs_pow2_factor: u8,
28    pub elem_lhs: MatmulElemType,
29    pub elem_rhs: MatmulElemType,
30    pub elem_out: MatmulElemType,
31    pub matrix_layout_lhs: MatrixBatchLayout,
32    pub matrix_layout_rhs: MatrixBatchLayout,
33}
34
35#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
36pub enum MatmulGlobalScale {
37    Large,
38    Medium,
39    Small,
40}
41
42#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize)]
43pub struct MatmulAutotuneAnalysis {
44    pub scale_global: MatmulGlobalScale,
45    pub kind: MatmulKind,
46}
47
48impl MatmulGlobalScale {
49    pub fn from_size(m: usize, n: usize, k: usize) -> Self {
50        if m < 512 && k < 512 && n < 512 {
51            MatmulGlobalScale::Small
52        } else if m < 2048 && k < 2048 && n < 2048 {
53            MatmulGlobalScale::Medium
54        } else {
55            MatmulGlobalScale::Large
56        }
57    }
58}
59
60/// Whether it's a good idea to try and run double-buffered matmul.
61pub fn should_tune_double_buffering(fused: bool, key: &MatmulAutotuneKey) -> bool {
62    matches!(key.analysis.kind, MatmulKind::General)
63        && match key.analysis.scale_global {
64            MatmulGlobalScale::Large => true,
65            MatmulGlobalScale::Medium => true,
66            MatmulGlobalScale::Small => fused,
67        }
68}
69
70#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)]
71pub struct MatmulElemType {
72    pub elem: ElemType,
73    pub quantized: bool,
74}
75
76impl MatmulAutotuneKey {
77    /// Create the autotune key based on the shape of both lhs and rhs as well as the element type
78    /// used for the calculation.
79    #[allow(clippy::too_many_arguments)]
80    pub fn generate<R: Runtime>(
81        _client: &ComputeClient<R::Server>,
82        lhs_shape: &[usize],
83        rhs_shape: &[usize],
84        lhs_strides: &[usize],
85        rhs_strides: &[usize],
86        elem_lhs: MatmulElemType,
87        elem_rhs: MatmulElemType,
88        elem_out: MatmulElemType,
89    ) -> MatmulAutotuneKey {
90        let ndims = lhs_shape.len();
91        let m = lhs_shape[ndims - 2];
92        let k = lhs_shape[ndims - 1];
93        let n = rhs_shape[ndims - 1];
94
95        let matrix_layout_lhs = matrix_batch_layout(lhs_strides);
96        let matrix_layout_rhs = matrix_batch_layout(rhs_strides);
97
98        let kind = MatmulKind::from(MatmulProblemSize {
99            m: m as u32,
100            n: n as u32,
101            k: k as u32,
102        });
103
104        let lhs_pow2_factor = match matrix_layout_lhs {
105            MatrixBatchLayout::Contiguous => pow2_factor(k),
106            MatrixBatchLayout::MildlyPermuted { transposed, .. } => match transposed {
107                true => pow2_factor(m),
108                false => pow2_factor(k),
109            },
110            MatrixBatchLayout::HighlyPermuted => 0,
111        };
112        let rhs_pow2_factor = match matrix_layout_rhs {
113            MatrixBatchLayout::Contiguous => pow2_factor(n),
114            MatrixBatchLayout::MildlyPermuted { transposed, .. } => match transposed {
115                true => pow2_factor(k),
116                false => pow2_factor(n),
117            },
118            MatrixBatchLayout::HighlyPermuted => 0,
119        };
120
121        let definition = MatmulProblemDefinition::new(
122            m,
123            n,
124            k,
125            lhs_pow2_factor,
126            rhs_pow2_factor,
127            elem_lhs,
128            elem_rhs,
129            elem_out,
130            matrix_layout_lhs,
131            matrix_layout_rhs,
132        );
133        let analysis = MatmulAutotuneAnalysis {
134            scale_global: MatmulGlobalScale::from_size(m, n, k),
135            kind,
136        };
137
138        Self::new(definition, analysis)
139    }
140}
141
142/// Defines the potential vectorization.
143fn pow2_factor(axis: usize) -> u8 {
144    for i in (1..4).rev() {
145        if axis.is_multiple_of(2usize.pow(i as u32)) {
146            return i;
147        }
148    }
149
150    0
151}
cubecl_matmul/tune_key.rs

cubecl_matmul/
tune_key.rs
