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> {
206        #[allow(clippy::type_complexity)]
207        register:
208            Box<dyn FnOnce(&mut KernelLauncher) -> VirtualLayoutCompilationArg<C, S> + Send + Sync>,
209    }
210
211    impl<C: Coordinates, S: Coordinates> VirtualLayoutLaunch<C, S> {
212        pub fn new<L: Layout<Coordinates = C, SourceCoordinates = S> + LaunchArg>(
213            layout: L::RuntimeArg,
214        ) -> Self {
215            Self {
216                register: Box::new(move |launcher| {
217                    let comp_arg = L::register(layout, launcher);
218                    let comp_arg_2 = comp_arg.clone();
219                    let expand = move |builder: &mut KernelBuilder| {
220                        VirtualLayoutExpand::new(L::expand(&comp_arg_2, builder))
221                    };
222                    VirtualLayoutCompilationArg::new::<L::CompilationArg>(
223                        comp_arg,
224                        Rc::new(RefCell::new(expand)),
225                    )
226                }),
227            }
228        }
229    }
230
231    #[derive(Clone)]
232    pub struct VirtualLayoutCompilationArg<C: Coordinates, S: Coordinates> {
233        type_name: String,
234        debug: Rc<dyn core::fmt::Debug>,
235        hash: StableHash,
236        expand: ExpandFn<C, S>,
237    }
238
239    // SAFETY: The struct is readonly, so `Sync` is safe to implement
240    unsafe impl<C: Coordinates, S: Coordinates> Send for VirtualLayoutCompilationArg<C, S> {}
241    unsafe impl<C: Coordinates, S: Coordinates> Sync for VirtualLayoutCompilationArg<C, S> {}
242
243    impl<C: Coordinates, S: Coordinates> VirtualLayoutCompilationArg<C, S> {
244        pub fn new<L: CompilationArg + 'static>(arg: L, expand: ExpandFn<C, S>) -> Self {
245            // Hash ahead of time so we don't need to store the actual data, which would be far
246            // more complex
247            let hash = StableHasher::hash_one(&arg);
248            Self {
249                type_name: core::any::type_name::<L>().to_string(),
250                debug: Rc::new(arg),
251                hash,
252                expand,
253            }
254        }
255    }
256
257    impl<C: Coordinates, S: Coordinates> PartialEq for VirtualLayoutCompilationArg<C, S> {
258        fn eq(&self, other: &Self) -> bool {
259            self.type_name == other.type_name && self.hash == other.hash
260        }
261    }
262    impl<C: Coordinates, S: Coordinates> Eq for VirtualLayoutCompilationArg<C, S> {}
263
264    impl<C: Coordinates, S: Coordinates> core::hash::Hash for VirtualLayoutCompilationArg<C, S> {
265        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
266            self.type_name.hash(state);
267            self.hash.hash(state);
268        }
269    }
270
271    impl<C: Coordinates, S: Coordinates> core::fmt::Debug for VirtualLayoutCompilationArg<C, S> {
272        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
273            f.debug_struct(stringify!(VirtualLayout))
274                .field("type", &DebugRaw(&self.type_name))
275                .field("value", &self.debug)
276                .finish()
277        }
278    }
279
280    impl<C: Coordinates + 'static, S: Coordinates + 'static> LaunchArg for VirtualLayout<C, S> {
281        type RuntimeArg = VirtualLayoutLaunch<C, S>;
282        type CompilationArg = VirtualLayoutCompilationArg<C, S>;
283
284        fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
285            let func = arg.register;
286            func(launcher)
287        }
288        fn expand(
289            arg: &Self::CompilationArg,
290            builder: &mut KernelBuilder,
291        ) -> <Self as CubeType>::ExpandType {
292            let mut expand = arg.expand.borrow_mut();
293            expand(builder)
294        }
295    }
296}
297
298pub use launch::*;
cubecl_std/tensor/layout/virtual.rs

cubecl_std/tensor/layout/
virtual.rs
