cubecl_std/
swizzle.rs

1use cubecl::prelude::*;
2use cubecl_core::{self as cubecl};
3
4/// Swizzling strategy for a buffer.
5/// See the following docs from cutlass:
6///
7/// 0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx
8///                                ^--^ `MBase` is the number of least-sig bits to keep constant
9///                   ^-^       ^-^     `BBits` is the number of bits in the mask
10///                     ^---------^     `SShift` is the distance to shift the YYY mask
11///                                        (pos shifts YYY to the right, neg shifts YYY to the left)
12///
13/// # Example
14/// Given:
15/// 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx
16/// the result is:
17/// 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY
18///
19///
20/// Some newer features, as well as cutlass in places, use a different terminology of `span` and
21/// `atom`. For shared memory swizzle specifically, the parameters map as follows:
22/// * `bits` = `log2(span / atom)`, or the number of atoms within one span, converted to address bits
23/// * `base` = `log2(atom)`, the size of the atom, converted to address bits
24/// * `shift` = `log2(all_banks_bytes / atom)`, or the total number of atoms in all 32 shared memory banks, converted to address bits
25///
26/// For example:
27/// * 32-byte span with a 16-byte atom = `[1, 4, 3]`
28/// * 128-byte span with a 32-byte atom = `[3, 5, 2]`
29///
30#[derive(CubeType, CubeLaunch, Clone, Copy)]
31#[expand(derive(Clone, Copy))]
32pub struct Swizzle {
33    #[cube(comptime)]
34    yyy_mask: u32,
35    #[cube(comptime)]
36    shift: u32,
37    #[cube(comptime)]
38    invert_shift: bool,
39    /// Precalculate repeat after so we don't need to keep all the parts
40    #[cube(comptime)]
41    repeats_after: u32,
42}
43
44#[cube]
45impl Swizzle {
46    /// Create a new swizzle with comptime parameters
47    pub fn new(#[comptime] bits: u32, #[comptime] base: u32, #[comptime] shift: i32) -> Self {
48        let invert_shift = shift < 0;
49        let mask = (1u32 << bits) - 1;
50        let yyy_mask = comptime![mask << (base + Ord::max(shift, 0) as u32)];
51        let repeats_after = comptime![if bits > 0 {
52            1u32 << (base + bits + Ord::max(shift, 0) as u32)
53        } else {
54            1u32 << base
55        }];
56        Swizzle {
57            yyy_mask,
58            shift: comptime![shift.unsigned_abs()],
59            invert_shift,
60            repeats_after,
61        }
62    }
63
64    /// Create a new noop swizzle object
65    pub fn none() -> Self {
66        Swizzle {
67            yyy_mask: 0u32,
68            shift: 0u32,
69            invert_shift: false,
70            repeats_after: 1u32,
71        }
72    }
73
74    /// Apply the swizzle to a coordinate with a given item size. This is the size of the full type,
75    /// including vectorization.
76    /// `offset` should be in terms of vectors from the start of the buffer, and the buffer should be
77    /// aligned to `repeats_after`. This is to work around the fact we don't currently support
78    /// retrieving the actual address of an offset.
79    /// If you're using absolute/unvectorized indices, pass `E::Scalar::type_size()` instead of the full
80    /// vector size.
81    pub fn apply(&self, offset: u32, #[comptime] type_size: usize) -> u32 {
82        // Special case here so we don't need to special case in kernels that can have no swizzle.
83        // If `yyy_mask == 0`, the whole thing is a noop.
84        if comptime![self.yyy_mask == 0] {
85            offset
86        } else {
87            let offset_bytes = offset * type_size as u32;
88            let offset_masked = offset_bytes & self.yyy_mask;
89            let offset_shifted =
90                shift_right(offset_masked, self.shift, comptime![self.invert_shift]);
91            let offset_bytes = offset_bytes ^ offset_shifted;
92            offset_bytes / type_size as u32
93        }
94    }
95
96    /// After how many elements this pattern repeats. Can be used to align the buffer (i.e. smem)
97    /// so offsets match addresses.
98    pub fn repeats_after(&self) -> comptime_type!(u32) {
99        self.repeats_after
100    }
101}
102
103#[cube]
104fn shift_right(value: u32, shift: u32, #[comptime] invert: bool) -> u32 {
105    if invert {
106        value << shift
107    } else {
108        value >> shift
109    }
110}
cubecl_std/swizzle.rs

cubecl_std/
swizzle.rs
