t4a_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, and
83    /// return its completion: `shared_memory` may only be written again once it resolves.
84    /// Panics if it's unsupported.
85    #[allow(unused)]
86    fn tensor_map_store(&self, shared_memory: &[T], pos: C) -> Pending<()> {
87        unexpanded!()
88    }
89}
90
91macro_rules! view_ops_read_ref {
92    ($ty: ty) => {
93        impl<T: CubePrimitive, C: Coordinates, V: ViewOperations<T, C> + ?Sized>
94            ViewOperations<T, C> for $ty
95        {
96        }
97        impl<T: CubePrimitive, C: Coordinates, V: ViewOperationsExpand<T, C> + ?Sized>
98            ViewOperationsExpand<T, C> for $ty
99        {
100            #[allow(clippy::too_many_arguments)]
101            fn __expand_read_method(&self, scope: &Scope, pos: C::ExpandType) -> T::ExpandType {
102                (**self).__expand_read_method(scope, pos)
103            }
104
105            #[allow(clippy::too_many_arguments)]
106            fn __expand_read_checked_method(
107                &self,
108                scope: &Scope,
109                pos: C::ExpandType,
110            ) -> T::ExpandType {
111                (**self).__expand_read_checked_method(scope, pos)
112            }
113
114            #[allow(clippy::too_many_arguments)]
115            fn __expand_read_masked_method(
116                &self,
117                scope: &Scope,
118                pos: C::ExpandType,
119                value: T::ExpandType,
120            ) -> T::ExpandType {
121                (**self).__expand_read_masked_method(scope, pos, value)
122            }
123
124            #[allow(clippy::too_many_arguments)]
125            fn __expand_read_unchecked_method(
126                &self,
127                scope: &Scope,
128                pos: C::ExpandType,
129            ) -> T::ExpandType {
130                (**self).__expand_read_unchecked_method(scope, pos)
131            }
132
133            #[allow(clippy::too_many_arguments)]
134            fn __expand_as_linear_slice_method(
135                &self,
136                scope: &Scope,
137                pos: C::ExpandType,
138                size: C::ExpandType,
139            ) -> &SliceExpand<T> {
140                (**self).__expand_as_linear_slice_method(scope, pos, size)
141            }
142
143            #[allow(clippy::too_many_arguments)]
144            fn __expand_tensor_map_load_method(
145                &self,
146                scope: &Scope,
147                barrier: &NativeExpand<Barrier>,
148                shared_memory: &mut SliceExpand<T>,
149                pos: C::ExpandType,
150            ) -> () {
151                (**self).__expand_tensor_map_load_method(scope, barrier, shared_memory, pos);
152            }
153
154            #[allow(clippy::too_many_arguments)]
155            fn __expand_shape_method(&self, scope: &Scope) -> C::ExpandType {
156                (**self).__expand_shape_method(scope)
157            }
158
159            #[allow(clippy::too_many_arguments)]
160            fn __expand_is_in_bounds_method(
161                &self,
162                scope: &Scope,
163                pos: C::ExpandType,
164            ) -> NativeExpand<bool> {
165                (**self).__expand_is_in_bounds_method(scope, pos)
166            }
167        }
168    };
169}
170
171view_ops_read_ref!(&V);
172view_ops_read_ref!(&mut V);
173
174impl<T: CubePrimitive, C: Coordinates, V: ViewOperationsMut<T, C> + ?Sized> ViewOperationsMut<T, C>
175    for &mut V
176{
177}
178impl<T: CubePrimitive, C: Coordinates, V: ViewOperationsMutExpand<T, C> + ?Sized>
179    ViewOperationsMutExpand<T, C> for &mut V
180{
181    #[allow(clippy::too_many_arguments)]
182    fn __expand_write_method(&self, scope: &Scope, pos: C::ExpandType, value: T::ExpandType) {
183        (**self).__expand_write_method(scope, pos, value);
184    }
185
186    #[allow(clippy::too_many_arguments)]
187    fn __expand_write_checked_method(
188        &self,
189        scope: &Scope,
190        pos: C::ExpandType,
191        value: T::ExpandType,
192    ) {
193        (**self).__expand_write_checked_method(scope, pos, value);
194    }
195
196    #[allow(clippy::too_many_arguments)]
197    fn __expand_as_linear_slice_mut_method(
198        &self,
199        scope: &Scope,
200        pos: C::ExpandType,
201        size: C::ExpandType,
202    ) -> &mut SliceExpand<T> {
203        (**self).__expand_as_linear_slice_mut_method(scope, pos, size)
204    }
205
206    #[allow(clippy::too_many_arguments)]
207    fn __expand_tensor_map_store_method(
208        &self,
209        scope: &Scope,
210        shared_memory: &SliceExpand<T>,
211        pos: C::ExpandType,
212    ) -> PendingExpand<()> {
213        (**self).__expand_tensor_map_store_method(scope, shared_memory, pos)
214    }
215}
t4a_cubecl_std/tensor/view/operations/base.rs

t4a_cubecl_std/tensor/view/operations/
base.rs
