Skip to main content

cubecl_std/tensor/view/operations/
view.rs

1use super::*;
2use crate::tensor::{View, ViewExpand};
3use crate::tensor::{ViewMut, ViewMutExpand, layout::Coordinates};
4use cubecl::prelude::*;
5use cubecl_core::{self as cubecl, prelude::barrier::Barrier};
6
7macro_rules! impl_base {
8    ($ty: ident, $expand: ident) => {
9        impl<'a, T: CubePrimitive, C: Coordinates + 'a> Vectorized for $ty<'a, T, C> {}
10        impl<'a, T: CubePrimitive, C: Coordinates + 'a> VectorizedExpand for $expand<'a, T, C> {
11            fn vector_size(&self) -> VectorSize {
12                self.inner.vector_size()
13            }
14        }
15
16        impl<'a, T: CubePrimitive, C: Coordinates + 'a> ViewOperations<T, C> for $ty<'a, T, C> {}
17        impl<'a, T: CubePrimitive, C: Coordinates + 'a> ViewOperationsExpand<T, C>
18            for $expand<'a, T, C>
19        {
20            fn __expand_read_method(&self, scope: &Scope, pos: <C>::ExpandType) -> <T>::ExpandType {
21                $expand::__expand_read_method(self, scope, pos)
22            }
23
24            fn __expand_read_checked_method(
25                &self,
26                scope: &Scope,
27                pos: <C>::ExpandType,
28            ) -> <T>::ExpandType {
29                $expand::__expand_read_checked_method(self, scope, pos)
30            }
31
32            fn __expand_read_masked_method(
33                &self,
34                scope: &Scope,
35                pos: <C>::ExpandType,
36                mask_value: <T>::ExpandType,
37            ) -> <T>::ExpandType {
38                $expand::__expand_read_masked_method(self, scope, pos, mask_value)
39            }
40
41            fn __expand_read_unchecked_method(
42                &self,
43                scope: &Scope,
44                pos: <C>::ExpandType,
45            ) -> <T>::ExpandType {
46                $expand::__expand_read_unchecked_method(self, scope, pos)
47            }
48
49            fn __expand_as_linear_slice_method(
50                &self,
51                scope: &Scope,
52                pos: <C>::ExpandType,
53                end: <C>::ExpandType,
54            ) -> &SliceExpand<T> {
55                $expand::__expand_as_linear_slice_inner_method(self, scope, pos, end)
56            }
57
58            fn __expand_shape_method(&self, scope: &Scope) -> <C>::ExpandType {
59                $expand::__expand_shape_method(self, scope)
60            }
61
62            fn __expand_is_in_bounds_method(
63                &self,
64                scope: &Scope,
65                pos: <C>::ExpandType,
66            ) -> NativeExpand<bool> {
67                $expand::__expand_is_in_bounds_method(self, scope, pos)
68            }
69
70            fn __expand_tensor_map_load_method(
71                &self,
72                scope: &Scope,
73                barrier: &NativeExpand<Barrier>,
74                shared_memory: &mut SliceExpand<T>,
75                pos: C::ExpandType,
76            ) {
77                $expand::__expand_tensor_map_load_method(self, scope, barrier, shared_memory, pos)
78            }
79        }
80    };
81}
82
83impl_base!(View, ViewExpand);
84impl_base!(ViewMut, ViewMutExpand);
85
86impl<'a, T: CubePrimitive, C: Coordinates + 'a> ViewOperationsMut<T, C> for ViewMut<'a, T, C> {}
87impl<'a, T: CubePrimitive, C: Coordinates + 'a> ViewOperationsMutExpand<T, C>
88    for ViewMutExpand<'a, T, C>
89{
90    fn __expand_write_method(&self, scope: &Scope, pos: <C>::ExpandType, value: <T>::ExpandType) {
91        self.inner.__expand_write_method(scope, pos, value);
92    }
93
94    fn __expand_write_checked_method(
95        &self,
96        scope: &Scope,
97        pos: <C>::ExpandType,
98        value: <T>::ExpandType,
99    ) {
100        self.inner.__expand_write_checked_method(scope, pos, value);
101    }
102
103    fn __expand_as_linear_slice_mut_method(
104        &self,
105        scope: &Scope,
106        pos: <C>::ExpandType,
107        end: <C>::ExpandType,
108    ) -> &mut SliceExpand<T> {
109        self.inner
110            .__expand_as_linear_slice_mut_method(scope, pos, end)
111    }
112
113    fn __expand_tensor_map_store_method(
114        &self,
115        scope: &Scope,
116        shared_memory: &SliceExpand<T>,
117        pos: C::ExpandType,
118    ) {
119        self.inner
120            .__expand_tensor_map_store_method(scope, shared_memory, pos)
121    }
122}