cubek_std/tile/variants/
shared.rs

1//! `TileKind::SharedTile` — the smem stage slot variant.
2//!
3//! [`SharedTile`] is the enum payload (vectorization erased from the type).
4//! [`StridedTile`] is the typed form readers/writers consume.
5//! [`SharedTile::wrap`] / [`SharedTile::view`] are pure retypes between them.
6
7use cubecl::{intrinsic, prelude::*, std::Swizzle};
8
9use crate::MatrixLayout;
10use crate::stage::{StageMemoryConfig, as_swizzle_object};
11use crate::tile::variants::instruction::{
12    cmma::cmma_write_to_shared,
13    interleaved::interleaved_write_to_shared,
14    mma::{MmaFragment, MmaFragmentExpand, mma_write_to_shared},
15    plane_vec::planevec_write_to_shared,
16    register::register_write_to_shared,
17};
18use crate::tile::{Tile, TileKind, TileKindExpand, TileScope};
19
20#[derive(CubeType, Clone)]
21#[expand(derive(Clone))]
22/// Typed form of the smem stage slot. `start`/`end`/`stride` are in vector
23/// units.
24pub struct StridedTile<ES: Numeric, N: Size> {
25    /// Slice containing all data for the stage
26    pub container: Box<[Vector<ES, N>]>,
27    /// Offset of the tile in the stage
28    pub start: u32,
29    /// End of the tile in the stage, may be wrong with swizzle
30    pub end: u32,
31    /// Stride between each row/col, depending on MatrixLayout (the other is assumed to be 1)
32    pub stride: u32,
33    /// Swizzle object to transform the index
34    pub swizzle: Swizzle,
35    #[cube(comptime)]
36    /// Layout of the tile (row-major or column-major).
37    pub layout: MatrixLayout,
38}
39
40#[cube]
41impl<ES: Numeric, N: Size> StridedTile<ES, N> {
42    /// Creates a tile from a contiguous slice of data.
43    ///
44    /// The slice length must exactly match the tile size.
45    pub fn new_contiguous(
46        container: &[Vector<ES, N>],
47        start: u32,
48        #[comptime] config: StageMemoryConfig,
49    ) -> StridedTile<ES, N> {
50        let len = config.elements_per_tile() / config.vector_size;
51        let layout = config.matrix_layout;
52        let stride = match layout {
53            MatrixLayout::RowMajor => config.elements_per_tile_along_col,
54            MatrixLayout::ColMajor => config.elements_per_tile_along_row,
55        };
56
57        let stride = stride / config.vector_size;
58
59        StridedTile::<ES, N> {
60            container: unsafe { container.as_boxed_unchecked() },
61            start,
62            end: start + len,
63            stride,
64            swizzle: as_swizzle_object(config.swizzle),
65            layout,
66        }
67    }
68
69    /// Creates a tile from a strided slice of data.
70    ///
71    /// The slice must include all elements of the tile, though it may include unused gaps.
72    pub fn new_strided(
73        container: &[Vector<ES, N>],
74        start: u32,
75        end: u32,
76        stride: u32,
77        swizzle: Swizzle,
78        #[comptime] layout: MatrixLayout,
79    ) -> StridedTile<ES, N> {
80        StridedTile::<ES, N> {
81            container: unsafe { container.as_boxed_unchecked() },
82            start,
83            end,
84            stride,
85            swizzle,
86            layout,
87        }
88    }
89}
90
91#[cube]
92impl<ES: Numeric, N: Size> StridedTile<ES, N> {
93    pub fn unvectorized_stride(&self) -> u32 {
94        let stage_vector_size = self.container.vector_size();
95        self.stride * stage_vector_size as u32
96    }
97}
98
99#[cube]
100impl<ES: Numeric, N: Size> StridedTile<ES, N> {
101    /// Returns the tile as an offset read-only slice. Should only be used when swizzling is
102    /// definitely not applicable.
103    pub fn as_slice(&self) -> &[Vector<ES, N>] {
104        &self.container[self.start as usize..self.end as usize]
105    }
106
107    /// Returns the tile as an offset slice. Should only be used when swizzling is definitely not
108    /// applicable.
109    pub fn as_slice_mut(&mut self) -> &mut [Vector<ES, N>] {
110        &mut self.container[self.start as usize..self.end as usize]
111    }
112}
113
114#[cube]
115impl<ES: Numeric, N: Size> StridedTile<ES, N> {
116    /// Returns a specific vector from the tile based on coordinates.
117    pub fn get_vector(&self, coor_strided: u32, coor_contiguous: u32) -> Vector<ES, N> {
118        let offset = coor_strided * self.stride + coor_contiguous;
119        let offset_abs = self.start + offset;
120        let type_size = Vector::<ES, N>::type_size();
121        let offset_swizzled = self.swizzle.apply(offset_abs, type_size);
122        self.container[offset_swizzled as usize]
123    }
124
125    pub fn stage_offset(&self, relative_offset: u32) -> u32 {
126        let offset = self.start + relative_offset;
127        let type_size = Vector::<ES, N>::type_size();
128        self.swizzle.apply(offset, type_size)
129    }
130
131    #[allow(unused_variables)]
132    pub fn with_vector_size<N2: Size>(&self) -> StridedTile<ES, N2> {
133        let vector_size = N2::value();
134        intrinsic!(|scope| {
135            let stage_vector_size = self.container.vector_size();
136
137            if vector_size == self.container.vector_size() {
138                return self.__expand_with_stage_vector_size_method(scope);
139            }
140
141            let current = stage_vector_size;
142            let mut out: StridedTileExpand<ES, N2> =
143                self.clone().__expand_with_stage_vector_size_method(scope);
144
145            if current < vector_size {
146                let ratio = ((vector_size / current) as u32).into_expand(scope);
147                let start = self.start.__expand_div_method(scope, ratio);
148                let end = self.end.__expand_div_method(scope, ratio);
149                let stride = self.stride.__expand_div_method(scope, ratio);
150                out.start = start;
151                out.end = end;
152                out.stride = stride;
153            } else {
154                let ratio = ((current / vector_size) as u32).into_expand(scope);
155                let start = self.start.__expand_mul_method(scope, ratio);
156                let end = self.end.__expand_mul_method(scope, ratio);
157                let stride = self.stride.__expand_mul_method(scope, ratio);
158                out.start = start;
159                out.end = end;
160                out.stride = stride;
161            }
162
163            out
164        })
165    }
166
167    /// Cast only the stage vector size. This leaves the tile in an invalid state - start, end and
168    /// stride must be adjusted accordingly.
169    /// # Safety
170    /// Must not be used without further metadata adjustments
171    #[allow(unused)]
172    unsafe fn with_stage_vector_size<N2: Size>(&self) -> StridedTile<ES, N2> {
173        StridedTile::<ES, N2> {
174            container: unsafe { self.container.with_vector_size::<N2>().as_boxed_unchecked() },
175            start: self.start,
176            end: self.end,
177            stride: self.stride,
178            swizzle: self.swizzle,
179            layout: self.layout,
180        }
181    }
182}
183
184/// Payload of [`TileKind::SharedTile`]. Vectorization is erased from the
185/// type but kept on the runtime slice; project back with [`view`](Self::view).
186#[derive(CubeType, Clone)]
187pub struct SharedTile<E: Numeric> {
188    pub(crate) container: Box<[E]>,
189    pub(crate) start: u32,
190    pub(crate) end: u32,
191    pub(crate) stride: u32,
192    pub(crate) swizzle: Swizzle,
193    #[cube(comptime)]
194    pub(crate) layout: MatrixLayout,
195}
196
197#[cube]
198impl<E: Numeric> SharedTile<E> {
199    /// Erase the vectorization from a [`StridedTile`].
200    pub fn wrap<V: Size>(tile: StridedTile<E, V>) -> SharedTile<E> {
201        let container = unsafe { tile.container.downcast_unchecked::<E>() };
202        SharedTile::<E> {
203            container: unsafe { container.as_boxed_unchecked() },
204            start: tile.start,
205            end: tile.end,
206            stride: tile.stride,
207            swizzle: tile.swizzle,
208            layout: tile.layout,
209        }
210    }
211
212    /// Project back to a typed [`StridedTile`]. Must match the original
213    /// vectorization.
214    pub fn view<V: Size>(&self) -> StridedTile<E, V> {
215        let container = unsafe { self.container.downcast_unchecked::<Vector<E, V>>() };
216        StridedTile::<E, V> {
217            container: unsafe { container.as_boxed_unchecked() },
218            start: self.start,
219            end: self.end,
220            stride: self.stride,
221            swizzle: self.swizzle,
222            layout: self.layout,
223        }
224    }
225}
226
227#[cube]
228impl<E: Numeric> SharedTile<E> {
229    /// Write-back leg of `Tile::copy_from`: routes to the source variant's
230    /// `*_write_to_shared` helper.
231    pub fn copy_from<SE: Numeric, SS: Size, L: Numeric, R: Numeric, Sc: TileScope>(
232        &mut self,
233        source: &Tile<SE, Sc>,
234    ) {
235        match &source.kind {
236            TileKind::Cmma(t) => cmma_write_to_shared::<E, SS, SE>(self, &t.matrix),
237            TileKind::Bounce(b) => cmma_write_to_shared::<E, SS, SE>(self, &b.cmma.matrix),
238            TileKind::Mma(t) => match &t.fragment {
239                MmaFragment::Acc(f) => {
240                    mma_write_to_shared::<E, SS, SE, L, R>(self, f, t.tile_size, t.mma_io_config);
241                }
242                MmaFragment::Lhs(_) | MmaFragment::Rhs(_) => {
243                    panic!("Mma write_to_shared only supported for Acc role")
244                }
245            },
246            TileKind::Register(t) => {
247                register_write_to_shared::<E, SS, SE>(self, &t.tile.data, t.tile_size);
248            }
249            TileKind::PlaneVec(t) => {
250                planevec_write_to_shared::<SE, E, SS>(
251                    self,
252                    &t.data,
253                    t.tile_size,
254                    t.reduce_vector_size,
255                );
256            }
257            TileKind::Interleaved(t) => {
258                interleaved_write_to_shared::<E, SS, SE>(self, &t.data, t.tile_size);
259            }
260            _ => panic!("SharedTile::copy_from: unsupported source variant"),
261        }
262    }
263}
cubek_std/tile/variants/shared.rs

cubek_std/tile/variants/
shared.rs
