1use std::marker::PhantomData;
2
3use super::*;
4use crate::tensor::layout::{Coordinates, VirtualLayout, VirtualLayoutExpand};
5use cubecl::prelude::*;
6use cubecl_core::{self as cubecl, prelude::barrier::Barrier};
7
8#[derive(CubeType)]
9pub struct VirtualView<T: CubePrimitive, C: Coordinates, S: Coordinates, V: ViewOperations<T, S>> {
10 #[allow(unused)]
11 view: V,
12 #[allow(unused)]
13 layout: VirtualLayout<C, S>,
14 #[cube(comptime)]
15 _ty: PhantomData<T>,
16}
17
18#[cube]
19impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V: ViewOperations<T, S>>
20 VirtualView<T, C, S, V>
21{
22 pub fn new(view: V, layout: VirtualLayout<C, S>) -> Self {
23 VirtualView::<T, C, S, V> {
24 view,
25 layout,
26 _ty: PhantomData,
27 }
28 }
29}
30
31impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V: ViewOperations<T, S>>
32 VirtualViewExpand<T, C, S, V>
33{
34 pub fn new(view: V::ExpandType, layout: VirtualLayoutExpand<C, S>) -> Self {
35 VirtualViewExpand::<T, C, S, V> {
36 view,
37 layout,
38 _ty: PhantomData,
39 }
40 }
41}
42
43#[derive(CubeType)]
44pub struct VirtualViewMut<
45 T: CubePrimitive,
46 C: Coordinates,
47 S: Coordinates,
48 V: ViewOperationsMut<T, S>,
49> {
50 #[allow(unused)]
51 view: V,
52 #[allow(unused)]
53 layout: VirtualLayout<C, S>,
54 #[cube(comptime)]
55 _ty: PhantomData<T>,
56}
57
58#[cube]
59impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V: ViewOperationsMut<T, S>>
60 VirtualViewMut<T, C, S, V>
61{
62 pub fn new(view: V, layout: VirtualLayout<C, S>) -> Self {
63 VirtualViewMut::<T, C, S, V> {
64 view,
65 layout,
66 _ty: PhantomData,
67 }
68 }
69}
70
71impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V: ViewOperationsMut<T, S>>
72 VirtualViewMutExpand<T, C, S, V>
73{
74 pub fn new(view: V::ExpandType, layout: VirtualLayoutExpand<C, S>) -> Self {
75 VirtualViewMutExpand::<T, C, S, V> {
76 view,
77 layout,
78 _ty: PhantomData,
79 }
80 }
81}
82
83macro_rules! impl_virtual_read {
84 ($ty: ident, $expand: ident, $trait: ident) => {
85 impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V> Vectorized for $ty<T, C, S, V> where
86 V: $trait<T, S>
87 {
88 }
89 impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V> VectorizedExpand
90 for $expand<T, C, S, V>
91 where
92 V: $trait<T, S>,
93 {
94 fn vector_size(&self) -> VectorSize {
95 self.view.vector_size()
96 }
97 }
98
99 impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V> ViewOperations<T, C>
100 for $ty<T, C, S, V>
101 where
102 V: $trait<T, S>,
103 {
104 }
105
106 impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V> ViewOperationsExpand<T, C>
107 for $expand<T, C, S, V>
108 where
109 V: $trait<T, S>,
110 {
111 fn __expand_read_method(&self, scope: &Scope, pos: <C>::ExpandType) -> <T>::ExpandType {
112 let pos = self
113 .layout
114 .clone()
115 .__expand_to_source_pos_method(scope, pos);
116 self.view.__expand_read_method(scope, pos)
117 }
118
119 fn __expand_read_checked_method(
120 &self,
121 scope: &Scope,
122 pos: <C>::ExpandType,
123 ) -> <T>::ExpandType {
124 let (read_pos, in_bounds) = self
125 .layout
126 .clone()
127 .__expand_to_source_pos_checked_method(scope, pos);
128 let zero = T::__expand_cast_from(scope, 0.into());
129 let value = self.view.__expand_read_checked_method(scope, read_pos);
130 select::expand::<T>(scope, in_bounds, value, zero)
131 }
132
133 fn __expand_read_masked_method(
134 &self,
135 scope: &Scope,
136 pos: <C>::ExpandType,
137 mask_value: <T>::ExpandType,
138 ) -> <T>::ExpandType {
139 let (read_pos, in_bounds) = self
140 .layout
141 .clone()
142 .__expand_to_source_pos_checked_method(scope, pos);
143 let value = self.view.__expand_read_checked_method(scope, read_pos);
144 select::expand::<T>(scope, in_bounds, value, mask_value)
145 }
146
147 fn __expand_read_unchecked_method(
148 &self,
149 scope: &Scope,
150 pos: <C>::ExpandType,
151 ) -> <T>::ExpandType {
152 let pos = self
153 .layout
154 .clone()
155 .__expand_to_source_pos_method(scope, pos);
156 self.view.__expand_read_unchecked_method(scope, pos)
157 }
158
159 fn __expand_as_linear_slice_method(
160 &self,
161 scope: &Scope,
162 pos: <C>::ExpandType,
163 end: <C>::ExpandType,
164 ) -> &SliceExpand<T> {
165 let pos = self
166 .layout
167 .clone()
168 .__expand_to_source_pos_method(scope, pos);
169 let end = self
170 .layout
171 .clone()
172 .__expand_to_source_pos_method(scope, end);
173 self.view.__expand_as_linear_slice_method(scope, pos, end)
174 }
175
176 fn __expand_shape_method(&self, scope: &Scope) -> <C>::ExpandType {
177 self.layout.clone().__expand_shape_method(scope)
178 }
179
180 fn __expand_is_in_bounds_method(
181 &self,
182 scope: &Scope,
183 pos: C::ExpandType,
184 ) -> NativeExpand<bool> {
185 let (pos, in_bounds_layout) = self
186 .layout
187 .clone()
188 .__expand_to_source_pos_checked_method(scope, pos);
189 let in_bounds_view = self.view.__expand_is_in_bounds_method(scope, pos);
190 in_bounds_layout.__expand_and_method(scope, in_bounds_view)
191 }
192
193 fn __expand_tensor_map_load_method(
194 &self,
195 scope: &Scope,
196 barrier: &NativeExpand<Barrier>,
197 shared_memory: &mut SliceExpand<T>,
198 pos: C::ExpandType,
199 ) {
200 let pos = self
201 .layout
202 .clone()
203 .__expand_to_source_pos_method(scope, pos);
204 self.view
205 .__expand_tensor_map_load_method(scope, barrier, shared_memory, pos);
206 }
207 }
208 };
209}
210
211impl_virtual_read!(VirtualView, VirtualViewExpand, ViewOperations);
212impl_virtual_read!(VirtualViewMut, VirtualViewMutExpand, ViewOperationsMut);
213
214impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V> ViewOperationsMut<T, C>
215 for VirtualViewMut<T, C, S, V>
216where
217 V: ViewOperationsMut<T, S>,
218{
219}
220
221impl<T: CubePrimitive, C: Coordinates, S: Coordinates, V> ViewOperationsMutExpand<T, C>
222 for VirtualViewMutExpand<T, C, S, V>
223where
224 V: ViewOperationsMut<T, S>,
225{
226 fn __expand_write_method(&self, scope: &Scope, pos: <C>::ExpandType, value: <T>::ExpandType) {
227 let pos = self
228 .layout
229 .clone()
230 .__expand_to_source_pos_method(scope, pos);
231 self.view.__expand_write_method(scope, pos, value);
232 }
233
234 fn __expand_write_checked_method(
235 &self,
236 scope: &Scope,
237 pos: <C>::ExpandType,
238 value: <T>::ExpandType,
239 ) {
240 let (pos, in_bounds) = self
241 .layout
242 .clone()
243 .__expand_to_source_pos_checked_method(scope, pos);
244 if_expand(scope, in_bounds, |scope| {
245 self.view.__expand_write_checked_method(scope, pos, value);
246 });
247 }
248
249 fn __expand_as_linear_slice_mut_method(
250 &self,
251 scope: &Scope,
252 pos: <C>::ExpandType,
253 end: <C>::ExpandType,
254 ) -> &mut SliceExpand<T> {
255 let pos = self
256 .layout
257 .clone()
258 .__expand_to_source_pos_method(scope, pos);
259 let end = self
260 .layout
261 .clone()
262 .__expand_to_source_pos_method(scope, end);
263 self.view
264 .__expand_as_linear_slice_mut_method(scope, pos, end)
265 }
266
267 fn __expand_tensor_map_store_method(
268 &self,
269 scope: &Scope,
270 shared_memory: &SliceExpand<T>,
271 pos: C::ExpandType,
272 ) {
273 let pos = self
274 .layout
275 .clone()
276 .__expand_to_source_pos_method(scope, pos);
277 self.view
278 .__expand_tensor_map_store_method(scope, shared_memory, pos);
279 }
280}