burn_std/tensor/
matmul.rs

1use crate::tensor::{Metadata, Shape, Strides};
2
3/// Accelerated matmul tiles consume rows in chunks of this size; a row count
4/// that is a multiple of it already tiles cleanly.
5const ROW_TILE: usize = 32;
6
7/// The action to take for a batched matmul with a broadcast rhs.
8#[derive(Debug, PartialEq, Eq, Clone, Copy)]
9pub enum MatmulTransformAction {
10    /// Launch the matmul as-is.
11    Keep,
12    /// Fold every lhs/out batch dim into the rows: `[.., b, m, k] @ [.., 1, k, n]`
13    /// runs as one `[b*m, k] @ [k, n]` matmul against the shared rhs, instead of
14    /// `b` matmuls that each re-read it.
15    MergeBatches {
16        /// The merged row count (batches × rows).
17        rows: usize,
18    },
19}
20
21impl MatmulTransformAction {
22    /// Apply the action to one merged operand (lhs or out): fold the batch dims
23    /// into the rows, in place. A pure metadata rewrite — the buffer is shared.
24    ///
25    /// Requires batch-contiguous rows (see [MatmulTransformAnalysis]); the merge
26    /// keeps the rank, with every batch dim set to 1.
27    pub fn apply(&self, meta: &mut Metadata) {
28        let rows = match self {
29            MatmulTransformAction::Keep => return,
30            MatmulTransformAction::MergeBatches { rows } => *rows,
31        };
32
33        let rank = meta.rank();
34        let stride_rows = merged_row_stride(meta.shape(), meta.strides())
35            .expect("The action requires batch-contiguous rows");
36
37        for i in 0..rank - 2 {
38            meta.shape_mut()[i] = 1;
39            meta.strides_mut()[i] = rows * stride_rows;
40        }
41        meta.shape_mut()[rank - 2] = rows;
42        meta.strides_mut()[rank - 2] = stride_rows;
43    }
44}
45
46/// The facts of a batched matmul `lhs @ rhs` relevant to folding its batches
47/// into its rows.
48#[derive(Debug, PartialEq, Eq, Clone, Copy)]
49pub struct MatmulTransformAnalysis {
50    /// Product of the lhs batch dims.
51    batches: usize,
52    /// Rows of one batch (`m`).
53    rows: usize,
54    /// Columns of the output (`n`).
55    cols: usize,
56    /// The rhs is shared by every batch, and each operand's rows advance with a
57    /// single stride across batch boundaries (no holes), so the fold is a pure
58    /// view.
59    mergeable: bool,
60}
61
62impl MatmulTransformAnalysis {
63    /// Analyze from shapes alone: the operands are assumed batch-contiguous, as
64    /// holds for freshly materialized tensors. Use [Self::from_metadata] when
65    /// the layouts are known.
66    ///
67    /// The rhs may have a lower rank than the lhs (implicit broadcast).
68    pub fn from_shapes(lhs: &Shape, rhs: &Shape) -> Self {
69        Self::new(lhs, rhs, true)
70    }
71
72    /// Analyze from full metadata: on top of the shape facts, the lhs and out
73    /// rows must advance with a single stride across batch boundaries.
74    pub fn from_metadata(lhs: &Metadata, rhs: &Metadata, out: &Metadata) -> Self {
75        let mergeable = merged_row_stride(lhs.shape(), lhs.strides()).is_some()
76            && merged_row_stride(out.shape(), out.strides()).is_some();
77
78        Self::new(lhs.shape(), rhs.shape(), mergeable)
79    }
80
81    fn new(lhs: &Shape, rhs: &Shape, rows_contiguous: bool) -> Self {
82        let rank = lhs.num_dims();
83        let rank_rhs = rhs.num_dims();
84
85        let batches = lhs[..rank - 2].iter().product();
86        let rows = lhs[rank - 2];
87        let cols = rhs[rank_rhs - 1];
88
89        // Every batch dim the rhs actually has must be broadcast.
90        let rhs_shared = rhs[..rank_rhs - 2].iter().all(|&dim| dim == 1);
91
92        Self {
93            batches,
94            rows,
95            cols,
96            mergeable: rhs_shared && rows_contiguous,
97        }
98    }
99}
100
101/// Decides the [action](MatmulTransformAction) to take for a batched matmul,
102/// given its [analysis](MatmulTransformAnalysis).
103#[derive(Debug, Default, Clone, Copy)]
104pub enum MatmulTransformPolicy {
105    /// Merge the batches into the rows when the fold is a pure view and the
106    /// merged problem tiles better: per-batch rows that already fill row tiles
107    /// stay batched (a batched matmul beats one big matmul at that scale), tiny
108    /// row counts merge when that brings the output closer to square.
109    #[default]
110    BetterTiling,
111    /// Never transform.
112    Never,
113}
114
115impl MatmulTransformPolicy {
116    /// The action to take for a matmul with the given analysis.
117    pub fn action(&self, analysis: &MatmulTransformAnalysis) -> MatmulTransformAction {
118        match self {
119            MatmulTransformPolicy::Never => MatmulTransformAction::Keep,
120            MatmulTransformPolicy::BetterTiling => {
121                if !analysis.mergeable || analysis.batches == 1 {
122                    return MatmulTransformAction::Keep;
123                }
124
125                // Rows already fill the tiles: keep the batched form.
126                if analysis.rows.is_multiple_of(ROW_TILE) {
127                    return MatmulTransformAction::Keep;
128                }
129
130                let rows = analysis.batches * analysis.rows;
131
132                // The merge must bring the output closer to square.
133                if !squarer(rows, analysis.rows, analysis.cols) {
134                    return MatmulTransformAction::Keep;
135                }
136
137                MatmulTransformAction::MergeBatches { rows }
138            }
139        }
140    }
141}
142
143/// Whether an output of `rows_new` x `cols` is closer to square than
144/// `rows` x `cols`.
145fn squarer(rows_new: usize, rows: usize, cols: usize) -> bool {
146    // min(a, c) / max(a, c) > min(b, c) / max(b, c), cross-multiplied to stay
147    // in integers.
148    rows_new.min(cols) * rows.max(cols) > rows.min(cols) * rows_new.max(cols)
149}
150
151/// The single stride with which rows advance across batch boundaries, when the
152/// batch and row dims form a hole-free chain — the requirement for folding the
153/// batches into the rows as a pure view. Dims of size 1 carry arbitrary strides
154/// and never anchor the chain.
155fn merged_row_stride(shape: &Shape, strides: &Strides) -> Option<usize> {
156    let rank = shape.num_dims();
157    let mut chained: Option<(usize, usize)> = None;
158
159    for i in (0..rank - 1).rev() {
160        if shape[i] == 1 {
161            continue;
162        }
163        match chained {
164            None => chained = Some((strides[i], strides[i] * shape[i])),
165            Some((row_stride, expected)) => {
166                if strides[i] != expected {
167                    return None;
168                }
169                chained = Some((row_stride, strides[i] * shape[i]));
170            }
171        }
172    }
173
174    match chained {
175        Some((row_stride, _)) => Some(row_stride),
176        // Every merged dim is a unit dim: a single row, any stride works.
177        None => Some(shape[rank - 1] * strides[rank - 1]),
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    use crate::shape;
185
186    fn analysis(lhs: &[usize], rhs: &[usize]) -> MatmulTransformAnalysis {
187        MatmulTransformAnalysis::from_shapes(&Shape::from(lhs), &Shape::from(rhs))
188    }
189
190    fn action(analysis: &MatmulTransformAnalysis) -> MatmulTransformAction {
191        MatmulTransformPolicy::default().action(analysis)
192    }
193
194    #[test]
195    fn merges_batched_vec_mat() {
196        // A decode step: 16 broadcast vec-mats fold into one 16-row matmul.
197        let analysis = analysis(&[16, 1, 4096], &[4096, 14336]);
198
199        assert_eq!(
200            action(&analysis),
201            MatmulTransformAction::MergeBatches { rows: 16 }
202        );
203    }
204
205    #[test]
206    fn keeps_row_tiled_batches() {
207        // A prefill step: 512 rows per batch already fill the row tiles; the
208        // batched matmul beats one big matmul.
209        let analysis = analysis(&[8, 512, 4096], &[4096, 14336]);
210
211        assert_eq!(action(&analysis), MatmulTransformAction::Keep);
212    }
213
214    #[test]
215    fn keeps_single_batch() {
216        let analysis = analysis(&[1, 1, 4096], &[4096, 14336]);
217
218        assert_eq!(action(&analysis), MatmulTransformAction::Keep);
219    }
220
221    #[test]
222    fn keeps_matrix() {
223        let analysis = analysis(&[16, 4096], &[4096, 14336]);
224
225        assert_eq!(action(&analysis), MatmulTransformAction::Keep);
226    }
227
228    #[test]
229    fn keeps_batched_rhs() {
230        // The rhs differs per batch: nothing is shared, nothing to fold.
231        let analysis = analysis(&[8, 1, 64], &[8, 64, 32]);
232
233        assert_eq!(action(&analysis), MatmulTransformAction::Keep);
234    }
235
236    #[test]
237    fn merges_with_broadcast_rhs_rank() {
238        let analysis = analysis(&[8, 1, 64], &[1, 64, 32]);
239
240        assert_eq!(
241            action(&analysis),
242            MatmulTransformAction::MergeBatches { rows: 8 }
243        );
244    }
245
246    #[test]
247    fn keeps_when_merge_leaves_square() {
248        // 1000 rows tower over 64 cols; folding to 4000 only makes the output
249        // less square.
250        let analysis = analysis(&[4, 1000, 64], &[64, 64]);
251
252        assert_eq!(action(&analysis), MatmulTransformAction::Keep);
253    }
254
255    #[test]
256    fn merges_multiple_batch_dims() {
257        let analysis = analysis(&[2, 8, 1, 4096], &[4096, 14336]);
258
259        assert_eq!(
260            action(&analysis),
261            MatmulTransformAction::MergeBatches { rows: 16 }
262        );
263    }
264
265    #[test]
266    fn never_policy_keeps() {
267        let analysis = analysis(&[16, 1, 4096], &[4096, 14336]);
268
269        assert_eq!(
270            MatmulTransformPolicy::Never.action(&analysis),
271            MatmulTransformAction::Keep
272        );
273    }
274
275    #[test]
276    fn metadata_merges_pitched_rows() {
277        // Padded batches with one row each still fold: the merged matrix reads
278        // its rows with the single (pitched) stride 8192.
279        let lhs = Metadata::new(shape![16, 1, 4096], crate::strides![8192, 4096, 1]);
280        let rhs = Metadata::new(shape![1, 4096, 14336], crate::strides![0, 14336, 1]);
281        let out = Metadata::new(shape![16, 1, 14336], crate::strides![14336, 14336, 1]);
282
283        let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out);
284
285        assert_eq!(
286            action(&analysis),
287            MatmulTransformAction::MergeBatches { rows: 16 }
288        );
289    }
290
291    #[test]
292    fn metadata_keeps_batch_holes() {
293        // Rows are contiguous within a batch but batches leave a gap: no single
294        // row stride can express the fold.
295        let lhs = Metadata::new(shape![4, 3, 8], crate::strides![48, 8, 1]);
296        let rhs = Metadata::new(shape![1, 8, 32], crate::strides![0, 32, 1]);
297        let out = Metadata::new(shape![4, 3, 32], crate::strides![96, 32, 1]);
298
299        let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out);
300
301        assert_eq!(action(&analysis), MatmulTransformAction::Keep);
302    }
303
304    #[test]
305    fn metadata_merges_contiguous() {
306        let lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 4096, 1]);
307        let rhs = Metadata::new(shape![1, 4096, 14336], crate::strides![0, 14336, 1]);
308        let out = Metadata::new(shape![16, 1, 14336], crate::strides![14336, 14336, 1]);
309
310        let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out);
311
312        assert_eq!(
313            action(&analysis),
314            MatmulTransformAction::MergeBatches { rows: 16 }
315        );
316    }
317
318    #[test]
319    fn metadata_ignores_unit_dim_strides() {
320        // Size-1 dims carry arbitrary strides; they must not anchor the chain.
321        let lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 12345, 1]);
322        let rhs = Metadata::new(shape![1, 4096, 14336], crate::strides![0, 14336, 1]);
323        let out = Metadata::new(shape![16, 1, 14336], crate::strides![14336, 99999, 1]);
324
325        let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out);
326
327        assert_eq!(
328            action(&analysis),
329            MatmulTransformAction::MergeBatches { rows: 16 }
330        );
331    }
332
333    #[test]
334    fn apply_folds_batches_in_place() {
335        let mut lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 4096, 1]);
336
337        MatmulTransformAction::MergeBatches { rows: 16 }.apply(&mut lhs);
338
339        assert_eq!(lhs.shape(), &shape![1, 16, 4096]);
340        assert_eq!(lhs.strides()[1], 4096);
341        assert_eq!(lhs.strides()[2], 1);
342    }
343
344    #[test]
345    fn apply_keep_is_noop() {
346        let mut lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 4096, 1]);
347
348        MatmulTransformAction::Keep.apply(&mut lhs);
349
350        assert_eq!(lhs.shape(), &shape![16, 1, 4096]);
351    }
352}
burn_std/tensor/matmul.rs

burn_std/tensor/
matmul.rs
