cubecl_std/tensor/view/
base.rs

1use std::marker::PhantomData;
2
3use cubecl::prelude::*;
4use cubecl_core::{self as cubecl, ir::VectorSize, prelude::barrier::Barrier, unexpanded};
5
6use crate::tensor::{
7    ViewOperations, ViewOperationsExpand, ViewOperationsMut, ViewOperationsMutExpand, VirtualView,
8    VirtualViewMut,
9    layout::{Coordinates, Layout, VirtualLayout, VirtualLayoutExpand, slice::SliceLayout},
10};
11
12/// A conceptual view of an underlying linear storage.
13/// Allows abstract indexing in multiple dimensions, without having to know the data layout or
14/// location.
15#[derive(Clone, Copy)]
16pub struct View<'a, E: CubePrimitive, C: Coordinates> {
17    _layout: PhantomData<C>,
18    _ty: PhantomData<E>,
19    _lifetime: PhantomData<&'a ()>,
20}
21
22/// Expand type of [`View`]
23#[derive(Clone, Copy)]
24pub struct ViewExpand<'a, E: CubePrimitive, C: Coordinates> {
25    pub(super) inner: &'a (dyn ViewOperationsExpand<E, C> + 'a),
26}
27
28/// Mutable view
29/// Note: `Clone` and `Copy` should ideally not be there, but are required for good ergonomics until
30/// `Reborrow` and `CoerceShared` are implemented and stabilized.
31#[derive(Clone, Copy)]
32pub struct ViewMut<'a, E: CubePrimitive, C: Coordinates> {
33    _layout: PhantomData<C>,
34    _ty: PhantomData<E>,
35    _lifetime: PhantomData<&'a mut ()>,
36}
37
38/// Expand type of [`ViewMutExpand`]
39#[derive(Clone, Copy)]
40pub struct ViewMutExpand<'a, E: CubePrimitive, C: Coordinates> {
41    pub(super) inner: &'a (dyn ViewOperationsMutExpand<E, C> + 'a),
42}
43
44macro_rules! impl_cube_type {
45    ($ty: ident, $expand: ident) => {
46        impl<'a, E: CubePrimitive, C: Coordinates + 'a> CubeType for $ty<'a, E, C> {
47            type ExpandType = $expand<'a, E, C>;
48        }
49
50        impl<'a, E: CubePrimitive, C: Coordinates> IntoExpand for $expand<'a, E, C> {
51            type Expand = $expand<'a, E, C>;
52
53            fn into_expand(self, _: &Scope) -> Self::Expand {
54                self
55            }
56        }
57
58        impl<'a, E: CubePrimitive, C: Coordinates> ExpandTypeClone for $expand<'a, E, C> {
59            fn clone_unchecked(&self) -> Self {
60                self.clone()
61            }
62        }
63
64        impl<'a, E: CubePrimitive, C: Coordinates> IntoMut for $expand<'a, E, C> {
65            fn into_mut(self, _scope: &Scope) -> Self {
66                self
67            }
68        }
69
70        impl<'a, E: CubePrimitive, C: Coordinates> CubeDebug for $expand<'a, E, C> {}
71
72        impl<'a, E: CubePrimitive, C: Coordinates> AsRefExpand for $expand<'a, E, C> {
73            fn __expand_ref_method(&self, _: &Scope) -> &Self {
74                self
75            }
76        }
77        impl<'a, E: CubePrimitive, C: Coordinates> AsMutExpand for $expand<'a, E, C> {
78            fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
79                self
80            }
81        }
82        impl<'a, E: CubePrimitive, C: Coordinates> DerefExpand for $expand<'a, E, C> {
83            type Target = Self;
84
85            fn __expand_deref_method(&self, _: &Scope) -> Self::Target {
86                self.clone()
87            }
88        }
89    };
90}
91
92impl_cube_type!(View, ViewExpand);
93impl_cube_type!(ViewMut, ViewMutExpand);
94
95impl<'a, E: CubePrimitive, C: Coordinates + 'a> View<'a, E, C> {
96    /// Create a new tensor view from an underlying concrete storage and a layout to map it into
97    /// the target coordinate space
98    #[allow(unused_variables)]
99    pub fn new<V: ViewOperations<E, S> + 'a, S: Coordinates>(
100        view: V,
101        layout: impl Into<VirtualLayout<C, S>>,
102    ) -> Self {
103        View {
104            _layout: PhantomData,
105            _ty: PhantomData,
106            _lifetime: PhantomData,
107        }
108    }
109
110    /// Expand function for [`View::new`]
111    pub fn __expand_new<V: ViewOperations<E, S> + 'a, S: Coordinates + 'a>(
112        scope: &Scope,
113        view: V::ExpandType,
114        layout: VirtualLayoutExpand<C, S>,
115    ) -> ViewExpand<'a, E, C> {
116        ViewExpand::new(
117            scope,
118            VirtualView::<E, C, S, V>::__expand_new(scope, view, layout),
119        )
120    }
121}
122
123impl<'a, E: CubePrimitive, C: Coordinates + 'a> ViewMut<'a, E, C> {
124    /// Create a new tensor view from an underlying concrete storage and a layout to map it into
125    /// the target coordinate space
126    #[allow(unused_variables)]
127    pub fn new<V: ViewOperationsMut<E, S> + 'a, S: Coordinates>(
128        view: V,
129        layout: impl Into<VirtualLayout<C, S>>,
130    ) -> Self {
131        ViewMut {
132            _layout: PhantomData,
133            _ty: PhantomData,
134            _lifetime: PhantomData,
135        }
136    }
137
138    /// Expand function for [`View::new`]
139    pub fn __expand_new<V: ViewOperationsMut<E, S> + 'a, S: Coordinates + 'a>(
140        scope: &Scope,
141        view: V::ExpandType,
142        layout: VirtualLayoutExpand<C, S>,
143    ) -> ViewMutExpand<'a, E, C> {
144        ViewMutExpand::new(
145            scope,
146            VirtualViewMut::<E, C, S, V>::__expand_new(scope, view, layout),
147        )
148    }
149}
150
151macro_rules! impl_read {
152    ($ty: ident, $expand: ident) => {
153        impl<'a, E: CubePrimitive, C: Coordinates + 'a> $ty<'a, E, C> {
154            pub fn view<T: Coordinates + 'a>(
155                self,
156                _layout: impl Into<VirtualLayout<T, C>>,
157            ) -> $ty<'a, E, T> {
158                unexpanded!()
159            }
160
161            pub fn __expand_view<T: Coordinates + 'a>(
162                scope: &Scope,
163                this: $expand<'a, E, C>,
164                layout: VirtualLayoutExpand<T, C>,
165            ) -> $expand<'a, E, T> {
166                this.__expand_view_method(scope, layout)
167            }
168        }
169
170        impl<'a, E: CubePrimitive, C: Coordinates + 'a> $expand<'a, E, C> {
171            pub fn __expand_view_method<T: Coordinates + 'a>(
172                self,
173                scope: &Scope,
174                layout: VirtualLayoutExpand<T, C>,
175            ) -> $expand<'a, E, T> {
176                $ty::__expand_new::<$ty<'a, E, C>, C>(scope, self, layout)
177            }
178        }
179
180        #[cube]
181        impl<'a, E: CubePrimitive, C: Coordinates> $ty<'a, E, C> {
182            /// Calls [`Layout::shape`] on the view's layout
183            pub fn shape(&self) -> C {
184                intrinsic!(|scope| self.inner.__expand_shape_method(scope))
185            }
186
187            /// Calls [`Layout::is_in_bounds`] on the view's layout
188            #[allow(unused_variables)]
189            pub fn is_in_bounds(&self, pos: C) -> bool {
190                intrinsic!(|scope| self.inner.__expand_is_in_bounds_method(scope, pos))
191            }
192        }
193
194        #[cube]
195        impl<'a, E: CubePrimitive, C: Coordinates> $ty<'a, E, C> {
196            /// Read a value at `pos`. The layout handles translation into a concrete index.
197            #[allow(unused_variables)]
198            pub fn read(&self, pos: C) -> E {
199                intrinsic!(|scope| self.inner.__expand_read_method(scope, pos))
200            }
201
202            /// Read a value at `pos`. The layout handles translation into a concrete index.
203            /// Reading is done unchecked
204            #[allow(unused_variables)]
205            pub fn read_unchecked(&self, pos: C) -> E {
206                intrinsic!(|scope| self.inner.__expand_read_unchecked_method(scope, pos))
207            }
208
209            /// Read a value at `pos` if it's in bounds. The layout handles translation into a concrete index.
210            #[allow(unused_variables)]
211            pub fn read_checked(&self, pos: C) -> E {
212                intrinsic!(|scope| self.inner.__expand_read_checked_method(scope, pos))
213            }
214
215            /// Read a value at `pos` if it's in bounds, returning `mask_value` otherwise. The layout handles translation into a concrete index.
216            #[allow(unused_variables)]
217            pub fn read_masked(&self, pos: C, mask_value: E) -> E {
218                intrinsic!(|scope| self
219                    .inner
220                    .__expand_read_masked_method(scope, pos, mask_value))
221            }
222
223            /// Interpret this view as a linear slice encompassing the entire view.
224            ///
225            /// # Safety
226            ///
227            /// No checking is done on whether the slice is contiguous in memory.
228            pub fn as_linear_slice(&self) -> &'a [E] {
229                intrinsic!(|scope| {
230                    let shape = self.inner.__expand_shape_method(scope);
231                    let origin = C::__expand_from_int(scope, shape.clone_unchecked(), 0);
232                    // Inclusive end so clamping works correctly
233                    let one = C::__expand_from_int(scope, shape.clone_unchecked(), 1);
234                    let shape = C::__expand_max(scope, shape, one.clone_unchecked());
235                    let end = C::__expand_sub(scope, shape, one);
236                    let slice = self
237                        .inner
238                        .__expand_as_linear_slice_method(scope, origin, end);
239                    scope.create_kernel_ref(slice.expand.into())
240                })
241            }
242
243            pub fn vector_size(&self) -> comptime_type!(VectorSize) {
244                intrinsic!(|scope| self.inner.vector_size())
245            }
246        }
247
248        impl<'a, E: CubePrimitive, C: Coordinates> $expand<'a, E, C> {
249            pub(super) fn __expand_as_linear_slice_inner_method(
250                &self,
251                scope: &Scope,
252                pos: C::ExpandType,
253                end: C::ExpandType,
254            ) -> &SliceExpand<E> {
255                self.inner.__expand_as_linear_slice_method(scope, pos, end)
256            }
257        }
258
259        #[cube]
260        impl<'a, E: CubePrimitive, C: Coordinates + 'static> $ty<'a, E, C> {
261            /// Create a slice starting from `pos`, with `size`.
262            /// The layout handles translation into concrete indices.
263            /// Size will be clamped to the current layout size.
264            #[allow(unused_variables)]
265            pub fn slice(self, pos: C, size: C) -> $ty<'a, E, C> {
266                intrinsic!(|scope| self.slice(scope, pos, size, true))
267            }
268
269            /// Create a slice starting from `pos`, with `size`.
270            /// The layout handles translation into concrete indices.
271            /// Size and pos will be clamped to the current layout size.
272            /// #Safety
273            /// Access is always unchecked
274            #[allow(unused_variables)]
275            pub fn slice_unchecked(self, pos: C, size: C) -> $ty<'a, E, C> {
276                intrinsic!(|scope| self.slice(scope, pos, size, false))
277            }
278        }
279
280        impl<'a, E: CubePrimitive, C: Coordinates + 'static> $expand<'a, E, C> {
281            fn slice(
282                self,
283                scope: &Scope,
284                pos: C::ExpandType,
285                size: C::ExpandType,
286                checked: bool,
287            ) -> $expand<'a, E, C> {
288                let shape = self.__expand_shape_method(scope);
289                let pos = C::__expand_min(scope, pos, shape.clone_unchecked());
290                let max_size = C::__expand_sub(scope, shape, pos.clone_unchecked());
291                let size = C::__expand_min(scope, size, max_size);
292                let layout = SliceLayout::__expand_new(scope, pos, size, checked);
293                $ty::__expand_new::<$ty<'a, E, C>, _>(scope, self, layout.into())
294            }
295        }
296
297        #[cube]
298        impl<'a, E: CubePrimitive, C: Coordinates + 'a> $ty<'a, E, C> {
299            /// Execute a TMA load into shared memory, if the underlying storage supports it.
300            /// Panics if it's unsupported.
301            #[allow(unused_variables)]
302            pub fn tensor_map_load(&self, barrier: &Barrier, shared_memory: &mut [E], pos: C) {
303                intrinsic!(|scope| {
304                    self.inner
305                        .__expand_tensor_map_load_method(scope, barrier, shared_memory, pos)
306                })
307            }
308        }
309    };
310}
311
312impl_read!(View, ViewExpand);
313impl_read!(ViewMut, ViewMutExpand);
314
315impl<'a, E: CubePrimitive, C: Coordinates + 'a> ViewMut<'a, E, C> {
316    /// Reborrow this mutable view as a read-only [`View`] over the same storage. Lets a single
317    /// `ViewMut` carrier serve both read and write paths without a second handle.
318    #[allow(unused_variables)]
319    pub fn as_read(self) -> View<'a, E, C> {
320        unexpanded!()
321    }
322
323    pub fn __expand_as_read(scope: &Scope, this: ViewMutExpand<'a, E, C>) -> ViewExpand<'a, E, C> {
324        this.__expand_as_read_method(scope)
325    }
326}
327
328impl<'a, E: CubePrimitive, C: Coordinates + 'a> ViewMutExpand<'a, E, C> {
329    pub fn __expand_as_read_method(self, _scope: &Scope) -> ViewExpand<'a, E, C> {
330        let inner: &'a (dyn ViewOperationsExpand<E, C> + 'a) = self.inner;
331        ViewExpand { inner }
332    }
333}
334
335impl<'a, E: CubePrimitive, C: Coordinates> ViewExpand<'a, E, C> {
336    pub fn new<V: ViewOperationsExpand<E, C> + 'a>(scope: &Scope, view: V) -> Self {
337        let inner: &dyn ViewOperationsExpand<E, C> = scope.create_kernel_ref(view);
338        ViewExpand { inner }
339    }
340}
341
342impl<'a, E: CubePrimitive, C: Coordinates> ViewMutExpand<'a, E, C> {
343    pub fn new<V: ViewOperationsMutExpand<E, C> + 'a>(scope: &Scope, view: V) -> Self {
344        let inner: &mut dyn ViewOperationsMutExpand<E, C> = scope.create_kernel_ref(view);
345        ViewMutExpand { inner }
346    }
347}
348
349impl<'a, E: CubePrimitive, C: Coordinates + 'a> ViewMut<'a, E, C> {
350    pub fn view_mut<'b, T: Coordinates + 'a>(
351        self,
352        _layout: impl Layout<Coordinates = T, SourceCoordinates = C>,
353    ) -> ViewMut<'b, E, T>
354    where
355        'a: 'b,
356    {
357        unexpanded!()
358    }
359
360    pub fn __expand_view_mut<T: Coordinates + 'a>(
361        scope: &Scope,
362        this: ViewMutExpand<'a, E, C>,
363        layout: VirtualLayoutExpand<T, C>,
364    ) -> ViewMutExpand<'a, E, T> {
365        this.__expand_view_mut_method(scope, layout)
366    }
367}
368
369impl<'a, E: CubePrimitive, C: Coordinates + 'a> ViewMutExpand<'a, E, C> {
370    pub fn __expand_view_mut_method<'b, T: Coordinates + 'a>(
371        self,
372        scope: &Scope,
373        layout: VirtualLayoutExpand<T, C>,
374    ) -> ViewMutExpand<'b, E, T>
375    where
376        'a: 'b,
377    {
378        ViewMut::__expand_new::<ViewMut<'a, E, C>, C>(scope, self, layout)
379    }
380}
381
382#[cube]
383impl<'a, E: CubePrimitive, C: Coordinates> ViewMut<'a, E, C> {
384    /// Write a value to `pos`. The layout handles translation into a concrete index.
385    #[allow(unused_variables)]
386    pub fn write(&mut self, pos: C, value: E) {
387        intrinsic!(|scope| self.inner.__expand_write_method(scope, pos, value));
388    }
389
390    /// Write a value to `pos` if it's in bounds. The layout handles translation into a concrete index.
391    #[allow(unused_variables)]
392    pub fn write_checked(&mut self, pos: C, value: E) {
393        intrinsic!(|scope| self.inner.__expand_write_checked_method(scope, pos, value));
394    }
395
396    /// Interpret this view as a mutable linear slice encompassing the entire view.
397    ///
398    /// # Safety
399    ///
400    /// No checking is done on whether the slice is contiguous in memory.
401    pub fn as_linear_slice_mut(&mut self) -> &'a mut [E] {
402        intrinsic!(|scope| {
403            let shape = self.inner.__expand_shape_method(scope);
404            let origin = C::__expand_from_int(scope, shape.clone_unchecked(), 0);
405            // Inclusive end so clamping works correctly
406            let one = C::__expand_from_int(scope, shape.clone_unchecked(), 1);
407            let shape = C::__expand_max(scope, shape, one.clone_unchecked());
408            let end = C::__expand_sub(scope, shape, one);
409            let slice = self
410                .inner
411                .__expand_as_linear_slice_mut_method(scope, origin, end);
412            scope.create_kernel_ref(slice.expand.into())
413        })
414    }
415}
416
417impl<'a, E: CubePrimitive, C: Coordinates> ViewMutExpand<'a, E, C> {
418    pub(super) fn __expand_to_linear_slice_mut_inner_method(
419        &mut self,
420        scope: &Scope,
421        pos: C::ExpandType,
422        end: C::ExpandType,
423    ) -> &mut SliceExpand<E> {
424        self.inner
425            .__expand_as_linear_slice_mut_method(scope, pos, end)
426    }
427}
428
429#[cube]
430impl<'a, E: CubePrimitive, C: Coordinates + 'static> ViewMut<'a, E, C> {
431    /// Create a mutable slice starting from `pos`, with `size`.
432    /// The layout handles translation into concrete indices.
433    /// Size and pos will be clamped to the current layout size.
434    #[allow(unused_variables)]
435    pub fn slice_mut(self, pos: C, size: C) -> ViewMut<'a, E, C> {
436        intrinsic!(|scope| self.slice_mut(scope, pos, size, true))
437    }
438
439    /// Create a mutable slice starting from `pos`, with `size`.
440    /// The layout handles translation into concrete indices.
441    /// Size and pos will be clamped to the current layout size.
442    ///
443    /// # Safety
444    /// Access is always unchecked.
445    #[allow(unused_variables)]
446    pub fn slice_mut_unchecked(self, pos: C, size: C) -> ViewMut<'a, E, C> {
447        intrinsic!(|scope| self.slice_mut(scope, pos, size, false))
448    }
449}
450
451impl<'a, E: CubePrimitive, C: Coordinates + 'static> ViewMutExpand<'a, E, C> {
452    fn slice_mut(
453        &self,
454        scope: &Scope,
455        pos: C::ExpandType,
456        size: C::ExpandType,
457        checked: bool,
458    ) -> ViewMutExpand<'a, E, C> {
459        let shape = self.__expand_shape_method(scope);
460        let pos = C::__expand_min(scope, pos, shape.clone_unchecked());
461        let max_size = C::__expand_sub(scope, shape, pos.clone_unchecked());
462        let size = C::__expand_min(scope, size, max_size);
463        let layout = SliceLayout::__expand_new(scope, pos, size, checked);
464        self.clone().__expand_view_mut_method(scope, layout.into())
465    }
466}
467
468#[cube]
469impl<'a, E: CubePrimitive, C: Coordinates> ViewMut<'a, E, C> {
470    /// Execute a TMA store into global memory, if the underlying storage supports it.
471    /// Panics if it's unsupported.
472    #[allow(unused_variables)]
473    pub fn tensor_map_store(&self, shared_memory: &[E], pos: C) {
474        intrinsic!(|scope| {
475            self.inner
476                .__expand_tensor_map_store_method(scope, shared_memory, pos)
477        })
478    }
479}
cubecl_std/tensor/view/base.rs

cubecl_std/tensor/view/
base.rs
