cubecl_std/tensor/view/operations/
base.rs

1use crate::{CubeOption, tensor::layout::Coordinates};
2use cubecl::prelude::*;
3use cubecl_core::{self as cubecl, prelude::barrier::Barrier, unexpanded};
4
5/// Type from which we can read values in cube functions.
6/// For a mutable version, see [ListMut].
7#[allow(clippy::len_without_is_empty)]
8#[cube(self_type = "ref", expand_base_traits = "LinedExpand")]
9pub trait ViewOperations<T: CubePrimitive, C: Coordinates>: Lined {
10    #[allow(unused)]
11    fn read(&self, pos: C) -> T {
12        unexpanded!()
13    }
14
15    #[allow(unused)]
16    fn read_checked(&self, pos: C) -> T {
17        unexpanded!()
18    }
19
20    #[allow(unused)]
21    fn read_masked(&self, pos: C, value: T) -> T {
22        unexpanded!()
23    }
24
25    #[allow(unused)]
26    fn read_unchecked(&self, pos: C) -> T {
27        unexpanded!()
28    }
29
30    /// Create a slice starting from `pos`, with `size`.
31    /// The layout handles translation into concrete indices.
32    #[allow(unused)]
33    fn to_linear_slice(&self, pos: C, size: C) -> Slice<T, ReadOnly> {
34        unexpanded!()
35    }
36
37    fn as_tensor_map(&self) -> CubeOption<TensorMap<T>> {
38        unexpanded!()
39    }
40
41    ///.Execute a TMA load into shared memory, if the underlying storage supports it.
42    /// Panics if it's unsupported.
43    #[allow(unused)]
44    fn tensor_map_load(&self, barrier: &Barrier, shared_memory: &mut Slice<T, ReadWrite>, pos: C) {
45        unexpanded!()
46    }
47
48    #[allow(unused)]
49    fn shape(&self) -> C {
50        unexpanded!();
51    }
52
53    #[allow(unused)]
54    fn is_in_bounds(&self, pos: C) -> bool {
55        unexpanded!();
56    }
57}
58
59/// Type for which we can read and write values in cube functions.
60/// For an immutable version, see [List].
61#[cube(expand_base_traits = "ViewOperationsExpand<T, C>", self_type = "ref")]
62pub trait ViewOperationsMut<T: CubePrimitive, C: Coordinates>: ViewOperations<T, C> {
63    #[allow(unused)]
64    fn write(&self, pos: C, value: T) {
65        unexpanded!()
66    }
67
68    #[allow(unused)]
69    fn write_checked(&self, pos: C, value: T) {
70        unexpanded!()
71    }
72
73    /// Create a mutable slice starting from `pos`, with `size`.
74    /// The layout handles translation into concrete indices.
75    #[allow(unused, clippy::wrong_self_convention)]
76    fn to_linear_slice_mut(&self, pos: C, size: C) -> Slice<T, ReadWrite> {
77        unexpanded!()
78    }
79
80    ///.Execute a TMA store into global memory, if the underlying storage supports it.
81    /// Panics if it's unsupported.
82    #[allow(unused)]
83    fn tensor_map_store(&self, shared_memory: &Slice<T>, pos: C) {
84        unexpanded!()
85    }
86}
87
88// Automatic implementation for references to List.
89impl<'a, T: CubePrimitive, C: Coordinates, V: ViewOperations<T, C>> ViewOperations<T, C> for &'a V
90where
91    &'a V: CubeType<ExpandType = V::ExpandType>,
92{
93    fn read(&self, pos: C) -> T {
94        V::read(self, pos)
95    }
96
97    fn __expand_read(
98        scope: &mut Scope,
99        this: Self::ExpandType,
100        pos: C::ExpandType,
101    ) -> <T as CubeType>::ExpandType {
102        V::__expand_read(scope, this, pos)
103    }
104}
105
106// Automatic implementation for mutable references to List.
107impl<'a, T: CubePrimitive, C: Coordinates, L: ViewOperations<T, C>> ViewOperations<T, C>
108    for &'a mut L
109where
110    &'a mut L: CubeType<ExpandType = L::ExpandType>,
111{
112    fn read(&self, pos: C) -> T {
113        L::read(self, pos)
114    }
115
116    fn __expand_read(
117        scope: &mut Scope,
118        this: Self::ExpandType,
119        pos: C::ExpandType,
120    ) -> <T as CubeType>::ExpandType {
121        L::__expand_read(scope, this, pos)
122    }
123}
124
125// Automatic implementation for references to ListMut.
126impl<'a, T: CubePrimitive, C: Coordinates, L: ViewOperationsMut<T, C>> ViewOperationsMut<T, C>
127    for &'a L
128where
129    &'a L: CubeType<ExpandType = L::ExpandType>,
130{
131    fn write(&self, pos: C, value: T) {
132        L::write(self, pos, value);
133    }
134
135    fn __expand_write(
136        scope: &mut Scope,
137        this: Self::ExpandType,
138        pos: C::ExpandType,
139        value: T::ExpandType,
140    ) {
141        L::__expand_write(scope, this, pos, value);
142    }
143}
144
145// Automatic implementation for mutable references to ListMut.
146impl<'a, T: CubePrimitive, C: Coordinates, L: ViewOperationsMut<T, C>> ViewOperationsMut<T, C>
147    for &'a mut L
148where
149    &'a mut L: CubeType<ExpandType = L::ExpandType>,
150{
151    fn write(&self, pos: C, value: T) {
152        L::write(self, pos, value);
153    }
154
155    fn __expand_write(
156        scope: &mut Scope,
157        this: Self::ExpandType,
158        pos: C::ExpandType,
159        value: T::ExpandType,
160    ) {
161        L::__expand_write(scope, this, pos, value);
162    }
163}
cubecl_std/tensor/view/operations/base.rs

cubecl_std/tensor/view/operations/
base.rs
