cubecl_std/tensor/layout/
virtual.rs

1use std::{marker::PhantomData, sync::Arc};
2
3use cubecl::prelude::*;
4use cubecl_core::{self as cubecl, intrinsic, ir::Scope, unexpanded};
5
6use crate::tensor::layout::{Coordinates, Layout, LayoutExpand};
7
8/// A virtual layout, to carry a layout without the need for generic parameters everywhere.
9/// `C` represents the coordinate space of the underlying layout.
10#[derive(Clone)]
11pub struct VirtualLayout<C: Coordinates, S: Coordinates> {
12    _coords: PhantomData<(C, S)>,
13}
14
15impl<C: Coordinates, S: Coordinates> Copy for VirtualLayout<C, S> {}
16unsafe impl<C: Coordinates, S: Coordinates> Send for VirtualLayout<C, S> {}
17unsafe impl<C: Coordinates, S: Coordinates> Sync for VirtualLayout<C, S> {}
18
19#[derive(Clone)]
20pub struct VirtualLayoutExpand<C: Coordinates, S: Coordinates> {
21    pub(crate) state: Arc<dyn VirtualLayoutOperationsExpand<C, S>>,
22}
23
24#[cube]
25impl<C: Coordinates, S: Coordinates> VirtualLayout<C, S> {
26    /// Virtual version of [`Layout::to_source_pos`]
27    #[allow(unused)]
28    pub fn to_source_pos(&self, pos: C) -> S {
29        intrinsic!(|scope| { self.state.__expand_to_source_pos_virt_method(scope, pos) })
30    }
31
32    /// Virtual version of [`Layout::to_source_pos_checked`]
33    #[allow(unused)]
34    pub fn to_source_pos_checked(&self, pos: C) -> (S, bool) {
35        intrinsic!(|scope| {
36            self.state
37                .__expand_to_source_pos_checked_virt_method(scope, pos)
38        })
39    }
40
41    /// Virtual version of [`Layout::shape`]
42    pub fn shape(&self) -> C {
43        intrinsic!(|scope| { self.state.__expand_shape_virt_method(scope) })
44    }
45
46    /// Virtual version of [`Layout::is_in_bounds`]
47    #[allow(unused)]
48    pub fn is_in_bounds(&self, pos: C) -> bool {
49        intrinsic!(|scope| { self.state.__expand_is_in_bounds_virt_method(scope, pos) })
50    }
51}
52
53impl<C: Coordinates, S: Coordinates> VirtualLayout<C, S> {
54    /// Create a new virtual layout from a concrete one
55    pub fn new<L: Layout<Coordinates = C, SourceCoordinates = S>>(
56        _layout: L,
57    ) -> VirtualLayout<C, S> {
58        unexpanded!()
59    }
60
61    /// Expand function of [`VirtualLayout::`__`expand_new`]
62    pub fn __expand_new<L: Layout<Coordinates = C, SourceCoordinates = S> + 'static>(
63        _scope: &Scope,
64        layout: L::ExpandType,
65    ) -> VirtualLayoutExpand<C, S> {
66        VirtualLayoutExpand::new::<L::ExpandType>(layout)
67    }
68}
69
70impl<C: Coordinates, S: Coordinates> VirtualLayoutExpand<C, S> {
71    /// Create a new virtual layout from a concrete one
72    pub fn new<L: VirtualLayoutOperationsExpand<C, S> + 'static>(
73        layout: L,
74    ) -> VirtualLayoutExpand<C, S> {
75        VirtualLayoutExpand::<C, S> {
76            state: Arc::new(layout),
77        }
78    }
79}
80
81impl<C: Coordinates, S: Coordinates> CubeType for VirtualLayout<C, S> {
82    type ExpandType = VirtualLayoutExpand<C, S>;
83}
84
85impl<C: Coordinates, S: Coordinates> IntoExpand for VirtualLayoutExpand<C, S> {
86    type Expand = VirtualLayoutExpand<C, S>;
87
88    fn into_expand(self, _: &Scope) -> Self::Expand {
89        self
90    }
91}
92
93impl<C: Coordinates, S: Coordinates> ExpandTypeClone for VirtualLayoutExpand<C, S> {
94    fn clone_unchecked(&self) -> Self {
95        self.clone()
96    }
97}
98
99impl<C: Coordinates, S: Coordinates> IntoMut for VirtualLayoutExpand<C, S> {
100    fn into_mut(self, _scope: &Scope) -> Self {
101        self
102    }
103}
104
105impl<C: Coordinates, S: Coordinates> CubeDebug for VirtualLayoutExpand<C, S> {}
106
107impl<C: Coordinates, S: Coordinates> AsRefExpand for VirtualLayoutExpand<C, S> {
108    fn __expand_ref_method(&self, _: &Scope) -> &Self {
109        self
110    }
111}
112impl<C: Coordinates, S: Coordinates> AsMutExpand for VirtualLayoutExpand<C, S> {
113    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
114        self
115    }
116}
117
118// We need to seal the trait to allow us to blanket implement `From<L>` below
119mod private {
120    pub trait Sealed {}
121}
122pub trait VirtualLayoutOperationsExpand<C: CubeType, S: CubeType>: private::Sealed {
123    fn __expand_to_source_pos_virt_method(
124        &self,
125        scope: &Scope,
126        pos: <C as CubeType>::ExpandType,
127    ) -> <S as CubeType>::ExpandType;
128    fn __expand_to_source_pos_checked_virt_method(
129        &self,
130        scope: &Scope,
131        pos: <C as CubeType>::ExpandType,
132    ) -> <(S, bool) as CubeType>::ExpandType;
133    fn __expand_shape_virt_method(&self, scope: &Scope) -> <C as CubeType>::ExpandType;
134    fn __expand_is_in_bounds_virt_method(
135        &self,
136        scope: &Scope,
137        pos: <C as CubeType>::ExpandType,
138    ) -> NativeExpand<bool>;
139}
140
141impl<L: LayoutExpand> private::Sealed for L {}
142impl<L: LayoutExpand> VirtualLayoutOperationsExpand<L::Coordinates, L::SourceCoordinates> for L {
143    fn __expand_to_source_pos_virt_method(
144        &self,
145        scope: &Scope,
146        pos: <L::Coordinates as CubeType>::ExpandType,
147    ) -> <L::SourceCoordinates as CubeType>::ExpandType {
148        <L as LayoutExpand>::__expand_to_source_pos_method(self, scope, pos)
149    }
150
151    fn __expand_to_source_pos_checked_virt_method(
152        &self,
153        scope: &Scope,
154        pos: <L::Coordinates as CubeType>::ExpandType,
155    ) -> <(L::SourceCoordinates, bool) as CubeType>::ExpandType {
156        <L as LayoutExpand>::__expand_to_source_pos_checked_method(self, scope, pos)
157    }
158
159    fn __expand_shape_virt_method(
160        &self,
161        scope: &Scope,
162    ) -> <L::Coordinates as CubeType>::ExpandType {
163        <L as LayoutExpand>::__expand_shape_method(self, scope)
164    }
165
166    fn __expand_is_in_bounds_virt_method(
167        &self,
168        scope: &Scope,
169        pos: <L::Coordinates as CubeType>::ExpandType,
170    ) -> NativeExpand<bool> {
171        <L as LayoutExpand>::__expand_is_in_bounds_method(self, scope, pos)
172    }
173}
174
175impl<C: Coordinates, S: Coordinates, L: VirtualLayoutOperationsExpand<C, S> + 'static> From<L>
176    for VirtualLayoutExpand<C, S>
177{
178    fn from(value: L) -> Self {
179        VirtualLayoutExpand::new(value)
180    }
181}
182
183impl<L: Layout + 'static> From<L> for VirtualLayout<L::Coordinates, L::SourceCoordinates> {
184    fn from(_value: L) -> Self {
185        VirtualLayout {
186            _coords: PhantomData,
187        }
188    }
189}
190
191mod launch {
192    use alloc::rc::Rc;
193    use core::cell::RefCell;
194
195    use cubecl_core::{
196        format::DebugRaw,
197        hash::{StableHash, StableHasher},
198    };
199
200    use super::*;
201
202    type ExpandFn<C, S> =
203        Rc<RefCell<dyn FnMut(&mut KernelBuilder) -> VirtualLayoutExpand<C, S> + Send>>;
204
205    pub struct VirtualLayoutLaunch<C: Coordinates, S: Coordinates, R: Runtime> {
206        _phantom_runtime: core::marker::PhantomData<R>,
207        #[allow(clippy::type_complexity)]
208        register: Box<
209            dyn FnOnce(&mut KernelLauncher<R>) -> VirtualLayoutCompilationArg<C, S> + Send + Sync,
210        >,
211    }
212
213    impl<C: Coordinates, S: Coordinates, R: Runtime> VirtualLayoutLaunch<C, S, R> {
214        pub fn new<L: Layout<Coordinates = C, SourceCoordinates = S> + LaunchArg>(
215            layout: L::RuntimeArg<R>,
216        ) -> Self {
217            Self {
218                _phantom_runtime: PhantomData,
219                register: Box::new(move |launcher| {
220                    let comp_arg = L::register(layout, launcher);
221                    let comp_arg_2 = comp_arg.clone();
222                    let expand = move |builder: &mut KernelBuilder| {
223                        VirtualLayoutExpand::new(L::expand(&comp_arg_2, builder))
224                    };
225                    VirtualLayoutCompilationArg::new::<L::CompilationArg>(
226                        comp_arg,
227                        Rc::new(RefCell::new(expand)),
228                    )
229                }),
230            }
231        }
232    }
233
234    #[derive(Clone)]
235    pub struct VirtualLayoutCompilationArg<C: Coordinates, S: Coordinates> {
236        type_name: String,
237        debug: Rc<dyn core::fmt::Debug>,
238        hash: StableHash,
239        expand: ExpandFn<C, S>,
240    }
241
242    // SAFETY: The struct is readonly, so `Sync` is safe to implement
243    unsafe impl<C: Coordinates, S: Coordinates> Send for VirtualLayoutCompilationArg<C, S> {}
244    unsafe impl<C: Coordinates, S: Coordinates> Sync for VirtualLayoutCompilationArg<C, S> {}
245
246    impl<C: Coordinates, S: Coordinates> VirtualLayoutCompilationArg<C, S> {
247        pub fn new<L: CompilationArg + 'static>(arg: L, expand: ExpandFn<C, S>) -> Self {
248            // Hash ahead of time so we don't need to store the actual data, which would be far
249            // more complex
250            let hash = StableHasher::hash_one(&arg);
251            Self {
252                type_name: core::any::type_name::<L>().to_string(),
253                debug: Rc::new(arg),
254                hash,
255                expand,
256            }
257        }
258    }
259
260    impl<C: Coordinates, S: Coordinates> PartialEq for VirtualLayoutCompilationArg<C, S> {
261        fn eq(&self, other: &Self) -> bool {
262            self.type_name == other.type_name && self.hash == other.hash
263        }
264    }
265    impl<C: Coordinates, S: Coordinates> Eq for VirtualLayoutCompilationArg<C, S> {}
266
267    impl<C: Coordinates, S: Coordinates> core::hash::Hash for VirtualLayoutCompilationArg<C, S> {
268        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
269            self.type_name.hash(state);
270            self.hash.hash(state);
271        }
272    }
273
274    impl<C: Coordinates, S: Coordinates> core::fmt::Debug for VirtualLayoutCompilationArg<C, S> {
275        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
276            f.debug_struct(stringify!(VirtualLayout))
277                .field("type", &DebugRaw(&self.type_name))
278                .field("value", &self.debug)
279                .finish()
280        }
281    }
282
283    impl<C: Coordinates + 'static, S: Coordinates + 'static> LaunchArg for VirtualLayout<C, S> {
284        type RuntimeArg<R: Runtime> = VirtualLayoutLaunch<C, S, R>;
285        type CompilationArg = VirtualLayoutCompilationArg<C, S>;
286
287        fn register<R: Runtime>(
288            arg: Self::RuntimeArg<R>,
289            launcher: &mut KernelLauncher<R>,
290        ) -> Self::CompilationArg {
291            let func = arg.register;
292            func(launcher)
293        }
294        fn expand(
295            arg: &Self::CompilationArg,
296            builder: &mut KernelBuilder,
297        ) -> <Self as CubeType>::ExpandType {
298            let mut expand = arg.expand.borrow_mut();
299            expand(builder)
300        }
301    }
302}
303
304pub use launch::*;
cubecl_std/tensor/layout/virtual.rs

cubecl_std/tensor/layout/
virtual.rs
