cubecl_std/tensor/view/operations/
virtual_tensor.rs

1use super::*;
2use crate::tensor::{
3    layout::Coords1d,
4    r#virtual::{VirtualTensor, VirtualTensorExpand},
5};
6use cubecl::prelude::*;
7use cubecl_core::{self as cubecl, io::read_masked, prelude::barrier::Barrier};
8
9impl<T: Numeric, N: Size, IO: Clone> ViewOperations<Vector<T, N>, Coords1d>
10    for VirtualTensor<T, N, IO>
11{
12}
13impl<T: Numeric, N: Size, IO: Clone> ViewOperationsExpand<Vector<T, N>, Coords1d>
14    for VirtualTensorExpand<T, N, IO>
15{
16    fn __expand_read_method(
17        &self,
18        scope: &Scope,
19        pos: NativeExpand<usize>,
20    ) -> <Vector<T, N> as CubeType>::ExpandType {
21        Self::__expand_read_method(self, scope, pos)
22    }
23
24    fn __expand_read_checked_method(
25        &self,
26        scope: &Scope,
27        pos: NativeExpand<usize>,
28    ) -> <Vector<T, N> as CubeType>::ExpandType {
29        let zero = Vector::__expand_cast_from(scope, 0.into());
30        self.__expand_read_masked_method(scope, pos, zero)
31    }
32
33    fn __expand_read_masked_method(
34        &self,
35        scope: &Scope,
36        pos: NativeExpand<usize>,
37        mask_value: <Vector<T, N> as CubeType>::ExpandType,
38    ) -> <Vector<T, N> as CubeType>::ExpandType {
39        let in_bounds = self.__expand_is_in_bounds_method(scope, pos);
40        let slice = self.__expand_as_slice_method(scope);
41        read_masked::expand::<Vector<T, N>>(scope, in_bounds, slice, pos, mask_value)
42    }
43
44    fn __expand_read_unchecked_method(
45        &self,
46        scope: &Scope,
47        pos: NativeExpand<usize>,
48    ) -> <Vector<T, N> as CubeType>::ExpandType {
49        self.__expand_read_method(scope, pos)
50    }
51
52    fn __expand_as_linear_slice_method(
53        &self,
54        scope: &Scope,
55        pos: NativeExpand<usize>,
56        end: NativeExpand<usize>,
57    ) -> &SliceExpand<Vector<T, N>> {
58        // Convert to exclusive end
59        let end = end.__expand_add_method(scope, 1usize.into_expand(scope));
60        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
61        // negative length.
62        let start = clamp_max::expand(scope, pos, end);
63        <Self as SliceOperatorExpand<Vector<T, N>>>::__expand_slice_method(self, scope, start, end)
64    }
65
66    fn __expand_shape_method(&self, scope: &Scope) -> NativeExpand<usize> {
67        self.clone().__expand_buffer_len_method(scope)
68    }
69
70    fn __expand_is_in_bounds_method(
71        &self,
72        scope: &Scope,
73        pos: NativeExpand<usize>,
74    ) -> NativeExpand<bool> {
75        let len = self.clone().__expand_buffer_len_method(scope);
76        pos.__expand_lt_method(scope, &len)
77    }
78
79    fn __expand_tensor_map_load_method(
80        &self,
81        _scope: &Scope,
82        _barrier: &NativeExpand<Barrier>,
83        _shared_memory: &mut SliceExpand<Vector<T, N>>,
84        _pos: NativeExpand<usize>,
85    ) {
86        unimplemented!("Not a tensor map");
87    }
88}
89
90impl<T: Numeric, N: Size> ViewOperationsMut<Vector<T, N>, Coords1d>
91    for VirtualTensor<T, N, ReadWrite>
92{
93}
94impl<T: Numeric, N: Size> ViewOperationsMutExpand<Vector<T, N>, Coords1d>
95    for VirtualTensorExpand<T, N, ReadWrite>
96{
97    fn __expand_write_method(
98        &self,
99        scope: &Scope,
100        pos: NativeExpand<usize>,
101        value: <Vector<T, N> as CubeType>::ExpandType,
102    ) {
103        self.state_write().__expand_write_method(scope, pos, value)
104    }
105
106    fn __expand_write_checked_method(
107        &self,
108        scope: &Scope,
109        pos: NativeExpand<usize>,
110        value: <Vector<T, N> as CubeType>::ExpandType,
111    ) {
112        let len = self.clone().__expand_buffer_len_method(scope);
113        let in_bounds = pos.__expand_lt_method(scope, &len);
114        if_expand(scope, in_bounds, |scope| {
115            self.__expand_write_method(scope, pos, value)
116        })
117    }
118
119    fn __expand_as_linear_slice_mut_method(
120        &self,
121        _scope: &Scope,
122        _pos: NativeExpand<usize>,
123        _end: NativeExpand<usize>,
124    ) -> &mut SliceExpand<Vector<T, N>> {
125        todo!("VirtualTensor don't support slice mut yet");
126    }
127
128    fn __expand_tensor_map_store_method(
129        &self,
130        _scope: &Scope,
131        _shared_memory: &SliceExpand<Vector<T, N>>,
132        _pos: <Coords1d as CubeType>::ExpandType,
133    ) {
134        unimplemented!("Not a tensor map");
135    }
136}
cubecl_std/tensor/view/operations/virtual_tensor.rs

cubecl_std/tensor/view/operations/
virtual_tensor.rs
