cubecl_tile/tile/kind/memory/view/
flat.rs

1//! The flat 1-D view over a [`Tile`]: [`FlatLayout`] and its masked [`FlatView`].
2
3use cubecl::prelude::*;
4
5use cubecl_std::tensor::layout::{Coords1d, CoordsDyn, Layout, LayoutExpand};
6
7use crate::*;
8
9/// A masked flat row-major scan over a [`Tile`].
10pub(crate) type FlatView<'a, T> = Masked<'a, T, Coords1d>;
11/// The mutable twin of [`FlatView`].
12pub(crate) type FlatViewMut<'a, T> = MaskedMut<'a, T, Coords1d>;
13
14/// Maps a flat row-major index to an N-D coordinate over `shape` ([`unravel`]).
15#[derive(CubeType, Clone)]
16#[expand(derive(Clone))]
17pub(crate) struct FlatLayout {
18    shape: Coords<u32>,
19}
20
21#[cube]
22impl FlatLayout {
23    pub(crate) fn new(shape: Coords<u32>) -> Self {
24        FlatLayout { shape }
25    }
26}
27
28#[cube]
29impl Layout for FlatLayout {
30    type Coordinates = Coords1d;
31    type SourceCoordinates = CoordsDyn;
32
33    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
34        self.shape.unravel(pos as u32).to_dyn()
35    }
36
37    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
38        (self.to_source_pos(pos), self.is_in_bounds(pos))
39    }
40
41    fn shape(&self) -> Self::Coordinates {
42        let rank = self.shape.len();
43        self.shape
44            .product(comptime!((0..rank).collect::<Vec<_>>()))
45            .cast::<usize>()
46    }
47
48    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
49        pos < self.shape()
50    }
51}
52
53#[cube]
54impl Tile {
55    /// A flat row-major view over `Vector<T, W>` lines of the tile's window, masking the overhang.
56    /// A packed store is refused.
57    pub fn flat<T: Numeric, W: Size>(&self) -> FlatView<'_, Vector<T, W>> {
58        let g = self.mem("flat");
59        if comptime!(g.store.packing != Packing::Plain) {
60            panic!("Tile::flat: a packed tile unpacks under Tile::copy_from or Tile::decoded")
61        }
62        g.flat::<T, W>()
63    }
64
65    /// The mutable twin of [`flat`](Tile::flat).
66    pub fn flat_mut<T: Numeric, W: Size>(&mut self) -> FlatViewMut<'_, Vector<T, W>> {
67        self.mem_mut("flat_mut").flat_mut::<T, W>()
68    }
69}
cubecl_tile/tile/kind/memory/view/flat.rs

cubecl_tile/tile/kind/memory/view/
flat.rs
