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
66    fn __expand_as_type(_scope: &Scope) -> TypeHandle {
67        unimplemented!()
68    }
69}
70
71impl NativeAssign for TensorLayout {}
72
73impl CubeType for TensorReinterpret {
74    type ExpandType = NativeExpand<TensorReinterpret>;
75}
76
77impl CubeDebug for TensorReinterpret {}
78impl CubePrimitive for TensorReinterpret {
79    type Scalar = u32;
80    type Size = Const<1>;
81    type WithScalar<S: Scalar> = S;
82
83    fn from_const_value(_: cubecl_ir::ConstantValue) -> Self {
84        panic!("Can't construct tensor layout from constant")
85    }
86
87    fn __expand_as_type(_scope: &Scope) -> TypeHandle {
88        unimplemented!()
89    }
90}
91
92impl NativeAssign for TensorReinterpret {}
93
94#[derive(CubeType, CubeLaunch)]
95pub struct TensorViewBuilder<T: CubePrimitive> {
96    #[allow(unused)]
97    buffer: Box<[T]>,
98    #[allow(unused)]
99    shape: Sequence<u32>,
100    /// Strides default to contiguous strides
101    strides: ComptimeOption<Sequence<u32>>,
102    #[cube(comptime)]
103    clamp_mode: TensorClampMode,
104}
105
106#[cube]
107impl<T: CubePrimitive> TensorView<T> {
108    #[allow(clippy::new_ret_no_self)]
109    pub fn new(buffer: &[T], shape: Sequence<u32>) -> TensorViewBuilder<T> {
110        TensorViewBuilder::<T> {
111            buffer: unsafe { buffer.as_boxed_unchecked() },
112            shape,
113            strides: ComptimeOption::new_None(),
114            clamp_mode: comptime![TensorClampMode::Constant(0)],
115        }
116    }
117
118    #[allow(unused)]
119    pub fn slice(&self, offs: Sequence<u32>, shape: Sequence<u32>) -> TensorView<T> {
120        intrinsic!(|scope| {
121            assert_eq!(
122                offs.len(),
123                self.layout.rank(scope),
124                "Offsets and view rank must match"
125            );
126            assert_eq!(
127                offs.len(),
128                shape.len(),
129                "Offsets and shape must have same rank"
130            );
131            let layout = self.layout.read_value(scope);
132            let offs = offs.iter_cloned().map(|it| it.read_value(scope)).collect();
133            let shape = shape.iter_cloned().map(|it| it.read_value(scope)).collect();
134            let slice_op = SliceOp::new(scope.ctx_mut(), layout, offs, shape);
135            scope.register(&slice_op);
136            let new_layout = slice_op.get_result(scope.ctx());
137            TensorViewExpand {
138                buffer: self.buffer.clone(),
139                layout: new_layout.into(),
140                view: self.view.clone(),
141            }
142        })
143    }
144}
145
146impl NativeExpand<TensorLayout> {
147    fn rank(&self, scope: &Scope) -> usize {
148        let ty = self.read_value(scope).get_type(scope.ctx());
149        let ctx = scope.ctx();
150        let ty = ty.deref(ctx);
151        let TensorLayoutType { rank, .. } = ty.downcast_ref().unwrap();
152        *rank
153    }
154}
155
156impl<T: CubePrimitive> TensorView<T> {
157    pub fn permuted(&self, _permutation: Sequence<usize>) -> TensorView<T> {
158        unexpanded!()
159    }
160}
161
162impl<T: CubePrimitive> TensorViewExpand<T> {
163    pub fn __expand_permuted_method(
164        self,
165        scope: &Scope,
166        permutation: SequenceExpand<usize>,
167    ) -> TensorViewExpand<T> {
168        let dims = permutation.len();
169        assert!(dims <= 5, "Max 5 dims allowed");
170        let permutation = permutation
171            .iter_cloned()
172            .map(|it| {
173                it.constant()
174                    .expect("permutation must be constant")
175                    .as_usize()
176            })
177            .collect::<alloc::vec::Vec<_>>();
178        let ty = TensorViewType::get(scope.ctx(), permutation.len(), false, permutation);
179        let op = CreateViewOp::new(scope.ctx_mut(), ty.into());
180        scope.register(&op);
181        let view = op.get_result(scope.ctx());
182
183        TensorViewExpand {
184            buffer: self.buffer,
185            layout: self.layout,
186            view: ComptimeOptionExpand::Some(view.into()),
187        }
188    }
189}
190
191impl<T: CubePrimitive> TensorViewBuilder<T> {
192    pub fn with_strides(mut self, strides: Sequence<u32>) -> Self {
193        self.strides = ComptimeOption::Some(strides);
194        self
195    }
196
197    pub fn with_clamp_mode(mut self, clamp_mode: TensorClampMode) -> Self {
198        self.clamp_mode = clamp_mode;
199        self
200    }
201
202    pub fn finish(self) -> TensorView<T> {
203        unexpanded!()
204    }
205}
206
207impl<T: CubePrimitive> TensorViewBuilderExpand<T> {
208    pub fn __expand_with_strides_method(
209        mut self,
210        _scope: &Scope,
211        strides: SequenceExpand<u32>,
212    ) -> Self {
213        self.strides = ComptimeOptionExpand::Some(strides);
214        self
215    }
216
217    pub fn __expand_with_clamp_mode_method(
218        mut self,
219        _scope: &Scope,
220        clamp_mode: TensorClampMode,
221    ) -> Self {
222        self.clamp_mode = clamp_mode;
223        self
224    }
225
226    pub fn __expand_finish_method(self, scope: &Scope) -> TensorViewExpand<T> {
227        let shape = self.shape.into_iter().map(|it| it.read_value(scope));
228        let strides = match self.strides {
229            ComptimeOptionExpand::None => None,
230            ComptimeOptionExpand::Some(strides) => {
231                Some(strides.into_iter().map(|it| it.read_value(scope)).collect())
232            }
233        };
234        let clamp_mode = ClampMode::from(self.clamp_mode);
235
236        let op = CreateLayoutOp::new(scope.ctx_mut(), shape.collect(), strides, clamp_mode);
237        scope.register(&op);
238        let layout = op.get_result(scope.ctx());
239
240        TensorViewExpand {
241            buffer: self.buffer,
242            layout: layout.into(),
243            view: ComptimeOptionExpand::None,
244        }
245    }
246}
247
248impl<T: CubePrimitive> LaunchArg for TensorView<T> {
249    type RuntimeArg<R: Runtime> = TensorViewBuilderLaunch<T, R>;
250    type CompilationArg = TensorViewBuilderCompilationArg<T>;
251
252    fn register<R: Runtime>(
253        arg: Self::RuntimeArg<R>,
254        launcher: &mut KernelLauncher<R>,
255    ) -> Self::CompilationArg {
256        TensorViewBuilder::<T>::register(arg, launcher)
257    }
258
259    fn expand(
260        arg: &Self::CompilationArg,
261        builder: &mut KernelBuilder,
262    ) -> <Self as CubeType>::ExpandType {
263        let build = TensorViewBuilder::<T>::expand(arg, builder);
264        build.__expand_finish_method(&builder.scope)
265    }
266}
cubecl_core/frontend/tensor_layout.rs

cubecl_core/frontend/
tensor_layout.rs
