cubecl_tile/space/
extent.rs

1//! One axis's size ([`Extent`]) and every axis's at once ([`Shape`]).
2
3use crate::{Axis, AxisMap};
4
5/// One axis's size: `Static` is comptime, `Dynamic` is resolved in-kernel from the tensor shape.
6#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
7pub(crate) enum Extent {
8    Static(usize),
9    Dynamic,
10}
11
12impl Extent {
13    /// The comptime size; panics on `Dynamic`.
14    pub(crate) fn get(self) -> usize {
15        match self {
16            Extent::Static(n) => n,
17            Extent::Dynamic => {
18                panic!("Extent::get: this axis is Dynamic; its size is only known at runtime")
19            }
20        }
21    }
22
23    pub(crate) fn is_dynamic(self) -> bool {
24        matches!(self, Extent::Dynamic)
25    }
26
27    /// The extent two spaces sharing an axis broadcast to when they are
28    /// [`merge`](crate::Space::merge)d.
29    pub(crate) fn broadcast(self, other: Extent) -> Extent {
30        match (self, other) {
31            (Extent::Static(1), b) => b,
32            (a, Extent::Static(1)) => a,
33            (Extent::Dynamic, _) | (_, Extent::Dynamic) => Extent::Dynamic,
34            (Extent::Static(a), Extent::Static(b)) if a == b => Extent::Static(a),
35            _ => panic!("Space::merge: axis appears with conflicting extents"),
36        }
37    }
38}
39
40/// Every axis with its comptime extent, in canonical order.
41#[derive(Clone, PartialEq, Eq, Hash, Debug)]
42pub struct Shape {
43    extents: AxisMap<Extent>,
44}
45
46impl Shape {
47    pub(crate) fn new(extents: &[(Axis, Extent)]) -> Self {
48        Shape {
49            extents: AxisMap::new(extents),
50        }
51    }
52
53    pub fn rank(&self) -> usize {
54        self.extents.len()
55    }
56
57    pub fn axis_at(&self, i: usize) -> Axis {
58        self.extents.axis_at(i)
59    }
60
61    pub fn position(&self, axis: Axis) -> usize {
62        self.extents.position(axis)
63    }
64
65    pub fn contains(&self, axis: Axis) -> bool {
66        self.extents.contains(axis)
67    }
68
69    pub fn axes(&self) -> impl Iterator<Item = Axis> + '_ {
70        self.extents.axes()
71    }
72
73    pub(crate) fn extent_raw(&self, axis: Axis) -> Extent {
74        self.extents.get(axis)
75    }
76
77    /// The axis's comptime size; panics on a dynamic axis.
78    pub fn extent(&self, axis: Axis) -> usize {
79        self.extent_raw(axis).get()
80    }
81
82    pub fn extent_at(&self, i: usize) -> usize {
83        self.extent(self.axis_at(i))
84    }
85
86    pub fn is_dynamic(&self, axis: Axis) -> bool {
87        self.extent_raw(axis).is_dynamic()
88    }
89
90    /// Every axis is [`Static`](Extent::Static), so a walk over it is fully comptime.
91    pub(crate) fn is_static(&self) -> bool {
92        self.axes().all(|axis| !self.is_dynamic(axis))
93    }
94
95    /// The listed axes with their extents, in the order listed.
96    pub(crate) fn subspace(&self, axes: &[Axis]) -> Shape {
97        let entries: Vec<_> = axes.iter().map(|&a| (a, self.extent_raw(a))).collect();
98        Shape::new(&entries)
99    }
100}
cubecl_tile/space/extent.rs

cubecl_tile/space/
extent.rs
