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_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| { self.state.__expand_to_source_pos_checked_method(scope, pos) })
36    }
37
38    /// Virtual version of [Layout::shape]
39    pub fn shape(&self) -> C {
40        intrinsic!(|scope| { self.state.__expand_shape_method(scope) })
41    }
42
43    /// Virtual version of [Layout::is_in_bounds]
44    #[allow(unused)]
45    pub fn is_in_bounds(&self, pos: C) -> bool {
46        intrinsic!(|scope| { self.state.__expand_is_in_bounds_method(scope, pos) })
47    }
48}
49
50impl<C: Coordinates, S: Coordinates> VirtualLayout<C, S> {
51    /// Create a new virtual layout from a concrete one
52    pub fn new<L: Layout<Coordinates = C, SourceCoordinates = S>>(
53        _layout: L,
54    ) -> VirtualLayout<C, S> {
55        unexpanded!()
56    }
57
58    /// Expand function of [VirtualLayout::__expand_new]
59    pub fn __expand_new<L: Layout<Coordinates = C, SourceCoordinates = S> + 'static>(
60        _scope: &mut Scope,
61        layout: L::ExpandType,
62    ) -> VirtualLayoutExpand<C, S> {
63        VirtualLayoutExpand::new::<L::ExpandType>(layout)
64    }
65}
66
67impl<C: Coordinates, S: Coordinates> VirtualLayoutExpand<C, S> {
68    /// Create a new virtual layout from a concrete one
69    pub fn new<L: VirtualLayoutOperationsExpand<C, S> + 'static>(
70        layout: L,
71    ) -> VirtualLayoutExpand<C, S> {
72        VirtualLayoutExpand::<C, S> {
73            state: Arc::new(layout),
74        }
75    }
76}
77
78impl<C: Coordinates, S: Coordinates> CubeType for VirtualLayout<C, S> {
79    type ExpandType = VirtualLayoutExpand<C, S>;
80}
81
82impl<C: Coordinates, S: Coordinates> IntoMut for VirtualLayoutExpand<C, S> {
83    fn into_mut(self, _scope: &mut Scope) -> Self {
84        self
85    }
86}
87
88impl<C: Coordinates, S: Coordinates> CubeDebug for VirtualLayoutExpand<C, S> {}
89
90// We need to seal the trait to allow us to blanket implement `From<L>` below
91mod private {
92    pub trait Sealed {}
93}
94pub trait VirtualLayoutOperationsExpand<C: CubeType, S: CubeType>: private::Sealed {
95    fn __expand_to_source_pos_method(
96        &self,
97        scope: &mut Scope,
98        pos: <C as CubeType>::ExpandType,
99    ) -> <S as CubeType>::ExpandType;
100    fn __expand_to_source_pos_checked_method(
101        &self,
102        scope: &mut Scope,
103        pos: <C as CubeType>::ExpandType,
104    ) -> <(S, bool) as CubeType>::ExpandType;
105    fn __expand_shape_method(&self, scope: &mut Scope) -> <C as CubeType>::ExpandType;
106    fn __expand_is_in_bounds_method(
107        &self,
108        scope: &mut Scope,
109        pos: <C as CubeType>::ExpandType,
110    ) -> ExpandElementTyped<bool>;
111}
112
113impl<L: LayoutExpand> private::Sealed for L {}
114impl<L: LayoutExpand> VirtualLayoutOperationsExpand<L::Coordinates, L::SourceCoordinates> for L {
115    fn __expand_to_source_pos_method(
116        &self,
117        scope: &mut Scope,
118        pos: <L::Coordinates as CubeType>::ExpandType,
119    ) -> <L::SourceCoordinates as CubeType>::ExpandType {
120        <L as LayoutExpand>::__expand_to_source_pos_method(self.clone(), scope, pos)
121    }
122
123    fn __expand_to_source_pos_checked_method(
124        &self,
125        scope: &mut Scope,
126        pos: <L::Coordinates as CubeType>::ExpandType,
127    ) -> <(L::SourceCoordinates, bool) as CubeType>::ExpandType {
128        <L as LayoutExpand>::__expand_to_source_pos_checked_method(self.clone(), scope, pos)
129    }
130
131    fn __expand_shape_method(&self, scope: &mut Scope) -> <L::Coordinates as CubeType>::ExpandType {
132        <L as LayoutExpand>::__expand_shape_method(self.clone(), scope)
133    }
134
135    fn __expand_is_in_bounds_method(
136        &self,
137        scope: &mut Scope,
138        pos: <L::Coordinates as CubeType>::ExpandType,
139    ) -> ExpandElementTyped<bool> {
140        <L as LayoutExpand>::__expand_is_in_bounds_method(self.clone(), scope, pos)
141    }
142}
143
144impl<C: Coordinates, S: Coordinates, L: VirtualLayoutOperationsExpand<C, S> + 'static> From<L>
145    for VirtualLayoutExpand<C, S>
146{
147    fn from(value: L) -> Self {
148        VirtualLayoutExpand::new(value)
149    }
150}
151
152impl<L: Layout + 'static> From<L> for VirtualLayout<L::Coordinates, L::SourceCoordinates> {
153    fn from(_value: L) -> Self {
154        VirtualLayout {
155            _coords: PhantomData,
156        }
157    }
158}
159
160mod launch {
161    use core::hash::BuildHasher;
162    use spin::Mutex;
163
164    use super::*;
165
166    type ExpandFn<C, S> =
167        Arc<Mutex<dyn FnMut(&mut KernelBuilder) -> VirtualLayoutExpand<C, S> + Send>>;
168
169    pub struct VirtualLayoutLaunch<'a, C: Coordinates, S: Coordinates, R: Runtime> {
170        _phantom_runtime: core::marker::PhantomData<R>,
171        _phantom_a: core::marker::PhantomData<&'a ()>,
172        inner: Arc<dyn ArgSettings<R> + 'a>,
173        hashed_arg: VirtualLayoutCompilationArg<C, S>,
174    }
175
176    impl<'a, C: Coordinates, S: Coordinates, R: cubecl::prelude::Runtime>
177        VirtualLayoutLaunch<'a, C, S, R>
178    {
179        pub fn new<L: Layout<Coordinates = C, SourceCoordinates = S> + LaunchArg>(
180            layout: L::RuntimeArg<'a, R>,
181        ) -> Self {
182            let comp_arg = L::compilation_arg(&layout);
183            let comp_arg_2 = comp_arg.clone();
184            let expand = move |builder: &mut KernelBuilder| {
185                let expand = L::expand(&comp_arg_2, builder);
186                VirtualLayoutExpand::new(expand)
187            };
188            let comp_arg_2 = comp_arg.clone();
189            let expand_out = move |builder: &mut KernelBuilder| {
190                let expand = L::expand_output(&comp_arg_2, builder);
191                VirtualLayoutExpand::new(expand)
192            };
193            let hashed_arg = VirtualLayoutCompilationArg::new::<L::CompilationArg>(
194                &comp_arg,
195                Arc::new(Mutex::new(expand)),
196                Arc::new(Mutex::new(expand_out)),
197            );
198
199            Self {
200                _phantom_runtime: PhantomData,
201                _phantom_a: PhantomData,
202                inner: Arc::new(layout),
203                hashed_arg,
204            }
205        }
206    }
207    impl<'a, C: Coordinates, S: Coordinates, R: cubecl::prelude::Runtime> ArgSettings<R>
208        for VirtualLayoutLaunch<'a, C, S, R>
209    {
210        fn register(&self, launcher: &mut cubecl::prelude::KernelLauncher<R>) {
211            self.inner.register(launcher);
212        }
213    }
214
215    #[derive(Clone)]
216    pub struct VirtualLayoutCompilationArg<C: Coordinates, S: Coordinates> {
217        type_name: String,
218        debug_string: String,
219        hash: u64,
220        expand: ExpandFn<C, S>,
221        expand_output: ExpandFn<C, S>,
222    }
223
224    impl<C: Coordinates, S: Coordinates> VirtualLayoutCompilationArg<C, S> {
225        pub fn new<L: CompilationArg>(
226            arg: &L,
227            expand: ExpandFn<C, S>,
228            expand_output: ExpandFn<C, S>,
229        ) -> Self {
230            // Hash ahead of time so we don't need to store the actual data, which would be far
231            // more complex
232            let state = foldhash::fast::FixedState::default();
233            let hash = state.hash_one(arg);
234            Self {
235                type_name: core::any::type_name::<L>().to_string(),
236                debug_string: format!("{arg:?}"),
237                hash,
238                expand,
239                expand_output,
240            }
241        }
242    }
243
244    impl<C: Coordinates, S: Coordinates> PartialEq for VirtualLayoutCompilationArg<C, S> {
245        fn eq(&self, other: &Self) -> bool {
246            self.type_name == other.type_name && self.hash == other.hash
247        }
248    }
249    impl<C: Coordinates, S: Coordinates> Eq for VirtualLayoutCompilationArg<C, S> {}
250
251    impl<C: Coordinates + 'static, S: Coordinates + 'static> CompilationArg
252        for VirtualLayoutCompilationArg<C, S>
253    {
254    }
255
256    impl<C: Coordinates, S: Coordinates> core::hash::Hash for VirtualLayoutCompilationArg<C, S> {
257        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
258            self.type_name.hash(state);
259            self.hash.hash(state);
260        }
261    }
262
263    impl<C: Coordinates, S: Coordinates> core::fmt::Debug for VirtualLayoutCompilationArg<C, S> {
264        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265            f.write_str(stringify!(VirtualLayout))?;
266            f.write_str("{")?;
267            f.write_fmt(format_args!("type: {:?},", &self.type_name))?;
268            f.write_fmt(format_args!("value: {:?},", &self.debug_string))?;
269            f.write_str("}")?;
270            Ok(())
271        }
272    }
273
274    impl<C: Coordinates + 'static, S: Coordinates + 'static> LaunchArg for VirtualLayout<C, S> {
275        type RuntimeArg<'a, R: Runtime> = VirtualLayoutLaunch<'a, C, S, R>;
276        type CompilationArg = VirtualLayoutCompilationArg<C, S>;
277
278        fn compilation_arg<'a, R: Runtime>(
279            runtime_arg: &Self::RuntimeArg<'a, R>,
280        ) -> Self::CompilationArg {
281            runtime_arg.hashed_arg.clone()
282        }
283        fn expand(
284            arg: &Self::CompilationArg,
285            builder: &mut KernelBuilder,
286        ) -> <Self as CubeType>::ExpandType {
287            let mut expand = arg.expand.as_ref().lock();
288            expand(builder)
289        }
290        fn expand_output(
291            arg: &Self::CompilationArg,
292            builder: &mut KernelBuilder,
293        ) -> <Self as CubeType>::ExpandType {
294            let mut expand = arg.expand_output.as_ref().lock();
295            expand(builder)
296        }
297    }
298}
299
300pub use launch::*;
cubecl_std/tensor/layout/virtual.rs

cubecl_std/tensor/layout/
virtual.rs
