cubek_convolution/components/global/layout/
spatial.rs

1use cubecl::std::tensor::{
2    layout::{
3        Coordinates, Coords1d, Layout, LayoutExpand,
4        as_dyn::{IntoDyn, IntoDynExpand},
5    },
6    r#virtual::VirtualTensor,
7};
8use cubecl::{
9    prelude::*,
10    std::tensor::launch::{MemoryArg, ViewLayoutLaunchArg},
11};
12use enumset::{EnumSet, EnumSetType};
13
14use crate::components::Dimensionality;
15
16#[derive(CubeType, CubeLaunch, Clone)]
17#[expand(derive(Clone))]
18pub struct NhwcCoords {
19    pub batch: u32,
20    pub spatial: Sequence<i32>,
21    pub channel: u32,
22}
23
24impl DerefExpand for NhwcCoordsExpand {
25    type Target = Self;
26
27    fn __expand_deref_method(&self, _: &Scope) -> Self::Target {
28        self.clone()
29    }
30}
31
32#[cube]
33impl IntoDyn for NhwcCoords {
34    fn into_dyn(self) -> Sequence<i32> {
35        let mut seq = Sequence::new();
36        seq.push(self.batch as i32);
37        for x in self.spatial {
38            seq.push(x);
39        }
40        seq.push(self.channel as i32);
41        seq
42    }
43}
44
45type NhwcTuple = (u32, Sequence<i32>, u32);
46
47#[cube]
48impl NhwcCoords {
49    pub fn new(batch: u32, spatial: Sequence<i32>, channel: u32) -> Self {
50        NhwcCoords {
51            batch,
52            spatial,
53            channel,
54        }
55    }
56
57    fn into_tuple(self) -> NhwcTuple {
58        (self.batch, self.spatial, self.channel)
59    }
60
61    fn from_tuple(tuple: NhwcTuple) -> Self {
62        NhwcCoords::new(tuple.0, tuple.1, tuple.2)
63    }
64}
65
66#[cube]
67impl Coordinates for NhwcCoords {
68    fn add(this: Self, other: Self) -> Self {
69        let tuple = NhwcTuple::add(this.into_tuple(), other.into_tuple());
70        NhwcCoords::from_tuple(tuple)
71    }
72
73    fn sub(this: Self, other: Self) -> Self {
74        let tuple = NhwcTuple::sub(this.into_tuple(), other.into_tuple());
75        NhwcCoords::from_tuple(tuple)
76    }
77
78    fn min(this: Self, other: Self) -> Self {
79        let tuple = <NhwcTuple as Coordinates>::min(this.into_tuple(), other.into_tuple());
80        NhwcCoords::from_tuple(tuple)
81    }
82
83    fn max(this: Self, other: Self) -> Self {
84        let tuple = <NhwcTuple as Coordinates>::max(this.into_tuple(), other.into_tuple());
85        NhwcCoords::from_tuple(tuple)
86    }
87
88    fn is_in_bounds(pos: Self, bounds: Self) -> bool {
89        NhwcTuple::is_in_bounds(pos.clone().into_tuple(), bounds.clone().into_tuple())
90    }
91
92    fn from_int(this: Self, #[comptime] value: i64) -> Self {
93        let tuple = NhwcTuple::from_int(this.clone().into_tuple(), value);
94        NhwcCoords::from_tuple(tuple)
95    }
96}
97
98#[derive(EnumSetType, Debug, Hash)]
99pub enum NhwcCheck {
100    Batch,
101    Spatial,
102    Channel,
103}
104
105/// Layout for a spatial (i.e. NHWC) tensor. Bounds check only applies to spatial dimensions, not
106/// channel or batch (because these are implicitly checked in the layouts used with spatial tensors).
107#[derive(CubeType, Clone)]
108pub struct NhwcLayout {
109    /// Stride for N
110    pub stride_batch: usize,
111    /// Strides for DHW
112    pub strides_spatial: Sequence<usize>,
113    /// Stride for C
114    pub stride_channel: usize,
115
116    /// Shape of N
117    pub shape_batch: u32,
118    /// Shape of DHW
119    pub shapes_spatial: Sequence<u32>,
120    /// Shape of C
121    pub shape_channel: u32,
122
123    #[cube(comptime)]
124    pub vector_size: VectorSize,
125    #[cube(comptime)]
126    pub checks: EnumSet<NhwcCheck>,
127}
128
129#[cube]
130impl NhwcLayout {
131    pub fn new<E: Numeric, N: Size, IO: Clone>(
132        tensor: VirtualTensor<E, N, IO>,
133        #[comptime] dim: Dimensionality,
134        #[comptime] checks: EnumSet<NhwcCheck>,
135    ) -> Self {
136        let spatial_dims = dim.num_dims();
137        let mut strides_spatial = Sequence::new();
138        let mut shapes_spatial = Sequence::new();
139
140        #[unroll]
141        for i in 0..spatial_dims {
142            strides_spatial.push(tensor.stride(i + 1));
143            shapes_spatial.push(tensor.shape(i + 1) as u32);
144        }
145
146        let stride_batch = tensor.stride(0);
147        let stride_channel = tensor.stride(spatial_dims + 1);
148
149        let shape_batch = tensor.shape(0) as u32;
150        let shape_channel = tensor.shape(spatial_dims + 1) as u32;
151
152        NhwcLayout {
153            stride_batch,
154            strides_spatial,
155            stride_channel,
156            shape_batch,
157            shapes_spatial,
158            shape_channel,
159            vector_size: tensor.vector_size(),
160            checks,
161        }
162    }
163}
164
165#[cube]
166impl Layout for NhwcLayout {
167    type Coordinates = NhwcCoords;
168    type SourceCoordinates = Coords1d;
169
170    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
171        let NhwcCoords {
172            batch,
173            spatial,
174            channel,
175        } = pos;
176
177        let spatial_dims = self.shapes_spatial.len();
178        let mut read_pos =
179            batch as usize * self.stride_batch + channel as usize * self.stride_channel;
180
181        #[unroll]
182        for i in 0..spatial_dims {
183            read_pos += spatial[i] as usize * self.strides_spatial[i];
184        }
185
186        read_pos / self.vector_size
187    }
188
189    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
190        (self.to_source_pos(pos.clone()), self.is_in_bounds(pos))
191    }
192
193    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
194        let mut in_bounds = true.runtime();
195        if self.checks.comptime().contains(NhwcCheck::Batch) {
196            in_bounds &= pos.batch < self.shape_batch;
197        }
198        if self.checks.comptime().contains(NhwcCheck::Spatial) {
199            let spatial_dims = self.shapes_spatial.len();
200
201            #[unroll]
202            for i in 0..spatial_dims {
203                let pos = pos.spatial[i];
204                in_bounds &= pos >= 0 && (pos as u32) < self.shapes_spatial[i];
205            }
206        }
207        if self.checks.comptime().contains(NhwcCheck::Channel) {
208            in_bounds &= pos.channel < self.shape_channel;
209        }
210
211        in_bounds
212    }
213
214    fn shape(&self) -> Self::Coordinates {
215        NhwcCoords {
216            batch: self.shape_batch,
217            spatial: cast_seq(self.shapes_spatial.clone()),
218            channel: self.shape_channel,
219        }
220    }
221}
222
223#[cube]
224pub(crate) fn cast_seq<From: CubePrimitive, To: CubePrimitive>(
225    seq: Sequence<From>,
226) -> Sequence<To> {
227    let num_elems = seq.len();
228    let mut out_seq = Sequence::new();
229    #[unroll]
230    for i in 0..num_elems {
231        let elem = To::cast_from(seq[i]);
232        out_seq.push(elem);
233    }
234    out_seq
235}
236
237pub struct NhwcLayoutLaunch {
238    checks: EnumSet<NhwcCheck>,
239}
240
241impl NhwcLayoutLaunch {
242    pub fn checked(checks: EnumSet<NhwcCheck>) -> Self {
243        Self { checks }
244    }
245
246    pub fn unchecked() -> Self {
247        Self {
248            checks: EnumSet::empty(),
249        }
250    }
251}
252
253#[derive_cube_comptime]
254pub struct NhwcLayoutCompilationArg {
255    pub spatial_rank: usize,
256    pub checks: EnumSet<NhwcCheck>,
257}
258
259impl ViewLayoutLaunchArg for NhwcLayout {
260    type RuntimeArg<R: Runtime> = NhwcLayoutLaunch;
261    type CompilationArg = NhwcLayoutCompilationArg;
262
263    fn register<R: Runtime, B: MemoryArg>(
264        arg: Self::RuntimeArg<R>,
265        buffer: &B,
266        _: Type,
267        launcher: &mut KernelLauncher<R>,
268    ) -> Self::CompilationArg {
269        let shape = buffer.shape();
270        let strides = buffer.strides();
271
272        let rank = shape.len();
273        let dim_c = rank - 1;
274
275        let stride_batch = strides[0];
276        let strides_spatial = strides[1..dim_c].iter().copied().collect();
277        let stride_channel = strides[dim_c];
278
279        let shape_batch = shape[0] as u32;
280        let shapes_spatial = shape[1..dim_c].iter().map(|s| *s as u32).collect();
281        let shape_channel = shape[dim_c] as u32;
282
283        <usize as LaunchArg>::register(stride_batch, launcher);
284        <Sequence<usize> as LaunchArg>::register(strides_spatial, launcher);
285        <usize as LaunchArg>::register(stride_channel, launcher);
286        <u32 as LaunchArg>::register(shape_batch, launcher);
287        <Sequence<u32> as LaunchArg>::register(shapes_spatial, launcher);
288        <u32 as LaunchArg>::register(shape_channel, launcher);
289
290        NhwcLayoutCompilationArg {
291            spatial_rank: buffer.shape().len() - 2,
292            checks: arg.checks,
293        }
294    }
295
296    fn expand(
297        arg: &Self::CompilationArg,
298        ty: Type,
299        builder: &mut KernelBuilder,
300    ) -> <Self as CubeType>::ExpandType {
301        let strides_comp_arg = (0..arg.spatial_rank).map(|_| ()).collect();
302        let shape_comp_arg = (0..arg.spatial_rank).map(|_| ()).collect();
303        NhwcLayoutExpand {
304            stride_batch: <usize as LaunchArg>::expand(&(), builder),
305            strides_spatial: <Sequence<usize> as LaunchArg>::expand(&strides_comp_arg, builder),
306            stride_channel: <usize as LaunchArg>::expand(&(), builder),
307            shape_batch: <u32 as LaunchArg>::expand(&(), builder),
308            shapes_spatial: <Sequence<u32> as LaunchArg>::expand(&shape_comp_arg, builder),
309            shape_channel: <u32 as LaunchArg>::expand(&(), builder),
310            vector_size: ty.vector_size(),
311            checks: arg.checks,
312        }
313    }
314}
cubek_convolution/components/global/layout/spatial.rs

cubek_convolution/components/global/layout/
spatial.rs
