cubecl_utils_rs/
layout.rs

1//! Row padding to a vectorisation boundary.
2
3/// Lane count for vectorised loads.
4///
5/// Kernels index a row as `Vector<F, LINE_SIZE>`, so every row has to be a
6/// whole number of lines. Four is what `vec4` gives on every backend worth
7/// supporting.
8pub const LINE_SIZE: usize = 4;
9
10/// Round a dimensionality up to a whole number of lines.
11///
12/// ### Params
13///
14/// * `dim` - Original dimensionality
15///
16/// ### Returns
17///
18/// The smallest multiple of [`LINE_SIZE`] that is at least `dim`.
19#[inline]
20pub fn padded_dim(dim: usize) -> usize {
21    dim.next_multiple_of(LINE_SIZE)
22}
23
24/// Pad rows to `dim_padded` by appending zeros to each.
25///
26/// ### Params
27///
28/// * `flat` - Flattened row-major data of size `n * dim`
29/// * `n` - Number of rows
30/// * `dim` - Original dimensionality
31/// * `dim_padded` - Target dimensionality, must be at least `dim`
32///
33/// ### Returns
34///
35/// Padded flat data of size `n * dim_padded`.
36///
37/// ### Note
38///
39/// Padding with zeros is only sound for the metrics where a zero component
40/// contributes nothing: dot products, squared Euclidean and the L2 norms
41/// underneath cosine. It is not sound for anything that counts components.
42pub fn pad_vectors<T: num_traits::Float>(
43    flat: &[T],
44    n: usize,
45    dim: usize,
46    dim_padded: usize,
47) -> Vec<T> {
48    let mut padded = vec![T::zero(); n * dim_padded];
49    for i in 0..n {
50        let src = &flat[i * dim..(i + 1) * dim];
51        let dst = &mut padded[i * dim_padded..i * dim_padded + dim];
52        dst.copy_from_slice(src);
53    }
54    padded
55}
56
57///////////
58// Tests //
59///////////
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn test_padded_dim() {
67        assert_eq!(padded_dim(0), 0);
68        assert_eq!(padded_dim(1), 4);
69        assert_eq!(padded_dim(4), 4);
70        assert_eq!(padded_dim(5), 8);
71        assert_eq!(padded_dim(128), 128);
72    }
73
74    #[test]
75    fn test_pad_vectors_appends_zeros() {
76        let flat = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
77        let padded = pad_vectors(&flat, 2, 3, 4);
78        assert_eq!(padded, vec![1.0, 2.0, 3.0, 0.0, 4.0, 5.0, 6.0, 0.0]);
79    }
80
81    #[test]
82    fn test_pad_vectors_noop_when_already_aligned() {
83        let flat = vec![1.0f64, 2.0, 3.0, 4.0];
84        assert_eq!(pad_vectors(&flat, 1, 4, 4), flat);
85    }
86}
cubecl_utils_rs/layout.rs

cubecl_utils_rs/
layout.rs
