cubecl_core/frontend/
tensor_layout.rs

1use alloc::boxed::Box;
2
3use cubecl_ir::{
4    dialect::spirv::{CreateLayoutOp, CreateViewOp, SliceOp},
5    pliron::{builtin::op_interfaces::OneResultInterface, r#type::Typed},
6    types::spirv::{ClampMode, TensorLayoutType, TensorViewType},
7};
8use pliron::r#type::TypeHandle;
9
10use crate::{self as cubecl, unexpanded};
11
12use crate::prelude::*;
13
14#[derive_cube_comptime]
15pub enum TensorClampMode {
16    Undefined,
17    Constant(u32),
18    ClampToEdge,
19    Repeat,
20    RepeatMirrored,
21}
22
23impl From<TensorClampMode> for ClampMode {
24    fn from(value: TensorClampMode) -> Self {
25        match value {
26            TensorClampMode::Undefined => ClampMode::Undefined,
27            TensorClampMode::Constant(val) => ClampMode::Constant(val),
28            TensorClampMode::ClampToEdge => ClampMode::ClampToEdge,
29            TensorClampMode::Repeat => ClampMode::Repeat,
30            TensorClampMode::RepeatMirrored => ClampMode::RepeatMirrored,
31        }
32    }
33}
34
35// OpTypeTensorLayoutNV with optional OpTypeTensorViewNV
36#[derive(CubeType, Clone)]
37pub struct TensorView<T: CubePrimitive> {
38    #[allow(unused)]
39    pub(crate) buffer: Box<[T]>,
40    #[allow(unused)]
41    pub(crate) layout: TensorLayout,
42    #[allow(unused)]
43    pub(crate) view: ComptimeOption<TensorReinterpret>,
44}
45
46#[derive_cube_comptime]
47pub struct TensorLayout;
48
49#[derive_cube_comptime]
50pub struct TensorReinterpret;
51
52impl CubeType for TensorLayout {
53    type ExpandType = NativeExpand<TensorLayout>;
54}
55
56impl CubeDebug for TensorLayout {}
57impl CubePrimitive for TensorLayout {
58    type Scalar = u32;
59    type Size = Const<1>;
60    type WithScalar<S: Scalar> = S;
61
62    fn from_const_value(_: cubecl_ir::ConstantValue) -> Self {
63        panic!("Can't construct tensor layout from constant")
64    }
65}
66impl ConstantTyped for TensorLayout {
67    fn __expand_as_type(_scope: &Scope) -> TypeHandle {
68        unimplemented!()
69    }
70}
71
72impl NativeAssign for TensorLayout {}
73
74impl CubeType for TensorReinterpret {
75    type ExpandType = NativeExpand<TensorReinterpret>;
76}
77
78impl CubeDebug for TensorReinterpret {}
79impl CubePrimitive for TensorReinterpret {
80    type Scalar = u32;
81    type Size = Const<1>;
82    type WithScalar<S: Scalar> = S;
83
84    fn from_const_value(_: cubecl_ir::ConstantValue) -> Self {
85        panic!("Can't construct tensor layout from constant")
86    }
87}
88impl ConstantTyped for TensorReinterpret {
89    fn __expand_as_type(_scope: &Scope) -> TypeHandle {
90        unimplemented!()
91    }
92}
93
94impl NativeAssign for TensorReinterpret {}
95
96#[derive(CubeType, CubeLaunch)]
97pub struct TensorViewBuilder<T: CubePrimitive> {
98    #[allow(unused)]
99    buffer: Box<[T]>,
100    #[allow(unused)]
101    shape: Sequence<u32>,
102    /// Strides default to contiguous strides
103    strides: ComptimeOption<Sequence<u32>>,
104    #[cube(comptime)]
105    clamp_mode: TensorClampMode,
106}
107
108#[cube]
109impl<T: CubePrimitive> TensorView<T> {
110    #[allow(clippy::new_ret_no_self)]
111    pub fn new(buffer: &[T], shape: Sequence<u32>) -> TensorViewBuilder<T> {
112        TensorViewBuilder::<T> {
113            buffer: unsafe { buffer.as_boxed_unchecked() },
114            shape,
115            strides: ComptimeOption::new_None(),
116            clamp_mode: comptime![TensorClampMode::Constant(0)],
117        }
118    }
119
120    #[allow(unused)]
121    pub fn slice(&self, offs: Sequence<u32>, shape: Sequence<u32>) -> TensorView<T> {
122        intrinsic!(|scope| {
123            assert_eq!(
124                offs.len(),
125                self.layout.rank(scope),
126                "Offsets and view rank must match"
127            );
128            assert_eq!(
129                offs.len(),
130                shape.len(),
131                "Offsets and shape must have same rank"
132            );
133            let layout = self.layout.read_value(scope);
134            let offs = offs.iter_cloned().map(|it| it.read_value(scope)).collect();
135            let shape = shape.iter_cloned().map(|it| it.read_value(scope)).collect();
136            let slice_op = SliceOp::new(scope.ctx_mut(), layout, offs, shape);
137            scope.register(&slice_op);
138            let new_layout = slice_op.get_result(scope.ctx());
139            TensorViewExpand {
140                buffer: self.buffer.clone(),
141                layout: new_layout.into(),
142                view: self.view.clone(),
143            }
144        })
145    }
146}
147
148impl NativeExpand<TensorLayout> {
149    fn rank(&self, scope: &Scope) -> usize {
150        let ty = self.read_value(scope).get_type(scope.ctx());
151        let ctx = scope.ctx();
152        let ty = ty.deref(ctx);
153        let TensorLayoutType { rank, .. } = ty.downcast_ref().unwrap();
154        *rank
155    }
156}
157
158impl<T: CubePrimitive> TensorView<T> {
159    pub fn permuted(&self, _permutation: Sequence<usize>) -> TensorView<T> {
160        unexpanded!()
161    }
162}
163
164impl<T: CubePrimitive> TensorViewExpand<T> {
165    pub fn __expand_permuted_method(
166        self,
167        scope: &Scope,
168        permutation: SequenceExpand<usize>,
169    ) -> TensorViewExpand<T> {
170        let dims = permutation.len();
171        assert!(dims <= 5, "Max 5 dims allowed");
172        let permutation = permutation
173            .iter_cloned()
174            .map(|it| {
175                it.constant()
176                    .expect("permutation must be constant")
177                    .as_usize()
178            })
179            .collect::<alloc::vec::Vec<_>>();
180        let ty = TensorViewType::get(scope.ctx(), permutation.len(), false, permutation);
181        let op = CreateViewOp::new(scope.ctx_mut(), ty.into());
182        scope.register(&op);
183        let view = op.get_result(scope.ctx());
184
185        TensorViewExpand {
186            buffer: self.buffer,
187            layout: self.layout,
188            view: ComptimeOptionExpand::Some(view.into()),
189        }
190    }
191}
192
193impl<T: CubePrimitive> TensorViewBuilder<T> {
194    pub fn with_strides(mut self, strides: Sequence<u32>) -> Self {
195        self.strides = ComptimeOption::Some(strides);
196        self
197    }
198
199    pub fn with_clamp_mode(mut self, clamp_mode: TensorClampMode) -> Self {
200        self.clamp_mode = clamp_mode;
201        self
202    }
203
204    pub fn finish(self) -> TensorView<T> {
205        unexpanded!()
206    }
207}
208
209impl<T: CubePrimitive> TensorViewBuilderExpand<T> {
210    pub fn __expand_with_strides_method(
211        mut self,
212        _scope: &Scope,
213        strides: SequenceExpand<u32>,
214    ) -> Self {
215        self.strides = ComptimeOptionExpand::Some(strides);
216        self
217    }
218
219    pub fn __expand_with_clamp_mode_method(
220        mut self,
221        _scope: &Scope,
222        clamp_mode: TensorClampMode,
223    ) -> Self {
224        self.clamp_mode = clamp_mode;
225        self
226    }
227
228    pub fn __expand_finish_method(self, scope: &Scope) -> TensorViewExpand<T> {
229        let shape = self.shape.into_iter().map(|it| it.read_value(scope));
230        let strides = match self.strides {
231            ComptimeOptionExpand::None => None,
232            ComptimeOptionExpand::Some(strides) => {
233                Some(strides.into_iter().map(|it| it.read_value(scope)).collect())
234            }
235        };
236        let clamp_mode = ClampMode::from(self.clamp_mode);
237
238        let op = CreateLayoutOp::new(scope.ctx_mut(), shape.collect(), strides, clamp_mode);
239        scope.register(&op);
240        let layout = op.get_result(scope.ctx());
241
242        TensorViewExpand {
243            buffer: self.buffer,
244            layout: layout.into(),
245            view: ComptimeOptionExpand::None,
246        }
247    }
248}
249
250impl<T: CubePrimitive> LaunchArg for TensorView<T> {
251    type RuntimeArg = TensorViewBuilderLaunch<T>;
252    type CompilationArg = TensorViewBuilderCompilationArg<T>;
253
254    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
255        TensorViewBuilder::<T>::register(arg, launcher)
256    }
257
258    fn expand(
259        arg: &Self::CompilationArg,
260        builder: &mut KernelBuilder,
261    ) -> <Self as CubeType>::ExpandType {
262        let build = TensorViewBuilder::<T>::expand(arg, builder);
263        build.__expand_finish_method(&builder.scope)
264    }
265}
cubecl_core/frontend/tensor_layout.rs

cubecl_core/frontend/
tensor_layout.rs
