cubecl_std/tensor/view/operations/
virtual_tensor.rs

1use super::*;
2use crate::{
3    CubeOption, CubeOptionExpand,
4    tensor::{
5        layout::Coords1d,
6        r#virtual::{VirtualTensor, VirtualTensorExpand},
7    },
8};
9use cubecl::prelude::*;
10use cubecl_core::{self as cubecl, io::read_masked, prelude::barrier::BarrierExpand};
11
12impl<T: Numeric, IO: Clone> ViewOperations<Line<T>, Coords1d> for VirtualTensor<T, IO> {}
13impl<T: Numeric, IO: Clone> ViewOperationsExpand<Line<T>, Coords1d> for VirtualTensorExpand<T, IO> {
14    fn __expand_read_method(
15        &self,
16        scope: &mut Scope,
17        pos: ExpandElementTyped<u32>,
18    ) -> <Line<T> as CubeType>::ExpandType {
19        <Self as ListExpand<Line<T>>>::__expand_read_method(self, scope, pos)
20    }
21
22    fn __expand_read_checked_method(
23        &self,
24        scope: &mut Scope,
25        pos: ExpandElementTyped<u32>,
26    ) -> <Line<T> as CubeType>::ExpandType {
27        let zero = Line::__expand_cast_from(scope, 0.into());
28        self.__expand_read_masked_method(scope, pos, zero)
29    }
30
31    fn __expand_read_masked_method(
32        &self,
33        scope: &mut Scope,
34        pos: ExpandElementTyped<u32>,
35        mask_value: <Line<T> as CubeType>::ExpandType,
36    ) -> <Line<T> as CubeType>::ExpandType {
37        let in_bounds = self.__expand_is_in_bounds_method(scope, pos.clone());
38        let slice = self.clone().__expand_to_slice_method(scope);
39        read_masked::expand::<Line<T>>(scope, in_bounds, slice, pos, mask_value)
40    }
41
42    fn __expand_read_unchecked_method(
43        &self,
44        scope: &mut Scope,
45        pos: ExpandElementTyped<u32>,
46    ) -> <Line<T> as CubeType>::ExpandType {
47        <Self as ListExpand<Line<T>>>::__expand_read_unchecked_method(self, scope, pos)
48    }
49
50    fn __expand_to_linear_slice_method(
51        &self,
52        scope: &mut Scope,
53        pos: ExpandElementTyped<u32>,
54        end: ExpandElementTyped<u32>,
55    ) -> SliceExpand<Line<T>, ReadOnly> {
56        // Convert to exclusive end
57        let end = add::expand(scope, end, 1u32.into());
58        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
59        // negative length.
60        let start = Min::__expand_min(scope, pos, end.clone());
61        <Self as SliceOperatorExpand<Line<T>>>::__expand_slice_method(self, scope, start, end)
62    }
63
64    fn __expand_as_tensor_map_method(
65        &self,
66        scope: &mut Scope,
67    ) -> CubeOptionExpand<TensorMap<Line<T>>> {
68        CubeOption::__expand_new_None(scope)
69    }
70
71    fn __expand_shape_method(&self, scope: &mut Scope) -> ExpandElementTyped<u32> {
72        self.clone().__expand_buffer_len_method(scope)
73    }
74
75    fn __expand_is_in_bounds_method(
76        &self,
77        scope: &mut Scope,
78        pos: ExpandElementTyped<u32>,
79    ) -> ExpandElementTyped<bool> {
80        let len = self.clone().__expand_buffer_len_method(scope);
81        lt::expand(scope, pos, len)
82    }
83
84    fn __expand_tensor_map_load_method(
85        &self,
86        _scope: &mut Scope,
87        _barrier: BarrierExpand,
88        _shared_memory: SliceExpand<Line<T>, ReadWrite>,
89        _pos: ExpandElementTyped<u32>,
90    ) {
91        unimplemented!("Not a tensor map");
92    }
93}
94
95impl<T: Numeric> ViewOperationsMut<Line<T>, Coords1d> for VirtualTensor<T, ReadWrite> {}
96impl<T: Numeric> ViewOperationsMutExpand<Line<T>, Coords1d> for VirtualTensorExpand<T, ReadWrite> {
97    fn __expand_write_method(
98        &self,
99        scope: &mut Scope,
100        pos: ExpandElementTyped<u32>,
101        value: <Line<T> as CubeType>::ExpandType,
102    ) {
103        <Self as ListMutExpand<Line<T>>>::__expand_write_method(self, scope, pos, value)
104    }
105
106    fn __expand_write_checked_method(
107        &self,
108        scope: &mut Scope,
109        pos: ExpandElementTyped<u32>,
110        value: <Line<T> 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.into(), |scope| {
115            <Self as ListMutExpand<Line<T>>>::__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: ExpandElementTyped<u32>,
123        end: ExpandElementTyped<u32>,
124    ) -> SliceExpand<Line<T>, ReadWrite> {
125        // Convert to exclusive end
126        let end = add::expand(scope, end, 1u32.into());
127        // Handling for shapes that are 0 in at least one dim, ensures the slice is not
128        // negative length.
129        let start = Min::__expand_min(scope, pos, end.clone());
130        <Self as SliceMutOperatorExpand<Line<T>>>::__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<Line<T>, 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
