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::BarrierExpand};
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: &mut Scope,
19        pos: NativeExpand<usize>,
20    ) -> <Vector<T, N> as CubeType>::ExpandType {
21        <Self as ListExpand<Vector<T, N>>>::__expand_read_method(self, scope, pos)
22    }
23
24    fn __expand_read_checked_method(
25        &self,
26        scope: &mut 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: &mut 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.clone());
40        let slice = self.clone().__expand_to_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: &mut Scope,
47        pos: NativeExpand<usize>,
48    ) -> <Vector<T, N> as CubeType>::ExpandType {
49        <Self as ListExpand<Vector<T, N>>>::__expand_read_unchecked_method(self, scope, pos)
50    }
51
52    fn __expand_to_linear_slice_method(
53        &self,
54        scope: &mut Scope,
55        pos: NativeExpand<usize>,
56        end: NativeExpand<usize>,
57    ) -> SliceExpand<Vector<T, N>, ReadOnly> {
58        // Convert to exclusive end
59        let end = add::expand(scope, end, 1usize.into());
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.clone());
63        <Self as SliceOperatorExpand<Vector<T, N>>>::__expand_slice_method(self, scope, start, end)
64    }
65
66    fn __expand_shape_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
67        self.clone().__expand_buffer_len_method(scope)
68    }
69
70    fn __expand_is_in_bounds_method(
71        &self,
72        scope: &mut Scope,
73        pos: NativeExpand<usize>,
74    ) -> NativeExpand<bool> {
75        let len = self.clone().__expand_buffer_len_method(scope);
76        lt::expand(scope, pos, len)
77    }
78
79    fn __expand_tensor_map_load_method(
80        &self,
81        _scope: &mut Scope,
82        _barrier: BarrierExpand,
83        _shared_memory: SliceExpand<Vector<T, N>, ReadWrite>,
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: &mut Scope,
100        pos: NativeExpand<usize>,
101        value: <Vector<T, N> as CubeType>::ExpandType,
102    ) {
103        <Self as ListMutExpand<Vector<T, N>>>::__expand_write_method(self, scope, pos, value)
104    }
105
106    fn __expand_write_checked_method(
107        &self,
108        scope: &mut 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 = lt::expand(scope, pos.clone(), len);
114        if_expand(scope, in_bounds, |scope| {
115            <Self as ListMutExpand<Vector<T, N>>>::__expand_write_method(self, scope, pos, value)
116        })
117    }
118
119    fn __expand_to_linear_slice_mut_method(
120        &self,
121        scope: &mut Scope,
122        pos: NativeExpand<usize>,
123        end: NativeExpand<usize>,
124    ) -> SliceExpand<Vector<T, N>, ReadWrite> {
125        // Convert to exclusive end
126        let end = add::expand(scope, end, 1usize.into());
127        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
128        // negative length.
129        let start = clamp_max::expand(scope, pos, end.clone());
130        <Self as SliceMutOperatorExpand<Vector<T, N>>>::__expand_slice_mut_method(
131            self, scope, start, end,
132        )
133    }
134
135    fn __expand_tensor_map_store_method(
136        &self,
137        _scope: &mut Scope,
138        _shared_memory: SliceExpand<Vector<T, N>, ReadOnly>,
139        _pos: <Coords1d as CubeType>::ExpandType,
140    ) {
141        unimplemented!("Not a tensor map");
142    }
143}
cubecl_std/tensor/view/operations/virtual_tensor.rs

cubecl_std/tensor/view/operations/
virtual_tensor.rs
