cubecl_std/tensor/
virtual.rs

1use alloc::sync::Arc;
2use core::{cell::UnsafeCell, marker::PhantomData};
3use cubecl::prelude::*;
4use cubecl_core::{self as cubecl, unexpanded};
5use std::ops::{Deref, DerefMut};
6
7use crate::tensor::{
8    ViewExpand, ViewMut, ViewMutExpand,
9    layout::{
10        Coordinates, Coords1d, Layout, VirtualLayout, VirtualLayoutExpand, simple::SimpleLayout,
11    },
12    view::View,
13};
14
15/// Tensor representation that is decoupled from how the tensor is stored.
16#[derive(Clone)]
17pub struct VirtualTensor<E: Numeric, N: Size, IO = ReadOnly> {
18    // state: Arc<dyn VirtualTensorOperations<E>>,
19    _e: PhantomData<E>,
20    _n: PhantomData<N>,
21    _p: PhantomData<IO>,
22}
23
24/// Expand type for [`VirtualTensor`].
25#[derive(Clone)]
26pub struct VirtualTensorExpand<E: Numeric, N: Size, IO> {
27    state: Arc<UnsafeCell<dyn VirtualTensorOperationsExpand<E, N>>>,
28    _p: PhantomData<IO>,
29}
30
31impl<E: Numeric, N: Size, IO> VirtualTensorExpand<E, N, IO> {
32    fn state_read(&self) -> &dyn VirtualTensorOperationsExpand<E, N> {
33        // SAFETY: State is valid for the entire lifetime of `self`
34        unsafe { &mut *self.state.get() }
35    }
36
37    #[allow(clippy::mut_from_ref)]
38    pub(crate) fn state_write(&self) -> &mut dyn VirtualTensorOperationsExpand<E, N> {
39        // SAFETY: State is a handle into memory and only the memory is written. The state itself
40        // is never mutated, so this is safe.
41        unsafe { &mut *self.state.get() }
42    }
43}
44
45#[cube]
46impl<E: Numeric, N: Size, IO: Clone> VirtualTensor<E, N, IO> {
47    #[allow(unused)]
48    pub fn read(&self, index: usize) -> Vector<E, N> {
49        intrinsic!(|scope| { self.state_read().__expand_read_method(scope, index) })
50    }
51
52    #[allow(unused, clippy::len_without_is_empty)]
53    pub fn len(&self) -> usize {
54        intrinsic!(|scope| { self.state_read().__expand_len_method(scope) })
55    }
56}
57
58impl<T: Numeric, N: Size, IO: Clone> Deref for VirtualTensor<T, N, IO> {
59    type Target = [Vector<T, N>];
60
61    fn deref(&self) -> &Self::Target {
62        unexpanded!()
63    }
64}
65
66impl<T: Numeric, N: Size> DerefMut for VirtualTensor<T, N, ReadWrite> {
67    fn deref_mut(&mut self) -> &mut Self::Target {
68        unexpanded!()
69    }
70}
71
72impl<E: Numeric, N: Size, IO: Clone> Vectorized for VirtualTensor<E, N, IO> {}
73impl<E: Numeric, N: Size, IO: Clone> VectorizedExpand for VirtualTensorExpand<E, N, IO> {
74    fn vector_size(&self) -> VectorSize {
75        self.state_read().vector_size()
76    }
77}
78
79impl<E: Numeric, N: Size, IO: Clone> SliceOperator<Vector<E, N>> for VirtualTensor<E, N, IO> {}
80impl<E: Numeric, N: Size, IO: Clone> SliceOperatorExpand<Vector<E, N>>
81    for VirtualTensorExpand<E, N, IO>
82{
83    fn __expand_slice_method(
84        &self,
85        scope: &Scope,
86        start: NativeExpand<usize>,
87        end: NativeExpand<usize>,
88    ) -> &SliceExpand<Vector<E, N>> {
89        self.state_read()
90            .__expand_read_window_method(scope, start, end)
91    }
92
93    fn __expand_slice_mut_method(
94        &mut self,
95        _: &Scope,
96        _: NativeExpand<usize>,
97        _: NativeExpand<usize>,
98    ) -> &mut SliceExpand<Vector<E, N>> {
99        todo!("VirtualTensor don't support slice mut yet");
100    }
101}
102
103#[cube]
104impl<E: Numeric, N: Size, IO: Clone> VirtualTensor<E, N, IO> {
105    pub fn as_slice(&self) -> &[Vector<E, N>] {
106        self.slice(0, self.len())
107    }
108
109    pub fn as_mut_slice(&mut self) -> &mut [Vector<E, N>] {
110        self.slice_mut(0, self.len())
111    }
112
113    pub fn as_tensor_map(&self) -> ComptimeOption<TensorMap<E, Tiled>> {
114        intrinsic!(|scope| self.state_read().__expand_as_tensor_map_method(scope))
115    }
116    #[allow(unused)]
117    pub fn slice(&self, start: usize, end: usize) -> &[Vector<E, N>] {
118        intrinsic!(|scope| self
119            .state_read()
120            .__expand_read_window_method(scope, start, end))
121    }
122    /// Get the shape of the tensor at the given axis.
123    #[allow(unused)]
124    pub fn shape(&self, axis: usize) -> usize {
125        intrinsic!(|scope| { self.state_read().__expand_shape_method(scope, axis) })
126    }
127    /// Get the stride of the tensor at the given axis.
128    #[allow(unused)]
129    pub fn stride(&self, axis: usize) -> usize {
130        intrinsic!(|scope| self.state_read().__expand_stride_method(scope, axis))
131    }
132    /// Get the rank of the tensor.
133    pub fn rank(&self) -> usize {
134        intrinsic!(|scope| self.state_read().__expand_rank_method(scope))
135    }
136
137    pub fn buffer_len(&self) -> usize {
138        intrinsic!(|scope| self.state_read().__expand_buffer_len_method(scope))
139    }
140}
141
142impl<E: Numeric, N: Size, IO: Clone + 'static> VirtualTensor<E, N, IO> {
143    /// Create a conceptual view over this tensor, allowing for multi-dimensional indexing with custom
144    /// layouts
145    pub fn view<C: Coordinates + 'static>(
146        &self,
147        layout: impl Into<VirtualLayout<C, Coords1d>>,
148    ) -> View<'_, Vector<E, N>, C> {
149        View::new::<&VirtualTensor<E, N, IO>, Coords1d>(self, layout)
150    }
151
152    /// Create a conceptual view over this tensor, allowing for multi-dimensional indexing with custom
153    /// layouts
154    pub fn into_view<C: Coordinates + 'static>(
155        self,
156        layout: impl Into<VirtualLayout<C, Coords1d>>,
157    ) -> View<'static, Vector<E, N>, C> {
158        View::new::<VirtualTensor<E, N, IO>, Coords1d>(self, layout)
159    }
160}
161
162#[cube]
163impl<E: Numeric, N: Size, IO: Clone + 'static> VirtualTensor<E, N, IO> {
164    /// Create a conceptual view over this tensor, with a simple linear layout
165    pub fn as_view(&self) -> View<'_, Vector<E, N>, usize> {
166        let vector_size = self.vector_size();
167        View::new::<&VirtualTensor<E, N, IO>, usize>(
168            self,
169            SimpleLayout::new(self.len() * vector_size, vector_size),
170        )
171    }
172}
173
174impl<E: Numeric, N: Size, IO: Clone + 'static> VirtualTensorExpand<E, N, IO> {
175    /// Create a conceptual view over this tensor, allowing for multi-dimensional indexing with custom
176    /// layouts
177    pub fn __expand_view_method<C: Coordinates + 'static>(
178        &self,
179        scope: &Scope,
180        layout: VirtualLayoutExpand<C, Coords1d>,
181    ) -> ViewExpand<'_, Vector<E, N>, C> {
182        View::__expand_new::<&VirtualTensor<E, N, IO>, Coords1d>(scope, self, layout)
183    }
184
185    pub fn __expand_into_view_method<C: Coordinates + 'static>(
186        self,
187        scope: &Scope,
188        layout: VirtualLayoutExpand<C, Coords1d>,
189    ) -> ViewExpand<'static, Vector<E, N>, C> {
190        View::__expand_new::<VirtualTensor<E, N, IO>, Coords1d>(scope, self, layout)
191    }
192}
193
194impl<E: Numeric, N: Size> VirtualTensor<E, N, ReadWrite> {
195    /// Create a mutable conceptual view over this tensor, allowing for multi-dimensional indexing
196    /// with custom layouts
197    pub fn view_mut<C: Coordinates + 'static>(
198        &mut self,
199        layout: impl Layout<Coordinates = C, SourceCoordinates = Coords1d> + 'static,
200    ) -> ViewMut<'_, Vector<E, N>, C> {
201        ViewMut::new::<&mut VirtualTensor<E, N, ReadWrite>, Coords1d>(self, layout)
202    }
203    pub fn __expand_view_mut<'a, C: Coordinates + 'static>(
204        scope: &Scope,
205        this: &'a mut VirtualTensorExpand<E, N, ReadWrite>,
206        layout: VirtualLayoutExpand<C, Coords1d>,
207    ) -> ViewMutExpand<'a, Vector<E, N>, C> {
208        this.__expand_view_mut_method::<C>(scope, layout)
209    }
210
211    pub fn into_view_mut<C: Coordinates + 'static>(
212        self,
213        layout: impl Layout<Coordinates = C, SourceCoordinates = Coords1d> + 'static,
214    ) -> ViewMut<'static, Vector<E, N>, C> {
215        ViewMut::new::<VirtualTensor<E, N, ReadWrite>, Coords1d>(self, layout)
216    }
217}
218
219impl<E: Numeric, N: Size> VirtualTensorExpand<E, N, ReadWrite> {
220    pub fn __expand_view_mut_method<C: Coordinates + 'static>(
221        &mut self,
222        scope: &Scope,
223        layout: VirtualLayoutExpand<C, Coords1d>,
224    ) -> ViewMutExpand<'_, Vector<E, N>, C> {
225        ViewMut::__expand_new::<&mut VirtualTensor<E, N, ReadWrite>, Coords1d>(scope, self, layout)
226    }
227
228    pub fn __expand_into_view_mut_method<C: Coordinates + 'static>(
229        self,
230        scope: &Scope,
231        layout: VirtualLayoutExpand<C, Coords1d>,
232    ) -> ViewMutExpand<'static, Vector<E, N>, C> {
233        ViewMut::__expand_new::<VirtualTensor<E, N, ReadWrite>, Coords1d>(scope, self, layout)
234    }
235}
236
237#[cube]
238impl<E: Numeric, N: Size> VirtualTensor<E, N, ReadWrite> {
239    /// Create a conceptual mutable view over this tensor, with a simple linear layout
240    pub fn as_view_mut(&mut self) -> ViewMut<'_, Vector<E, N>, usize> {
241        let vector_size = self.vector_size();
242        let layout = SimpleLayout::new(self.len() * vector_size, vector_size);
243        ViewMut::new::<&mut VirtualTensor<E, N, ReadWrite>, usize>(self, layout)
244    }
245}
246
247#[cube]
248impl<E: Numeric, N: Size, IO: Clone + 'static> VirtualTensor<E, N, IO> {
249    pub fn coordinate(&self, index: usize, dim: usize) -> usize {
250        let num_strides = index / self.stride(dim);
251        num_strides % self.shape(dim)
252    }
253}
254
255#[cube]
256impl<E: Numeric, N: Size> VirtualTensor<E, N, ReadWrite> {
257    #[allow(unused)]
258    pub fn write(&mut self, index: usize, value: Vector<E, N>) {
259        intrinsic!(|scope| self
260            .state_write()
261            .__expand_write_method(scope, index, value))
262    }
263}
264
265impl<E: Numeric, N: Size> VirtualTensor<E, N, ReadOnly> {
266    /// Create a new [read only](ReadOnly) [virtual tensor](VirtualTensor).
267    pub fn new<V: VirtualTensorOperations<E, N> + 'static>(_v: V) -> Self {
268        unexpanded!()
269    }
270
271    /// Expand function of [`Self::new`].
272    pub fn __expand_new<V: VirtualTensorOperations<E, N> + 'static>(
273        _scope: &Scope,
274        v: V::ExpandType,
275    ) -> VirtualTensorExpand<E, N, ReadOnly> {
276        VirtualTensorExpand {
277            state: Arc::new(UnsafeCell::new(v)),
278            _p: PhantomData,
279        }
280    }
281}
282
283impl<E: Numeric, N: Size> VirtualTensor<E, N, ReadWrite> {
284    /// Create a new [read write](ReadWrite) [virtual tensor](VirtualTensor).
285    pub fn new<V: VirtualTensorOperations<E, N> + 'static>(_v: V) -> Self {
286        unexpanded!()
287    }
288
289    /// Expand function of [`Self::new`].
290    pub fn __expand_new<V: VirtualTensorOperations<E, N> + 'static>(
291        _scope: &Scope,
292        v: V::ExpandType,
293    ) -> VirtualTensorExpand<E, N, ReadWrite> {
294        VirtualTensorExpand {
295            state: Arc::new(UnsafeCell::new(v)),
296            _p: PhantomData,
297        }
298    }
299}
300
301/// Trait to be implemented by a type that can become a [virtual tensor](VirtualTensor).
302///
303/// The [expand trait](VirtualTensorOperationsExpand) also need to be implemented for the type's
304/// expand type.
305///
306/// # Warning
307///
308/// This trait is kind of unsafe, [`VirtualTensorOperations::write`] doesn't follow the mutability
309/// rules, but it won't lead to any undefined behavior.
310#[cube(expand_base_traits = "VectorizedExpand")]
311pub trait VirtualTensorOperations<E: Numeric, N: Size>: Vectorized {
312    fn as_tensor_map(&self) -> ComptimeOption<TensorMap<E, Tiled>> {
313        unexpanded!()
314    }
315    /// Read the tensor at the given index.
316    fn read(&self, _index: usize) -> Vector<E, N> {
317        unexpanded!()
318    }
319    fn read_window(&self, _start: usize, _end: usize) -> &[Vector<E, N>] {
320        unexpanded!()
321    }
322    /// Write the tensor at the given index.
323    #[allow(unused)]
324    fn write(&mut self, _index: usize, value: Vector<E, N>) {
325        unexpanded!()
326    }
327    /// Get the shape of the tensor at the given axis.
328    fn shape(&self, _axis: usize) -> usize {
329        unexpanded!()
330    }
331    /// Get the stride of the tensor at the given axis.
332    fn stride(&self, _axis: usize) -> usize {
333        unexpanded!()
334    }
335    /// Get the rank of the tensor.
336    fn rank(&self) -> usize {
337        unexpanded!()
338    }
339    fn len(&self) -> usize {
340        unexpanded!()
341    }
342    fn buffer_len(&self) -> usize {
343        unexpanded!()
344    }
345}
346
347/// Making [virtual tensors](VirtualTensor) a proper [cube type](CubeType).
348mod __cube_type {
349    use super::*;
350
351    impl<E: Numeric, N: Size, IO: Clone> CubeType for VirtualTensor<E, N, IO> {
352        type ExpandType = VirtualTensorExpand<E, N, IO>;
353    }
354
355    impl<E: Numeric, N: Size, IO: Clone> IntoExpand for VirtualTensorExpand<E, N, IO> {
356        type Expand = VirtualTensorExpand<E, N, IO>;
357
358        fn into_expand(self, _: &Scope) -> Self::Expand {
359            self
360        }
361    }
362
363    impl<E: Numeric, N: Size, IO: Clone> ExpandTypeClone for VirtualTensorExpand<E, N, IO> {
364        fn clone_unchecked(&self) -> Self {
365            self.clone()
366        }
367    }
368
369    impl<E: Numeric, N: Size, IO> IntoMut for VirtualTensorExpand<E, N, IO> {
370        fn into_mut(self, _scope: &Scope) -> Self {
371            self
372        }
373    }
374
375    impl<E: Numeric, N: Size, IO> CubeDebug for VirtualTensorExpand<E, N, IO> {}
376
377    impl<E: Numeric, N: Size, IO: Clone> AsRefExpand for VirtualTensorExpand<E, N, IO> {
378        fn __expand_ref_method(&self, _: &Scope) -> &Self {
379            self
380        }
381    }
382
383    impl<E: Numeric, N: Size, IO: Clone> AsMutExpand for VirtualTensorExpand<E, N, IO> {
384        fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
385            self
386        }
387    }
388}
389
390/// Enable tensors to be virtual.
391mod __tensor {
392    use super::*;
393
394    impl<E: Numeric, N: Size> VirtualTensorOperations<E, N> for Tensor<Vector<E, N>> {}
395    impl<E: Numeric, N: Size> VirtualTensorOperationsExpand<E, N> for TensorExpand<Vector<E, N>> {
396        fn __expand_read_method(
397            &self,
398            scope: &Scope,
399            index: NativeExpand<usize>,
400        ) -> NativeExpand<Vector<E, N>> {
401            unsafe {
402                self.__expand_get_unchecked_method(scope, index)
403                    .__expand_deref_method(scope)
404            }
405        }
406        fn __expand_read_window_method(
407            &self,
408            context: &Scope,
409            start: NativeExpand<usize>,
410            end: NativeExpand<usize>,
411        ) -> &'static SliceExpand<Vector<E, N>> {
412            unsafe { core::mem::transmute(self.__expand_slice_method(context, start, end)) }
413        }
414
415        fn __expand_write_method(
416            &mut self,
417            scope: &Scope,
418            index: NativeExpand<usize>,
419            value: NativeExpand<Vector<E, N>>,
420        ) {
421            unsafe {
422                self.__expand_get_unchecked_mut_method(scope, index)
423                    .__expand_assign_method(scope, value)
424            };
425        }
426
427        fn __expand_shape_method(
428            &self,
429            scope: &Scope,
430            axis: NativeExpand<usize>,
431        ) -> NativeExpand<usize> {
432            self.__expand_shape_method(scope, axis)
433        }
434
435        fn __expand_stride_method(
436            &self,
437            scope: &Scope,
438            axis: NativeExpand<usize>,
439        ) -> NativeExpand<usize> {
440            self.__expand_stride_method(scope, axis)
441        }
442
443        fn __expand_rank_method(&self, scope: &Scope) -> NativeExpand<usize> {
444            self.__expand_rank_method(scope)
445        }
446        fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
447            self.__expand_len_method(scope)
448        }
449        fn __expand_buffer_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
450            self.__expand_buffer_len_method(scope)
451        }
452
453        fn __expand_as_tensor_map_method(
454            &self,
455            scope: &Scope,
456        ) -> ComptimeOptionExpand<TensorMap<E, Tiled>> {
457            ComptimeOption::__expand_new_None(scope)
458        }
459    }
460}
461
462/// Enable tensor maps to be virtual.
463mod __tensor_map {
464    use super::*;
465
466    impl<E: Numeric, N: Size> VirtualTensorOperations<E, N> for TensorMap<E, Tiled> {}
467    impl<E: Numeric, N: Size> VirtualTensorOperationsExpand<E, N>
468        for NativeExpand<TensorMap<E, Tiled>>
469    {
470        fn __expand_read_method(
471            &self,
472            _scope: &Scope,
473            _index: NativeExpand<usize>,
474        ) -> NativeExpand<Vector<E, N>> {
475            todo!()
476        }
477        fn __expand_read_window_method(
478            &self,
479            _context: &Scope,
480            _start: NativeExpand<usize>,
481            _end: NativeExpand<usize>,
482        ) -> &'static SliceExpand<Vector<E, N>> {
483            todo!()
484        }
485
486        fn __expand_write_method(
487            &mut self,
488            _scope: &Scope,
489            _index: NativeExpand<usize>,
490            _value: NativeExpand<Vector<E, N>>,
491        ) {
492            todo!()
493        }
494
495        fn __expand_shape_method(
496            &self,
497            _scope: &Scope,
498            _axis: NativeExpand<usize>,
499        ) -> NativeExpand<usize> {
500            todo!()
501        }
502
503        fn __expand_stride_method(
504            &self,
505            _scope: &Scope,
506            _axis: NativeExpand<usize>,
507        ) -> NativeExpand<usize> {
508            todo!()
509        }
510
511        fn __expand_rank_method(&self, _scope: &Scope) -> NativeExpand<usize> {
512            todo!()
513        }
514        fn __expand_len_method(&self, _scope: &Scope) -> NativeExpand<usize> {
515            todo!()
516        }
517        fn __expand_buffer_len_method(&self, _scope: &Scope) -> NativeExpand<usize> {
518            todo!()
519        }
520
521        fn __expand_as_tensor_map_method(
522            &self,
523            scope: &Scope,
524        ) -> ComptimeOptionExpand<TensorMap<E, Tiled>> {
525            ComptimeOption::__expand_new_Some(scope, *self)
526        }
527    }
528}
cubecl_std/tensor/virtual.rs

cubecl_std/tensor/
virtual.rs
