cubecl_tile/space/
axis.rs

1//! The opaque axis label every layout and space concept is keyed on.
2
3use cubecl::zspace::SmallVec;
4use serde::{Deserialize, Serialize};
5
6use crate::Space;
7
8/// A labeled axis. The `u8` is a client-assigned index, not a position.
9#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)]
10pub struct Axis(pub u8);
11
12/// A comptime map from [`Axis`] to a value, in declared (canonical) order.
13#[derive(Clone, PartialEq, Eq, Hash, Debug)]
14pub(crate) struct AxisMap<T: Copy> {
15    entries: SmallVec<[(Axis, T); Space::MAX_RANK]>,
16}
17
18impl<T: Copy> AxisMap<T> {
19    /// Order is significant.
20    pub(crate) fn new(entries: &[(Axis, T)]) -> Self {
21        AxisMap {
22            entries: SmallVec::from_slice(entries),
23        }
24    }
25
26    pub(crate) fn get(&self, axis: Axis) -> T {
27        self.entries
28            .iter()
29            .find(|(a, _)| *a == axis)
30            .expect("AxisMap::get: axis not present")
31            .1
32    }
33
34    pub(crate) fn axis_at(&self, i: usize) -> Axis {
35        self.entries[i].0
36    }
37
38    /// The axes in declared order.
39    pub(crate) fn axes(&self) -> impl Iterator<Item = Axis> + '_ {
40        self.entries.iter().map(|&(a, _)| a)
41    }
42
43    pub(crate) fn position(&self, axis: Axis) -> usize {
44        self.entries
45            .iter()
46            .position(|(a, _)| *a == axis)
47            .expect("AxisMap::position: axis not present")
48    }
49
50    pub(crate) fn contains(&self, axis: Axis) -> bool {
51        self.entries.iter().any(|(a, _)| *a == axis)
52    }
53
54    pub(crate) fn len(&self) -> usize {
55        self.entries.len()
56    }
57}
cubecl_tile/space/axis.rs

cubecl_tile/space/
axis.rs
