cubecl_std/tensor/view/operations/
tensor_map.rs

1use super::*;
2use crate::tensor::layout::*;
3use cubecl::prelude::*;
4use cubecl_core::{self as cubecl, prelude::barrier::Barrier};
5
6// We don't know the linear layout, so only implement TMA loads/stores
7macro_rules! impl_tensor_map {
8    ($dim: literal, $coords: ty, $($val: ident),*) => {
9        paste::paste! {
10            impl<T: CubePrimitive> ViewOperations<T, $coords> for TensorMap<T, Tiled> {}
11            impl<T: CubePrimitive> ViewOperationsExpand<T, $coords> for NativeExpand<TensorMap<T, Tiled>> {
12                fn __expand_read_method(
13                    &self,
14                    _scope: &Scope,
15                    _pos: <$coords as CubeType>::ExpandType,
16                ) -> <T as CubeType>::ExpandType {
17                    unimplemented!("Can't read from tensor map");
18                }
19
20                fn __expand_read_checked_method(
21                    &self,
22                    _scope: &Scope,
23                    _pos: <$coords as CubeType>::ExpandType,
24                ) -> <T as CubeType>::ExpandType {
25                    unimplemented!("Can't read from tensor map");
26                }
27
28                fn __expand_read_masked_method(
29                    &self,
30                    _scope: &Scope,
31                    _pos: <$coords as CubeType>::ExpandType,
32                    _mask_value: <T as CubeType>::ExpandType,
33                ) -> <T as CubeType>::ExpandType {
34                    unimplemented!("Can't read from tensor map");
35                }
36
37                fn __expand_read_unchecked_method(
38                    &self,
39                    _scope: &Scope,
40                    _pos: <$coords as CubeType>::ExpandType,
41                ) -> <T as CubeType>::ExpandType {
42                    unimplemented!("Can't read from tensor map");
43                }
44
45                fn __expand_as_linear_slice_method(
46                    &self,
47                    _scope: &Scope,
48                    _pos: <$coords as CubeType>::ExpandType,
49                    _end: <$coords as CubeType>::ExpandType,
50                ) -> &SliceExpand<T> {
51                    unimplemented!("Can't read from tensor map");
52                }
53
54                fn __expand_shape_method(&self, _scope: &Scope) -> <$coords as CubeType>::ExpandType {
55                    unimplemented!("Can't read from tensor map");
56                }
57
58                fn __expand_is_in_bounds_method(
59                    &self,
60                    _scope: &Scope,
61                    _pos: <$coords as CubeType>::ExpandType,
62                ) -> NativeExpand<bool> {
63                    // Bounds checks are done in hardware, so treat them as always in bounds for the kernels
64                    true.into()
65                }
66
67                #[allow(unused_parens)]
68                fn __expand_tensor_map_load_method(
69                    &self,
70                    scope: &Scope,
71                    barrier: &NativeExpand<Barrier>,
72                    shared_memory: &mut SliceExpand<T>,
73                    pos: <$coords as CubeType>::ExpandType,
74                ) {
75                    let shared = shared_memory.__expand_downcast_mut_method(scope);
76                    let ($($val),*) = pos;
77                    let ($($val),*) = ($(i32::__expand_cast_from(scope, $val)),*);
78                    barrier.[<__expand_tma_load_ $dim d_method>]::<T, T>(scope, self, shared, $($val),*);
79                }
80            }
81
82            impl<T: CubePrimitive> ViewOperationsMut<T, $coords> for TensorMap<T, Tiled> {}
83            impl<T: CubePrimitive> ViewOperationsMutExpand<T, $coords> for NativeExpand<TensorMap<T, Tiled>> {
84                fn __expand_write_method(
85                    &self,
86                    _scope: &Scope,
87                    _pos: <$coords as CubeType>::ExpandType,
88                    _value: <T as CubeType>::ExpandType,
89                ) {
90                    unimplemented!("Can't write to tensor map");
91                }
92
93                fn __expand_write_checked_method(
94                    &self,
95                    _scope: &Scope,
96                    _pos: <$coords as CubeType>::ExpandType,
97                    _value: <T as CubeType>::ExpandType,
98                ) {
99                    unimplemented!("Can't write to tensor map");
100                }
101
102                fn __expand_as_linear_slice_mut_method(
103                    &self,
104                    _scope: &Scope,
105                    _pos: <$coords as CubeType>::ExpandType,
106                    _end: <$coords as CubeType>::ExpandType,
107                ) -> &mut SliceExpand<T> {
108                    unimplemented!("Can't write to tensor map");
109                }
110
111                #[allow(unused_parens)]
112                fn __expand_tensor_map_store_method(
113                    &self,
114                    scope: &Scope,
115                    shared_memory: &SliceExpand<T>,
116                    pos: <$coords as CubeType>::ExpandType,
117                ) {
118                    let shared = shared_memory.__expand_downcast_method(scope);
119                    let ($($val),*) = pos;
120                    let ($($val),*) = ($(i32::__expand_cast_from(scope, $val)),*);
121                    let mut this = self.clone();
122                    [<tma_store_ $dim d>]::expand::<T, T>(scope, &shared, &mut this, $($val),*);
123                }
124            }
125        }
126    };
127}
128
129impl_tensor_map!(1, Coords1d, x);
130impl_tensor_map!(2, Coords2d, x, y);
131impl_tensor_map!(3, Coords3d, x, y, z);
132impl_tensor_map!(4, Coords4d, x, y, z, v);
133impl_tensor_map!(5, Coords5d, x, y, z, v, w);
134
135impl_tensor_map!(1, Coords1i, x);
136impl_tensor_map!(2, Coords2i, x, y);
137impl_tensor_map!(3, Coords3i, x, y, z);
138impl_tensor_map!(4, Coords4i, x, y, z, v);
139impl_tensor_map!(5, Coords5i, x, y, z, v, w);
140
141// We don't know the linear layout, so only implement TMA loads
142macro_rules! impl_tensor_map_im2col {
143    ($dim: literal, $coords: ty, $($pos: ident),*; $($offs: ident),*) => {
144        paste::paste! {
145            impl<T: CubePrimitive> ViewOperations<T, $coords> for TensorMap<T, Im2col> {}
146            impl<T: CubePrimitive> ViewOperationsExpand<T, $coords> for NativeExpand<TensorMap<T, Im2col>> {
147                fn __expand_read_method(
148                    &self,
149                    _scope: &Scope,
150                    _pos: <$coords as CubeType>::ExpandType,
151                ) -> <T as CubeType>::ExpandType {
152                    unimplemented!("Can't read from tensor map");
153                }
154
155                fn __expand_read_checked_method(
156                    &self,
157                    _scope: &Scope,
158                    _pos: <$coords as CubeType>::ExpandType,
159                ) -> <T as CubeType>::ExpandType {
160                    unimplemented!("Can't read from tensor map");
161                }
162
163                fn __expand_read_masked_method(
164                    &self,
165                    _scope: &Scope,
166                    _pos: <$coords as CubeType>::ExpandType,
167                    _mask_value: <T as CubeType>::ExpandType,
168                ) -> <T as CubeType>::ExpandType {
169                    unimplemented!("Can't read from tensor map");
170                }
171
172                fn __expand_read_unchecked_method(
173                    &self,
174                    _scope: &Scope,
175                    _pos: <$coords as CubeType>::ExpandType,
176                ) -> <T as CubeType>::ExpandType {
177                    unimplemented!("Can't read from tensor map");
178                }
179
180                fn __expand_as_linear_slice_method(
181                    &self,
182                    _scope: &Scope,
183                    _pos: <$coords as CubeType>::ExpandType,
184                    _end: <$coords as CubeType>::ExpandType,
185                ) -> &SliceExpand<T> {
186                    unimplemented!("Can't read from tensor map");
187                }
188
189                fn __expand_shape_method(&self, _scope: &Scope) -> <$coords as CubeType>::ExpandType {
190                    unimplemented!("Can't read from tensor map");
191                }
192
193                fn __expand_is_in_bounds_method(
194                    &self,
195                    _scope: &Scope,
196                    _pos: <$coords as CubeType>::ExpandType,
197                ) -> NativeExpand<bool> {
198                    // Bounds checks are done in hardware, so treat them as always in bounds for the kernels
199                    true.into()
200                }
201
202                #[allow(unused_parens)]
203                fn __expand_tensor_map_load_method(
204                    &self,
205                    scope: &Scope,
206                    barrier: &NativeExpand<Barrier>,
207                    shared_memory: &mut SliceExpand<T>,
208                    pos: <$coords as CubeType>::ExpandType,
209                ) {
210                    let shared = shared_memory.__expand_downcast_mut_method(scope);
211                    let ($($pos),*) = pos.0;
212                    let ($($pos),*) = ($(i32::__expand_cast_from(scope, $pos)),*);
213                    let ($($offs),*) = pos.1;
214                    let ($($offs),*) = ($(u16::__expand_cast_from(scope, $offs)),*);
215
216                    barrier.[<__expand_tma_load_im2col_ $dim d_method>]::<T, T>(scope, self, shared, $($pos),*, $($offs),*);
217                }
218            }
219        }
220    };
221}
222
223impl_tensor_map_im2col!(3, (Coords3d, Coords1d), n, w, c; x);
224impl_tensor_map_im2col!(4, (Coords4d, Coords2d), n, h, w, c; y, x);
225impl_tensor_map_im2col!(5, (Coords5d, Coords3d), n, d, h, w, c; z, y, x);
226
227impl_tensor_map_im2col!(3, (Coords3i, Coords1d), n, w, c; x);
228impl_tensor_map_im2col!(4, (Coords4i, Coords2d), n, h, w, c; y, x);
229impl_tensor_map_im2col!(5, (Coords5i, Coords3d), n, d, h, w, c; z, y, x);
230
231fn as_i32<T: CubePrimitive>(scope: &Scope, pos: &SequenceExpand<T>, i: usize) -> NativeExpand<i32> {
232    let x = pos
233        .__expand_index_method(scope, i.into_expand(scope))
234        .__expand_deref_method(scope);
235    i32::__expand_cast_from(scope, x)
236}
237
238fn as_u16<T: CubePrimitive>(
239    scope: &Scope,
240    offs: &SequenceExpand<T>,
241    i: usize,
242) -> NativeExpand<u16> {
243    let x = offs
244        .__expand_index_method(scope, i.into_expand(scope))
245        .__expand_deref_method(scope);
246    u16::__expand_cast_from(scope, x)
247}
248
249impl<T: CubePrimitive, N: CubePrimitive + Coordinates> ViewOperations<T, Sequence<N>>
250    for TensorMap<T, Tiled>
251{
252}
253impl<T: CubePrimitive, N: CubePrimitive + Coordinates> ViewOperationsExpand<T, Sequence<N>>
254    for NativeExpand<TensorMap<T, Tiled>>
255{
256    fn __expand_read_method(
257        &self,
258        _scope: &Scope,
259        _pos: SequenceExpand<N>,
260    ) -> <T as CubeType>::ExpandType {
261        unimplemented!("Can't read from tensor map");
262    }
263
264    fn __expand_read_checked_method(
265        &self,
266        _scope: &Scope,
267        _pos: SequenceExpand<N>,
268    ) -> <T as CubeType>::ExpandType {
269        unimplemented!("Can't read from tensor map");
270    }
271
272    fn __expand_read_masked_method(
273        &self,
274        _scope: &Scope,
275        _pos: SequenceExpand<N>,
276        _mask_value: <T as CubeType>::ExpandType,
277    ) -> <T as CubeType>::ExpandType {
278        unimplemented!("Can't read from tensor map");
279    }
280
281    fn __expand_read_unchecked_method(
282        &self,
283        _scope: &Scope,
284        _pos: SequenceExpand<N>,
285    ) -> <T as CubeType>::ExpandType {
286        unimplemented!("Can't read from tensor map");
287    }
288
289    fn __expand_as_linear_slice_method(
290        &self,
291        _scope: &Scope,
292        _pos: SequenceExpand<N>,
293        _end: SequenceExpand<N>,
294    ) -> &SliceExpand<T> {
295        unimplemented!("Can't read from tensor map");
296    }
297
298    fn __expand_shape_method(&self, _scope: &Scope) -> SequenceExpand<N> {
299        unimplemented!("Can't read from tensor map");
300    }
301
302    fn __expand_is_in_bounds_method(
303        &self,
304        _scope: &Scope,
305        _pos: SequenceExpand<N>,
306    ) -> NativeExpand<bool> {
307        // Bounds checks are done in hardware, so treat them as always in bounds for the kernels
308        true.into()
309    }
310
311    #[allow(unused_parens)]
312    fn __expand_tensor_map_load_method(
313        &self,
314        scope: &Scope,
315        barrier: &NativeExpand<Barrier>,
316        shared_memory: &mut SliceExpand<T>,
317        pos: SequenceExpand<N>,
318    ) {
319        let rank = pos.len();
320        let pos = &pos;
321        match rank {
322            1 => {
323                let x = as_i32(scope, pos, 0);
324                barrier.__expand_tma_load_1d_method(scope, self, shared_memory, x);
325            }
326            2 => {
327                let y = as_i32(scope, pos, 0);
328                let x = as_i32(scope, pos, 1);
329                barrier.__expand_tma_load_2d_method(scope, self, shared_memory, y, x);
330            }
331            3 => {
332                let z = as_i32(scope, pos, 0);
333                let y = as_i32(scope, pos, 1);
334                let x = as_i32(scope, pos, 2);
335                barrier.__expand_tma_load_3d_method(scope, self, shared_memory, z, y, x);
336            }
337            4 => {
338                let w = as_i32(scope, pos, 0);
339                let z = as_i32(scope, pos, 1);
340                let y = as_i32(scope, pos, 2);
341                let x = as_i32(scope, pos, 3);
342                barrier.__expand_tma_load_4d_method(scope, self, shared_memory, w, z, y, x);
343            }
344            5 => {
345                let v = as_i32(scope, pos, 0);
346                let w = as_i32(scope, pos, 1);
347                let z = as_i32(scope, pos, 2);
348                let y = as_i32(scope, pos, 3);
349                let x = as_i32(scope, pos, 4);
350                barrier.__expand_tma_load_5d_method(scope, self, shared_memory, v, w, z, y, x);
351            }
352            _ => panic!("TMA only supports 1D-5D loads"),
353        }
354    }
355}
356
357impl<T: CubePrimitive, N: CubePrimitive + Coordinates> ViewOperationsMut<T, Sequence<N>>
358    for TensorMap<T, Tiled>
359{
360}
361impl<T: CubePrimitive, N: CubePrimitive + Coordinates> ViewOperationsMutExpand<T, Sequence<N>>
362    for NativeExpand<TensorMap<T, Tiled>>
363{
364    fn __expand_write_method(
365        &self,
366        _scope: &Scope,
367        _pos: SequenceExpand<N>,
368        _value: <T as CubeType>::ExpandType,
369    ) {
370        unimplemented!("Can't write to tensor map");
371    }
372
373    fn __expand_write_checked_method(
374        &self,
375        _scope: &Scope,
376        _pos: SequenceExpand<N>,
377        _value: <T as CubeType>::ExpandType,
378    ) {
379        unimplemented!("Can't write to tensor map");
380    }
381
382    fn __expand_as_linear_slice_mut_method(
383        &self,
384        _scope: &Scope,
385        _pos: SequenceExpand<N>,
386        _end: SequenceExpand<N>,
387    ) -> &mut SliceExpand<T> {
388        unimplemented!("Can't write to tensor map");
389    }
390
391    #[allow(unused_parens)]
392    fn __expand_tensor_map_store_method(
393        &self,
394        scope: &Scope,
395        shared_memory: &SliceExpand<T>,
396        pos: SequenceExpand<N>,
397    ) {
398        let mut this = *self;
399        let rank = pos.len();
400        let pos = &pos;
401        match rank {
402            1 => {
403                let x = as_i32(scope, pos, 0);
404                tma_store_1d::expand(scope, shared_memory, &mut this, x);
405            }
406            2 => {
407                let y = as_i32(scope, pos, 0);
408                let x = as_i32(scope, pos, 1);
409                tma_store_2d::expand(scope, shared_memory, &mut this, y, x);
410            }
411            3 => {
412                let z = as_i32(scope, pos, 0);
413                let y = as_i32(scope, pos, 1);
414                let x = as_i32(scope, pos, 2);
415                tma_store_3d::expand(scope, shared_memory, &mut this, z, y, x);
416            }
417            4 => {
418                let w = as_i32(scope, pos, 0);
419                let z = as_i32(scope, pos, 1);
420                let y = as_i32(scope, pos, 2);
421                let x = as_i32(scope, pos, 3);
422                tma_store_4d::expand(scope, shared_memory, &mut this, w, z, y, x);
423            }
424            5 => {
425                let v = as_i32(scope, pos, 0);
426                let w = as_i32(scope, pos, 1);
427                let z = as_i32(scope, pos, 2);
428                let y = as_i32(scope, pos, 3);
429                let x = as_i32(scope, pos, 4);
430                tma_store_5d::expand(scope, shared_memory, &mut this, v, w, z, y, x);
431            }
432            _ => panic!("TMA store supports 1D-5D loads"),
433        }
434    }
435}
436
437impl<T: CubePrimitive, P: CubePrimitive + Coordinates, O: CubePrimitive + Coordinates>
438    ViewOperations<T, (Sequence<P>, Sequence<O>)> for TensorMap<T, Im2col>
439{
440}
441impl<T: CubePrimitive, P: CubePrimitive + Coordinates, O: CubePrimitive + Coordinates>
442    ViewOperationsExpand<T, (Sequence<P>, Sequence<O>)> for NativeExpand<TensorMap<T, Im2col>>
443{
444    fn __expand_read_method(
445        &self,
446        _scope: &Scope,
447        _pos: (SequenceExpand<P>, SequenceExpand<O>),
448    ) -> <T as CubeType>::ExpandType {
449        unimplemented!("Can't read from tensor map");
450    }
451
452    fn __expand_read_checked_method(
453        &self,
454        _scope: &Scope,
455        _pos: (SequenceExpand<P>, SequenceExpand<O>),
456    ) -> <T as CubeType>::ExpandType {
457        unimplemented!("Can't read from tensor map");
458    }
459
460    fn __expand_read_masked_method(
461        &self,
462        _scope: &Scope,
463        _pos: (SequenceExpand<P>, SequenceExpand<O>),
464        _mask_value: <T as CubeType>::ExpandType,
465    ) -> <T as CubeType>::ExpandType {
466        unimplemented!("Can't read from tensor map");
467    }
468
469    fn __expand_read_unchecked_method(
470        &self,
471        _scope: &Scope,
472        _pos: (SequenceExpand<P>, SequenceExpand<O>),
473    ) -> <T as CubeType>::ExpandType {
474        unimplemented!("Can't read from tensor map");
475    }
476
477    fn __expand_as_linear_slice_method(
478        &self,
479        _scope: &Scope,
480        _pos: (SequenceExpand<P>, SequenceExpand<O>),
481        _end: (SequenceExpand<P>, SequenceExpand<O>),
482    ) -> &SliceExpand<T> {
483        unimplemented!("Can't read from tensor map");
484    }
485
486    fn __expand_shape_method(&self, _scope: &Scope) -> (SequenceExpand<P>, SequenceExpand<O>) {
487        unimplemented!("Can't read from tensor map");
488    }
489
490    fn __expand_is_in_bounds_method(
491        &self,
492        _scope: &Scope,
493        _pos: (SequenceExpand<P>, SequenceExpand<O>),
494    ) -> NativeExpand<bool> {
495        // Bounds checks are done in hardware, so treat them as always in bounds for the kernels
496        true.into()
497    }
498
499    #[allow(unused_parens)]
500    fn __expand_tensor_map_load_method(
501        &self,
502        scope: &Scope,
503        barrier: &NativeExpand<Barrier>,
504        shared_memory: &mut SliceExpand<T>,
505        pos: (SequenceExpand<P>, SequenceExpand<O>),
506    ) {
507        let (pos, offs) = &pos;
508        let rank = pos.len();
509
510        match rank {
511            3 => {
512                let n = as_i32(scope, pos, 0);
513                let w = as_i32(scope, pos, 1);
514                let c = as_i32(scope, pos, 2);
515                let x = as_u16(scope, offs, 0);
516                barrier.__expand_tma_load_im2col_3d_method(scope, self, shared_memory, n, w, c, x);
517            }
518            4 => {
519                let n = as_i32(scope, pos, 0);
520                let h = as_i32(scope, pos, 1);
521                let w = as_i32(scope, pos, 2);
522                let c = as_i32(scope, pos, 3);
523                let y = as_u16(scope, offs, 0);
524                let x = as_u16(scope, offs, 1);
525                barrier.__expand_tma_load_im2col_4d_method(
526                    scope,
527                    self,
528                    shared_memory,
529                    n,
530                    h,
531                    w,
532                    c,
533                    y,
534                    x,
535                );
536            }
537            5 => {
538                let n = as_i32(scope, pos, 0);
539                let d = as_i32(scope, pos, 1);
540                let h = as_i32(scope, pos, 2);
541                let w = as_i32(scope, pos, 3);
542                let c = as_i32(scope, pos, 4);
543                let z = as_u16(scope, offs, 0);
544                let y = as_u16(scope, offs, 1);
545                let x = as_u16(scope, offs, 2);
546                barrier.__expand_tma_load_im2col_5d_method(
547                    scope,
548                    self,
549                    shared_memory,
550                    n,
551                    d,
552                    h,
553                    w,
554                    c,
555                    z,
556                    y,
557                    x,
558                );
559            }
560            _ => panic!("TMA im2col only supports 3D-5D loads"),
561        }
562    }
563}
cubecl_std/tensor/view/operations/tensor_map.rs

cubecl_std/tensor/view/operations/
tensor_map.rs
