cubecl_std/tensor/view/operations/
base.rs

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

cubecl_std/tensor/view/operations/
base.rs
