cubecl_tile/space/
coords.rs

1//! A coordinate or extent list whose entries keep their constness.
2
3use cubecl::ir::{ExpandValue, Scope};
4use cubecl::prelude::*;
5use cubecl::unexpanded;
6use cubecl_std::tensor::layout::CoordsDyn;
7
8use crate::algebra::{Integer, IntegerExpand, simplified_add, simplified_mul};
9
10/// An immutable coordinate/extent list whose expand's `IntoMut` is the identity.
11/// Copying into mutable slots (as `Sequence` does) would erase constness.
12pub struct Coords<C: Int> {
13    _c: core::marker::PhantomData<C>,
14}
15
16impl<C: Int> Clone for Coords<C> {
17    fn clone(&self) -> Self {
18        Coords {
19            _c: core::marker::PhantomData,
20        }
21    }
22}
23
24#[allow(clippy::new_without_default, clippy::len_without_is_empty)]
25impl<C: Int> Coords<C> {
26    #[allow(clippy::new_ret_no_self)]
27    pub fn new() -> Self {
28        unexpanded!()
29    }
30    pub fn push(&mut self, _v: C) {
31        unexpanded!()
32    }
33    /// The element at comptime `i`.
34    pub fn at(&self, _i: usize) -> C {
35        unexpanded!()
36    }
37    /// The comptime length.
38    pub fn len(&self) -> usize {
39        unexpanded!()
40    }
41    /// Re-view as boundary [`CoordsDyn`].
42    pub(crate) fn to_dyn(&self) -> CoordsDyn {
43        unexpanded!()
44    }
45    /// Product of the elements at comptime `picks` (empty picks give `1`).
46    pub(crate) fn product(&self, _picks: Vec<usize>) -> C {
47        unexpanded!()
48    }
49    /// Sum of the elements at comptime `picks` (empty picks give `0`).
50    pub(crate) fn sum(&self, _picks: Vec<usize>) -> C {
51        unexpanded!()
52    }
53
54    /// Copy these coordinates into mutable kernel registers.
55    pub(crate) fn stored(&self) -> Coords<C> {
56        unexpanded!()
57    }
58
59    /// Assign every coordinate into this carrier, produced by [`stored`](Coords::stored).
60    pub(crate) fn store_from(&mut self, _src: &Coords<C>) {
61        unexpanded!()
62    }
63
64    pub fn __expand_new(_scope: &Scope) -> CoordsExpand<C> {
65        CoordsExpand { values: Vec::new() }
66    }
67}
68
69#[cube]
70impl Coords<u32> {
71    /// Constant coordinates holding comptime `values`.
72    // `#[unroll]` needs a range loop; an iterator has no expansion.
73    #[allow(clippy::needless_range_loop)]
74    pub(crate) fn constant(#[comptime] values: Vec<usize>) -> Coords<u32> {
75        let mut out = Coords::<u32>::new();
76
77        #[unroll]
78        for p in 0..comptime!(values.len()) {
79            out.push(comptime!(values[p] as u32).runtime());
80        }
81
82        out
83    }
84
85    /// The coordinates of a flat row-major index `i` over these extents.
86    /// The leading entry skips the modulo so the divide stays constant for constant extents.
87    pub(crate) fn unravel(&self, i: u32) -> Coords<u32> {
88        let n = self.len();
89        let mut out = Coords::<u32>::new();
90
91        #[unroll]
92        for p in 0..n {
93            let quot = i.divided_by(self.product(comptime!(((p + 1)..n).collect::<Vec<_>>())));
94            if comptime!(p == 0) {
95                out.push(quot);
96            } else {
97                out.push(quot.remainder(self.at(p)));
98            }
99        }
100
101        out
102    }
103
104    /// Append `other`'s coordinates after these.
105    pub(crate) fn extend(&mut self, other: &Coords<u32>) {
106        #[unroll]
107        for p in 0..other.len() {
108            self.push(other.at(p));
109        }
110    }
111
112    /// Whether every coordinate of `pos` falls inside these extents.
113    pub(crate) fn within(&self, pos: CoordsDyn) -> bool {
114        let mut valid = true;
115
116        #[unroll]
117        for p in 0..self.len() {
118            valid = valid && pos[p] < self.at(p);
119        }
120
121        valid
122    }
123}
124
125pub struct CoordsExpand<C: Int> {
126    values: Vec<NativeExpand<C>>,
127}
128
129impl<C: Int> CubeType for Coords<C> {
130    type ExpandType = CoordsExpand<C>;
131}
132
133impl<C: Int> IntoExpand for CoordsExpand<C> {
134    type Expand = Self;
135    fn into_expand(self, _scope: &Scope) -> Self {
136        self
137    }
138}
139
140/// Identity, to keep constness (see [`Coords`]).
141impl<C: Int> IntoMut for CoordsExpand<C> {
142    fn into_mut(self, _scope: &Scope) -> Self {
143        self
144    }
145}
146
147impl<C: Int> CubeDebug for CoordsExpand<C> {}
148
149impl<C: Int> Clone for CoordsExpand<C> {
150    fn clone(&self) -> Self {
151        CoordsExpand {
152            values: self.values.clone(),
153        }
154    }
155}
156
157impl<C: Int> ExpandTypeClone for CoordsExpand<C> {
158    fn clone_unchecked(&self) -> Self {
159        self.clone()
160    }
161}
162
163impl<C: Int> AsRefExpand for CoordsExpand<C> {
164    fn __expand_ref_method(&self, _scope: &Scope) -> &Self {
165        self
166    }
167}
168
169impl<C: Int> AsMutExpand for CoordsExpand<C> {
170    fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self {
171        self
172    }
173}
174
175impl<C: Int> CoordsExpand<C> {
176    /// The entries known only at runtime, in order.
177    pub(crate) fn runtime_mut(&mut self) -> impl Iterator<Item = &mut NativeExpand<C>> {
178        self.values
179            .iter_mut()
180            .filter(|v| matches!(v.expand, ExpandValue::Value(_)))
181    }
182
183    pub fn __expand_assign_method(&mut self, _scope: &Scope, other: Self) {
184        self.values = other.values;
185    }
186    pub fn __expand_push_method(&mut self, _scope: &Scope, v: NativeExpand<C>) {
187        self.values.push(v);
188    }
189    pub fn __expand_at_method(&self, _scope: &Scope, i: NativeExpand<usize>) -> NativeExpand<C> {
190        let i = i
191            .expand
192            .as_const()
193            .expect("Coords::at: comptime index only")
194            .as_i64() as usize;
195        self.values[i]
196    }
197    pub fn __expand_len_method(&self, _scope: &Scope) -> usize {
198        self.values.len()
199    }
200    pub fn __expand_to_dyn_method(&self, scope: &Scope) -> SequenceExpand<u32> {
201        let mut out = Sequence::<u32>::__expand_new(scope);
202        for v in &self.values {
203            out.__expand_push_method(scope, unsafe { *v.as_type_ref_unchecked::<u32>() });
204        }
205        out
206    }
207    pub fn __expand_product_method(&self, scope: &Scope, picks: Vec<usize>) -> NativeExpand<C> {
208        let mut acc: NativeExpand<C> =
209            ExpandValue::constant(1u64.into(), C::elem_type(scope)).into();
210        for i in picks {
211            acc = simplified_mul(scope, acc, self.values[i]);
212        }
213        acc
214    }
215    pub fn __expand_sum_method(&self, scope: &Scope, picks: Vec<usize>) -> NativeExpand<C> {
216        let mut acc: NativeExpand<C> =
217            ExpandValue::constant(0u64.into(), C::elem_type(scope)).into();
218        for i in picks {
219            acc = simplified_add(scope, acc, self.values[i]);
220        }
221        acc
222    }
223
224    pub fn __expand_stored_method(&self, scope: &Scope) -> CoordsExpand<C> {
225        CoordsExpand {
226            values: self.values.iter().map(|v| (*v).into_mut(scope)).collect(),
227        }
228    }
229
230    pub fn __expand_store_from_method(&mut self, scope: &Scope, src: &CoordsExpand<C>) {
231        assert_eq!(
232            self.values.len(),
233            src.values.len(),
234            "Coords::store_from: source and destination lengths differ"
235        );
236        for (dst, src) in self.values.iter_mut().zip(&src.values) {
237            dst.__expand_assign_method(scope, *src);
238        }
239    }
240}
cubecl_tile/space/coords.rs

cubecl_tile/space/
coords.rs
