cubecl_tile/tile/kind/procedural/
memory_values.rs

1//! A memory tile's values read one at a time, at coordinates: what a recipe built over another
2//! operand reads of it.
3
4use cubecl::ir::{ElemType, Scope};
5use cubecl::prelude::*;
6
7use super::{Reads, Recipe, RecipeCoords};
8use crate::*;
9
10/// A memory tile's values, read one at a time at the coordinates they are asked at. An axis a
11/// read does not name is read where the tile's window sits on it, so a tile cut to one owner's
12/// batch and head reads that batch and head; a coordinate past the buffer reads zero. A recipe
13/// itself, evaluated at its own coordinates, so it composes with any other ([`Sum`](super::Sum)).
14#[derive(CubeType, Clone)]
15#[expand(derive(Clone))]
16pub struct MemoryValues {
17    memory: Memory,
18    #[cube(comptime)]
19    element: ElemType,
20}
21
22#[cube]
23impl Tile {
24    /// This tile's values, read one at a time ([`MemoryValues`]).
25    ///
26    /// # Panics
27    ///
28    /// At expansion, on a tile that is not memory.
29    pub fn values(&self) -> MemoryValues {
30        match &self.kind {
31            TileKind::Memory(memory) => MemoryValues {
32                memory: memory.clone(),
33                element: comptime!(self.element),
34            },
35            TileKind::TmaGmem(_)
36            | TileKind::PlaneTile(_)
37            | TileKind::PlanePartition(_)
38            | TileKind::Procedural(_)
39            | TileKind::Lines(_) => {
40                panic!("Tile::values: only a memory tile is read one value at a time")
41            }
42        }
43    }
44}
45
46#[cube]
47impl MemoryValues {
48    /// The value at `coordinates`, one along each of `axes` in order, as `T`.
49    pub fn read<T: Numeric>(&self, #[comptime] axes: Vec<Axis>, coordinates: &Coords<u32>) -> T {
50        let define!(Stored) = self.element;
51        let logical = comptime!(self.memory.projection.logical_axes().to_vec());
52        let mut position = Coords::<u32>::new();
53        #[allow(clippy::needless_range_loop)]
54        #[unroll]
55        for p in 0..logical.len() {
56            let axis = comptime!(logical[p]);
57            match comptime!(axes.iter().position(|&named| named == axis)) {
58                Some(named) => position.push(coordinates.at(named)),
59                None => position.push(u32::cast_from(self.memory.window.origin.at(p))),
60            }
61        }
62        let size!(One) = 1usize;
63        let view = self.memory.read_view::<Stored, One>(self.memory.base());
64        T::cast_from(view.read_checked(position.to_dyn()).extract(0usize))
65    }
66}
67
68#[cube]
69impl<T: Numeric> Recipe<T> for MemoryValues {
70    fn evaluate(&self, coordinates: &RecipeCoords) -> T {
71        let axes = comptime!(coordinates.space.axes().collect::<Vec<_>>());
72        let mut at = Coords::<u32>::new();
73        #[unroll]
74        for p in 0..axes.len() {
75            at.push(coordinates.at(p));
76        }
77        self.read::<T>(axes, &at)
78    }
79}
80
81impl Reads for MemoryValuesExpand {
82    fn reads(&self, _scope: &Scope, axis: Axis) -> bool {
83        self.memory.projection.logical_axes().contains(&axis)
84    }
85}
cubecl_tile/tile/kind/procedural/memory_values.rs

cubecl_tile/tile/kind/procedural/
memory_values.rs
