1use cubecl_core::{prelude::*, unexpanded};
2
3use crate::tensor::{View, ViewExpand, ViewMut, ViewMutExpand, layout::*};
4
5type ArrayExpand<T> = NativeExpand<Array<T>>;
6
7pub trait AsView<E: CubePrimitive>:
8 CubeType<ExpandType: AsViewExpand<E, SourceCoords = Self::SourceCoords>>
9{
10 type SourceCoords: Coordinates;
11
12 #[allow(unused)]
13 fn view<C: Coordinates + 'static>(
14 &self,
15 layout: impl Into<VirtualLayout<C, Self::SourceCoords>>,
16 ) -> View<'_, E, C> {
17 unexpanded!()
18 }
19
20 fn __expand_view<'a, C: Coordinates + 'static>(
21 scope: &Scope,
22 this: &'a Self::ExpandType,
23 layout: VirtualLayoutExpand<C, Self::SourceCoords>,
24 ) -> ViewExpand<'a, E, C> {
25 this.__expand_view_method(scope, layout)
26 }
27}
28
29pub trait AsViewExpand<E: CubePrimitive> {
30 type SourceCoords: Coordinates;
31
32 #[allow(unused)]
33 fn __expand_view_method<C: Coordinates + 'static>(
34 &self,
35 scope: &Scope,
36 layout: VirtualLayoutExpand<C, Self::SourceCoords>,
37 ) -> ViewExpand<'_, E, C>;
38}
39
40pub trait AsViewMut<E: CubePrimitive>: AsView<E> {
41 #[allow(unused)]
42 fn view_mut<C: Coordinates + 'static>(
43 &mut self,
44 layout: impl Into<VirtualLayout<C, Self::SourceCoords>>,
45 ) -> ViewMut<'_, E, C> {
46 unexpanded!()
47 }
48}
49
50pub trait AsViewMutExpand<E: CubePrimitive>: AsViewExpand<E> {
51 #[allow(clippy::too_many_arguments)]
52 fn __expand_view_mut_method<C: Coordinates + 'static>(
53 &mut self,
54 scope: &Scope,
55 layout: VirtualLayoutExpand<C, Self::SourceCoords>,
56 ) -> ViewMutExpand<'_, E, C>;
57}
58
59macro_rules! impl_as_view {
60 ($ty: ty, $expand: ty, $coords: ty) => {
61 impl<E: CubePrimitive> AsView<E> for $ty {
62 type SourceCoords = $coords;
63 }
64 impl<E: CubePrimitive> AsViewExpand<E> for $expand {
65 type SourceCoords = $coords;
66 fn __expand_view_method<C: Coordinates + 'static>(
67 &self,
68 scope: &Scope,
69 layout: VirtualLayoutExpand<C, $coords>,
70 ) -> super::ViewExpand<'_, E, C> {
71 View::__expand_new::<&[E], $coords>(
72 scope,
73 self.__expand_as_slice_method(scope),
74 layout,
75 )
76 }
77 }
78
79 impl<E: CubePrimitive> AsViewMut<E> for $ty {}
80 impl<E: CubePrimitive> AsViewMutExpand<E> for $expand {
81 fn __expand_view_mut_method<C: Coordinates + 'static>(
82 &mut self,
83 scope: &Scope,
84 layout: VirtualLayoutExpand<C, $coords>,
85 ) -> super::ViewMutExpand<'_, E, C> {
86 ViewMut::__expand_new::<&mut [E], $coords>(
87 scope,
88 self.__expand_as_mut_slice_method(scope),
89 layout,
90 )
91 }
92 }
93 };
94}
95
96impl_as_view!(Array<E>, ArrayExpand<E>, Coords1d);
97impl_as_view!(Tensor<E>, TensorExpand<E>, Coords1d);
98impl_as_view!(Shared<[E]>, SharedExpand<[E]>, Coords1d);
99
100impl<E: CubePrimitive> AsView<E> for [E] {
101 type SourceCoords = Coords1d;
102 fn view<C: Coordinates + 'static>(
103 &self,
104 layout: impl Into<VirtualLayout<C, Coords1d>>,
105 ) -> View<'_, E, C> {
106 View::new::<&[E], Coords1d>(self, layout)
107 }
108}
109
110impl<E: CubePrimitive> AsViewExpand<E> for SliceExpand<E> {
111 type SourceCoords = Coords1d;
112 fn __expand_view_method<C: Coordinates + 'static>(
113 &self,
114 scope: &Scope,
115 layout: VirtualLayoutExpand<C, Self::SourceCoords>,
116 ) -> ViewExpand<'_, E, C> {
117 View::__expand_new::<&[E], Self::SourceCoords>(scope, self, layout)
118 }
119}
120
121impl<E: CubePrimitive> AsViewMut<E> for [E] {
122 fn view_mut<C: Coordinates + 'static>(
123 &mut self,
124 layout: impl Into<VirtualLayout<C, Coords1d>>,
125 ) -> ViewMut<'_, E, C> {
126 ViewMut::new::<&mut [E], Coords1d>(self, layout)
127 }
128}
129impl<E: CubePrimitive> AsViewMutExpand<E> for SliceExpand<E> {
130 fn __expand_view_mut_method<C: Coordinates + 'static>(
131 &mut self,
132 scope: &Scope,
133 layout: VirtualLayoutExpand<C, Self::SourceCoords>,
134 ) -> ViewMutExpand<'_, E, C> {
135 ViewMut::__expand_new::<&mut [E], Coords1d>(scope, self, layout)
136 }
137}
138
139macro_rules! as_view_tensor_map {
140 ($($dim: literal),*) => {
141 paste::paste! {
142 pub trait AsTensorView<E: CubePrimitive>:
143 CubeType<ExpandType: AsTensorViewExpand<E>>
144 {
145 $(
146 #[allow(unused)]
147 fn [<view_ $dim>]<C: Coordinates + 'static>(
148 &self,
149 layout: impl Into<VirtualLayout<C, [<Coords $dim>]>>,
150 ) -> View<'_, E, C> {
151 unexpanded!()
152 }
153
154 fn [<__expand_view_ $dim>]<C: Coordinates + 'static>(
155 scope: &Scope,
156 this: Self::ExpandType,
157 layout: VirtualLayoutExpand<C, [<Coords $dim>]>,
158 ) -> ViewExpand<'_, E, C> {
159 this.[<__expand_view_ $dim _method>](scope, layout)
160 }
161 )*
162 }
163
164 pub trait AsTensorViewExpand<E: CubePrimitive> {
165 $(
166 #[allow(unused)]
167 fn [<__expand_view_ $dim _method>]<C: Coordinates + 'static>(
168 self,
169 scope: &Scope,
170 layout: VirtualLayoutExpand<C, [<Coords $dim>]>,
171 ) -> ViewExpand<'_, E, C>;
172 )*
173 }
174
175 pub trait AsTensorViewMut<E: CubePrimitive>: AsTensorView<E> {
176 $(
177 #[allow(unused)]
178 fn [<view_mut_ $dim>]<C: Coordinates + 'static>(
179 &mut self,
180 layout: impl Into<VirtualLayout<C, [<Coords $dim>]>>,
181 ) -> ViewMut<'_, E, C> {
182 unexpanded!()
183 }
184 )*
185 }
186
187 pub trait AsTensorViewMutExpand<E: CubePrimitive>: AsTensorViewExpand<E> {
188 $(
189 #[allow(clippy::too_many_arguments)]
190 fn [<__expand_view_mut_ $dim _method>]<C: Coordinates + 'static>(
191 self,
192 scope: &Scope,
193 layout: VirtualLayoutExpand<C, [<Coords $dim>]>,
194 ) -> ViewMutExpand<'_, E, C>;
195 )*
196 }
197
198 impl<E: CubePrimitive> AsTensorView<E> for TensorMap<E, Tiled> {}
199 impl<E: CubePrimitive> AsTensorViewExpand<E> for NativeExpand<TensorMap<E, Tiled>> {
200 $(
201 fn [<__expand_view_ $dim _method>]<C: Coordinates + 'static>(
202 self,
203 scope: &Scope,
204 layout: VirtualLayoutExpand<C, [<Coords $dim>]>,
205 ) -> super::ViewExpand<'_, E, C> {
206 View::__expand_new::<TensorMap<E, Tiled>, [<Coords $dim>]>(scope, self, layout)
207 }
208 )*
209 }
210
211 impl<E: CubePrimitive> AsTensorViewMut<E> for TensorMap<E, Tiled> {}
212 impl<E: CubePrimitive> AsTensorViewMutExpand<E> for NativeExpand<TensorMap<E, Tiled>> {
213 $(
214 fn [<__expand_view_mut_ $dim _method>]<C: Coordinates + 'static>(
215 self,
216 scope: &Scope,
217 layout: VirtualLayoutExpand<C, [<Coords $dim>]>,
218 ) -> super::ViewMutExpand<'_, E, C> {
219 ViewMut::__expand_new::<TensorMap<E, Tiled>, [<Coords $dim>]>(scope, self, layout)
220 }
221 )*
222 }
223 }
224 };
225}
226
227as_view_tensor_map!(1d, 2d, 3d, 4d, 5d, 1i, 2i, 3i, 4i, 5i);