cubecl_std/tensor/view/operations/
base.rs

1#![allow(clippy::mut_from_ref)]
2
3use crate::tensor::layout::Coordinates;
4use cubecl_core::{
5    self as cubecl,
6    frontend::{NativeExpand, barrier::Barrier},
7    prelude::*,
8    unexpanded,
9};
10
11/// Type from which we can read values in cube functions.
12/// For a mutable version, see [`ListMut`].
13#[allow(clippy::len_without_is_empty)]
14#[cube(expand_base_traits = "VectorizedExpand")]
15pub trait ViewOperations<T: CubePrimitive, C: Coordinates>: Vectorized {
16    #[allow(unused)]
17    fn read(&self, pos: C) -> T {
18        unexpanded!()
19    }
20
21    #[allow(unused)]
22    fn read_checked(&self, pos: C) -> T {
23        unexpanded!()
24    }
25
26    #[allow(unused)]
27    fn read_masked(&self, pos: C, value: T) -> T {
28        unexpanded!()
29    }
30
31    #[allow(unused)]
32    fn read_unchecked(&self, pos: C) -> T {
33        unexpanded!()
34    }
35
36    /// Create a slice starting from `pos`, with `size`.
37    /// The layout handles translation into concrete indices.
38    #[allow(unused)]
39    fn as_linear_slice(&self, pos: C, size: C) -> &[T] {
40        unexpanded!()
41    }
42
43    /// Execute a TMA load into shared memory, if the underlying storage supports it.
44    /// Panics if it's unsupported.
45    #[allow(unused)]
46    fn tensor_map_load(&self, barrier: &Barrier, shared_memory: &mut [T], pos: C) {
47        unexpanded!()
48    }
49
50    #[allow(unused)]
51    fn shape(&self) -> C {
52        unexpanded!();
53    }
54
55    #[allow(unused)]
56    fn is_in_bounds(&self, pos: C) -> bool {
57        unexpanded!();
58    }
59}
60
61/// Type for which we can read and write values in cube functions.
62/// For an immutable version, see [List].
63#[cube(expand_base_traits = "ViewOperationsExpand<T, C>")]
64pub trait ViewOperationsMut<T: CubePrimitive, C: Coordinates>: ViewOperations<T, C> {
65    #[allow(unused)]
66    fn write(&self, pos: C, value: T) {
67        unexpanded!()
68    }
69
70    #[allow(unused)]
71    fn write_checked(&self, pos: C, value: T) {
72        unexpanded!()
73    }
74
75    /// Create a mutable slice starting from `pos`, with `size`.
76    /// The layout handles translation into concrete indices.
77    #[allow(unused)]
78    fn as_linear_slice_mut(&self, pos: C, size: C) -> &mut [T] {
79        unexpanded!()
80    }
81
82    /// Execute a TMA store into global memory, if the underlying storage supports it.
83    /// Panics if it's unsupported.
84    #[allow(unused)]
85    fn tensor_map_store(&self, shared_memory: &[T], pos: C) {
86        unexpanded!()
87    }
88}
89
90macro_rules! view_ops_read_ref {
91    ($ty: ty) => {
92        impl<T: CubePrimitive, C: Coordinates, V: ViewOperations<T, C> + ?Sized>
93            ViewOperations<T, C> for $ty
94        {
95        }
96        impl<T: CubePrimitive, C: Coordinates, V: ViewOperationsExpand<T, C> + ?Sized>
97            ViewOperationsExpand<T, C> for $ty
98        {
99            #[allow(clippy::too_many_arguments)]
100            fn __expand_read_method(&self, scope: &Scope, pos: C::ExpandType) -> T::ExpandType {
101                (**self).__expand_read_method(scope, pos)
102            }
103
104            #[allow(clippy::too_many_arguments)]
105            fn __expand_read_checked_method(
106                &self,
107                scope: &Scope,
108                pos: C::ExpandType,
109            ) -> T::ExpandType {
110                (**self).__expand_read_checked_method(scope, pos)
111            }
112
113            #[allow(clippy::too_many_arguments)]
114            fn __expand_read_masked_method(
115                &self,
116                scope: &Scope,
117                pos: C::ExpandType,
118                value: T::ExpandType,
119            ) -> T::ExpandType {
120                (**self).__expand_read_masked_method(scope, pos, value)
121            }
122
123            #[allow(clippy::too_many_arguments)]
124            fn __expand_read_unchecked_method(
125                &self,
126                scope: &Scope,
127                pos: C::ExpandType,
128            ) -> T::ExpandType {
129                (**self).__expand_read_unchecked_method(scope, pos)
130            }
131
132            #[allow(clippy::too_many_arguments)]
133            fn __expand_as_linear_slice_method(
134                &self,
135                scope: &Scope,
136                pos: C::ExpandType,
137                size: C::ExpandType,
138            ) -> &SliceExpand<T> {
139                (**self).__expand_as_linear_slice_method(scope, pos, size)
140            }
141
142            #[allow(clippy::too_many_arguments)]
143            fn __expand_tensor_map_load_method(
144                &self,
145                scope: &Scope,
146                barrier: &NativeExpand<Barrier>,
147                shared_memory: &mut SliceExpand<T>,
148                pos: C::ExpandType,
149            ) -> () {
150                (**self).__expand_tensor_map_load_method(scope, barrier, shared_memory, pos);
151            }
152
153            #[allow(clippy::too_many_arguments)]
154            fn __expand_shape_method(&self, scope: &Scope) -> C::ExpandType {
155                (**self).__expand_shape_method(scope)
156            }
157
158            #[allow(clippy::too_many_arguments)]
159            fn __expand_is_in_bounds_method(
160                &self,
161                scope: &Scope,
162                pos: C::ExpandType,
163            ) -> NativeExpand<bool> {
164                (**self).__expand_is_in_bounds_method(scope, pos)
165            }
166        }
167    };
168}
169
170view_ops_read_ref!(&V);
171view_ops_read_ref!(&mut V);
172
173impl<T: CubePrimitive, C: Coordinates, V: ViewOperationsMut<T, C> + ?Sized> ViewOperationsMut<T, C>
174    for &mut V
175{
176}
177impl<T: CubePrimitive, C: Coordinates, V: ViewOperationsMutExpand<T, C> + ?Sized>
178    ViewOperationsMutExpand<T, C> for &mut V
179{
180    #[allow(clippy::too_many_arguments)]
181    fn __expand_write_method(&self, scope: &Scope, pos: C::ExpandType, value: T::ExpandType) {
182        (**self).__expand_write_method(scope, pos, value);
183    }
184
185    #[allow(clippy::too_many_arguments)]
186    fn __expand_write_checked_method(
187        &self,
188        scope: &Scope,
189        pos: C::ExpandType,
190        value: T::ExpandType,
191    ) {
192        (**self).__expand_write_checked_method(scope, pos, value);
193    }
194
195    #[allow(clippy::too_many_arguments)]
196    fn __expand_as_linear_slice_mut_method(
197        &self,
198        scope: &Scope,
199        pos: C::ExpandType,
200        size: C::ExpandType,
201    ) -> &mut SliceExpand<T> {
202        (**self).__expand_as_linear_slice_mut_method(scope, pos, size)
203    }
204
205    #[allow(clippy::too_many_arguments)]
206    fn __expand_tensor_map_store_method(
207        &self,
208        scope: &Scope,
209        shared_memory: &SliceExpand<T>,
210        pos: C::ExpandType,
211    ) {
212        (**self).__expand_tensor_map_store_method(scope, shared_memory, pos);
213    }
214}
cubecl_std/tensor/view/operations/base.rs

cubecl_std/tensor/view/operations/
base.rs
