cubecl_tile/layout/storage/elem_alignment.rs
1//! A buffer layout's alignment in elements, along its innermost dim.
2
3use super::geometry::stride_order;
4use crate::{Axis, Geometry, gcd, layout::StoragePartitioning};
5use cubecl::prelude::TensorBinding;
6
7/// The alignment of a buffer's layout, counted in elements along its innermost dim: the number
8/// every line width that reads the buffer's elements and no others divides.
9///
10/// Not a byte or address alignment: the tile launcher refuses a width that does not divide the
11/// innermost extent or any coarser stride, and every width past one where the innermost stride is
12/// not one. Those are all "divides", so the layout is aligned to their greatest common divisor, and
13/// that is the whole of what this holds, read once off the geometry so a derivation can ask it
14/// without the tensor. It is one of three things a bound width answers to: the device's ceiling,
15/// this, and a kernel's own requirements ([`Vectorization`](crate::launch::Vectorization)).
16///
17/// A view's start is the one address it reads: a view starting partway into its buffer is aligned
18/// to its offset in elements ([`at_start`](Self::at_start)), the buffer's own start being the
19/// runtime's to align.
20#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
21pub struct ElemAlignment(Widths);
22
23/// The widths that read a buffer's elements and no others.
24#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
25enum Widths {
26 /// Every width dividing this one.
27 Dividing(usize),
28 /// This width alone: a storage-tiled buffer stores its read as a piece, which a reader takes
29 /// whole.
30 Exactly(usize),
31}
32
33impl ElemAlignment {
34 /// A layout aligned to one element, read one element at a time: a strided or broadcast view,
35 /// whose innermost dim strides past one.
36 pub const ONE: Self = Self(Widths::Dividing(1));
37
38 /// The alignment of a buffer of `shape` at `strides`, its dims taken in the order the buffer
39 /// steps them, whatever order the view names them in.
40 pub fn in_stride_order(shape: &[usize], strides: &[usize]) -> Self {
41 let mut dims: Vec<(usize, usize)> =
42 shape.iter().copied().zip(strides.iter().copied()).collect();
43 dims.sort_by_key(|&dim| stride_order(dim));
44 let (shape, strides): (Vec<usize>, Vec<usize>) = dims.into_iter().unzip();
45 Self::new(&shape, &strides)
46 }
47
48 /// The alignment of rows of `columns` elements laid one after another: a buffer a launch
49 /// allocates or copies dense. Its rows count for nothing, as only the step between them does.
50 pub fn dense_rows(columns: usize) -> Self {
51 Self::new(&[columns], &[1])
52 }
53
54 /// The alignment of a buffer of `shape` at `strides`, innermost last.
55 ///
56 /// A dim of one element steps nowhere, so its stride reads no address and aligns nothing:
57 /// a view `[1, k]` swapped out of a column `[k, 1]` keeps the column's stride of one on its
58 /// single row, and is as aligned as the dense row it is.
59 pub fn new(shape: &[usize], strides: &[usize]) -> Self {
60 let (shape, strides): (Vec<usize>, Vec<usize>) = shape
61 .iter()
62 .zip(strides)
63 .filter(|&(&extent, _)| extent != 1)
64 .unzip();
65 let divides = match (shape.split_last(), strides.split_last()) {
66 (Some((&extent, _)), Some((&1, coarser))) => {
67 coarser.iter().fold(extent, |acc, &stride| gcd(acc, stride))
68 }
69 _ => 1,
70 };
71 Self(Widths::Dividing(divides))
72 }
73
74 /// The alignment a view starting `start_bytes` into its buffer gives lines of `element_bytes`
75 /// elements: the start's offset in elements, which every line's first element sits a whole
76 /// number of lines past, so a width dividing it keeps each line on its own boundary. A start
77 /// at no whole element is aligned to one; a view at its buffer's start, to any width.
78 pub fn at_start(start_bytes: u64, element_bytes: usize) -> Self {
79 let element_bytes = element_bytes.max(1) as u64;
80 let divides = match start_bytes.is_multiple_of(element_bytes) {
81 true => (start_bytes / element_bytes) as usize,
82 false => 1,
83 };
84 Self(Widths::Dividing(divides))
85 }
86
87 /// The alignment of a buffer stored in pieces `read` elements wide, which a reader takes whole:
88 /// a storage-tiled one's.
89 pub const fn stored_piece(read: usize) -> Self {
90 Self(Widths::Exactly(read))
91 }
92
93 /// Whether the buffer is read in the pieces it is stored in, a storage-tiled one's, rather than
94 /// along rows.
95 pub fn is_stored_piece(self) -> bool {
96 matches!(self.0, Widths::Exactly(_))
97 }
98
99 /// Whether a line `width` wide reads this buffer's elements and no others.
100 pub fn admits(self, width: usize) -> bool {
101 match self.0 {
102 _ if width == 1 => true,
103 Widths::Exactly(read) => width == read,
104 Widths::Dividing(divides) => divides.is_multiple_of(width),
105 }
106 }
107}
108
109impl From<&TensorBinding> for ElemAlignment {
110 /// The alignment of a bound matrix operand: its geometry's, or, for a storage-tiled one, the
111 /// piece its storage partitioning states as the finest along its columns, and a piece one
112 /// element wide where none does: no partitioning describes it, or its tiles step down a
113 /// column first, which the matmul kernels refuse before reading this.
114 fn from(binding: &TensorBinding) -> Self {
115 if !binding.tiling.is_tiled() {
116 return Self::new(&binding.shape, &binding.strides);
117 }
118 const ROWS: Axis = Axis(0);
119 const COLS: Axis = Axis(1);
120 let geometry = Geometry::from(binding);
121 let storage = StoragePartitioning::new(&geometry, &geometry.labels(&[ROWS, COLS]));
122 match storage.as_ref().and_then(|storage| storage.tiles().first()) {
123 Some(&(COLS, read)) => Self::stored_piece(read),
124 _ => Self::stored_piece(1),
125 }
126 }
127}
128
129#[cfg(test)]
130mod tests {
131 /// A row viewed out of a column keeps the column's stride on its single row, which steps
132 /// nowhere: it is as aligned as the dense row it is.
133 #[test]
134 fn a_dim_of_one_element_aligns_nothing() {
135 let k = 8;
136 let row = ElemAlignment::new(&[1, k], &[1, 1]);
137 assert_eq!(row, ElemAlignment::dense_rows(k));
138 assert!(ElemAlignment::in_stride_order(&[1, k], &[1, 1]).admits(k));
139 }
140
141 /// A column bound in buffer order runs its lines along the dim it steps, not along its dim of
142 /// one element, which ties its stride: `[k, 1]` at strides `[1, 1]` is a dense run of `k`.
143 #[test]
144 fn a_dim_of_one_element_is_never_innermost_in_stride_order() {
145 let k = 4;
146 assert_eq!(
147 ElemAlignment::in_stride_order(&[k, 1], &[1, 1]),
148 ElemAlignment::dense_rows(k)
149 );
150 let (_, axes) = Geometry::new(&[(k, 1), (1, 1)]).in_stride_order(&[Axis(2), Axis(1)]);
151 assert_eq!(axes, [Axis(1), Axis(2)]);
152 }
153
154 /// A col weight's rows run along `k`, a pitch apart, whichever way the view names them.
155 #[test]
156 fn a_view_is_aligned_in_the_order_its_buffer_steps() {
157 let (k, n, pitch) = (64, 48, 72);
158 let col = ElemAlignment::in_stride_order(&[4, k, n], &[n * pitch, 1, pitch]);
159 assert!(col.admits(8) && !col.admits(16));
160 assert!(ElemAlignment::in_stride_order(&[4, k, n], &[k * n, n, 1]).admits(16));
161 }
162
163 use super::*;
164
165 /// A dense row is aligned to its length; padding past it to what the pad
166 /// divides too; a strided innermost dim to one element.
167 /// A view's start aligns it to its offset in elements: any width at its buffer's start, the
168 /// offset's divisors partway in, one element where the start splits an element.
169 #[test]
170 fn a_start_is_aligned_to_its_offset_in_elements() {
171 assert!(ElemAlignment::at_start(0, 2).admits(16));
172 let partway = ElemAlignment::at_start(24, 2);
173 assert!(partway.admits(4) && !partway.admits(8));
174 assert!(!ElemAlignment::at_start(3, 2).admits(2));
175 }
176
177 #[test]
178 fn a_layout_is_aligned_to_the_divisors_of_its_extent_and_strides() {
179 let dense = ElemAlignment::new(&[4, 1500], &[1500, 1]);
180 assert!(dense.admits(4) && !dense.admits(8));
181 let pitched = ElemAlignment::new(&[4, 1504], &[1508, 1]);
182 assert!(pitched.admits(4) && !pitched.admits(8));
183 let col = ElemAlignment::new(&[4096, 4096], &[1, 4096]);
184 assert!(col.admits(1) && !col.admits(2));
185 let broadcast = ElemAlignment::new(&[8, 64], &[0, 1]);
186 assert!(broadcast.admits(64));
187 }
188}