cubecl_std/tensor/layout/
coordinates.rs

1use cubecl::prelude::*;
2use cubecl_core::{self as cubecl};
3use variadics_please::all_tuples;
4
5/// A set of coordinates used in layouts. Contains some utilities for comptime inspection.
6#[cube]
7pub trait Coordinates: CubeType<ExpandType: Clone> + Clone {
8    /// Add two coordinates together and return the result.
9    fn add(this: Self, other: Self) -> Self;
10    /// Subtract two coordinates from each other and return the result.
11    fn sub(this: Self, other: Self) -> Self;
12    /// Apply an elementwise minimum to the coordinates and return the result.
13    fn min(this: Self, other: Self) -> Self;
14    /// Apply an elementwise maximum to the coordinates and return the result.
15    fn max(this: Self, other: Self) -> Self;
16    /// Check whether `pos` is fully contained within `bounds`.
17    fn is_in_bounds(pos: Self, bounds: Self) -> bool;
18    /// Create a new coordinates object where all values are `value`.
19    /// `this` may be used as a reference coordinate for dynamically sized layouts.
20    fn from_int(this: Self, #[comptime] value: i64) -> Self;
21}
22
23// Aliases for convenience and semantic clarity
24pub type Coords1d = usize;
25pub type Coords1i = i32;
26pub type Coords2d = (u32, u32);
27pub type Coords2i = (i32, i32);
28pub type Coords3d = (u32, u32, u32);
29pub type Coords3i = (i32, i32, i32);
30pub type Coords4d = (u32, u32, u32, u32);
31pub type Coords4i = (i32, i32, i32, i32);
32pub type Coords5d = (u32, u32, u32, u32, u32);
33pub type Coords5i = (i32, i32, i32, i32, i32);
34pub type CoordsDyn = Sequence<u32>;
35pub type CoordsDynI = Sequence<i32>;
36
37macro_rules! impl_coordinates_tuple {
38    ($(($T:ident, $t:ident, $o: ident)),*) => {
39        // Need to force off debug symbols because of macro hygiene weirdness.
40        #[cube(no_debug_symbols)]
41        impl<$($T: Coordinates),*> Coordinates for ($($T),*) {
42            fn add(this: Self, other: Self) -> Self {
43                let ($($t),*) = this;
44                let ($($o),*) = other;
45                ($($T::add($t, $o)),*)
46            }
47            fn sub(this: Self, other: Self) -> Self {
48                let ($($t),*) = this;
49                let ($($o),*) = other;
50                ($($T::sub($t, $o)),*)
51            }
52            fn min(this: Self, other: Self) -> Self {
53                let ($($t),*) = this;
54                let ($($o),*) = other;
55                ($($T::min($t, $o)),*)
56            }
57            fn max(this: Self, other: Self) -> Self {
58                let ($($t),*) = this;
59                let ($($o),*) = other;
60                ($($T::max($t, $o)),*)
61            }
62            fn is_in_bounds(this: Self, other: Self) -> bool {
63                let ($($t),*) = this;
64                let ($($o),*) = other;
65                true $(&& $T::is_in_bounds($t, $o))*
66            }
67            fn from_int(this: Self, #[comptime] value: i64) -> Self {
68                let ($($t),*) = this;
69                ($($T::from_int($t, value)),*)
70            }
71        }
72    };
73}
74
75// Can't blanket implement because of trait rules
76macro_rules! impl_coordinates_primitive {
77    ($ty: ty) => {
78        #[cube]
79        impl Coordinates for $ty {
80            fn add(this: Self, other: Self) -> Self {
81                this + other
82            }
83            fn sub(this: Self, other: Self) -> Self {
84                this - other
85            }
86            fn min(this: Self, other: Self) -> Self {
87                this.min(other)
88            }
89            fn max(this: Self, other: Self) -> Self {
90                this.max(other)
91            }
92            fn is_in_bounds(pos: Self, bounds: Self) -> bool {
93                pos < bounds
94            }
95            fn from_int(_this: Self, #[comptime] value: i64) -> Self {
96                <$ty as Numeric>::from_int(value)
97            }
98        }
99    };
100    ($($ty: ty),*) => {
101        $(impl_coordinates_primitive!($ty);)*
102    }
103}
104
105impl_coordinates_primitive!(u8, u16, u32, u64, usize, i8, i16, i32, i64);
106all_tuples!(impl_coordinates_tuple, 2, 12, T, t, o);
107
108#[cube]
109impl<T: Coordinates<ExpandType: DerefExpand<Target = T::ExpandType>> + Copy> Coordinates
110    for Sequence<T>
111{
112    fn add(this: Self, other: Self) -> Self {
113        let rank = this.len();
114        let mut out = Sequence::new();
115
116        #[unroll]
117        for i in 0..rank {
118            out.push(T::add(this[i], other[i]));
119        }
120
121        out
122    }
123
124    fn sub(this: Self, other: Self) -> Self {
125        let rank = this.len();
126        let mut out = Sequence::new();
127
128        #[unroll]
129        for i in 0..rank {
130            out.push(T::sub(this[i], other[i]));
131        }
132
133        out
134    }
135
136    fn min(this: Self, other: Self) -> Self {
137        let rank = this.len();
138        let mut out = Sequence::new();
139
140        #[unroll]
141        for i in 0..rank {
142            out.push(T::min(this[i], other[i]));
143        }
144
145        out
146    }
147
148    fn max(this: Self, other: Self) -> Self {
149        let rank = this.len();
150        let mut out = Sequence::new();
151
152        #[unroll]
153        for i in 0..rank {
154            out.push(T::max(this[i], other[i]));
155        }
156
157        out
158    }
159
160    fn is_in_bounds(pos: Self, bounds: Self) -> bool {
161        let rank = pos.len();
162        let mut out = true;
163
164        #[unroll]
165        for i in 0..rank {
166            out &= T::is_in_bounds(pos[i], bounds[i]);
167        }
168
169        out
170    }
171
172    fn from_int(this: Self, #[comptime] value: i64) -> Self {
173        let rank = this.len();
174        let mut origin = Sequence::new();
175
176        #[unroll]
177        for i in 0..rank {
178            origin.push(T::from_int(this[i], value));
179        }
180
181        origin
182    }
183}
cubecl_std/tensor/layout/coordinates.rs

cubecl_std/tensor/layout/
coordinates.rs
