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

cubek_convolution/components/global/layout/
spatial.rs
