cubecl_std/tensor/view/
launch.rs

1use cubecl_core::prelude::*;
2use std::{marker::PhantomData, ops::Deref};
3
4use crate::tensor::{
5    View, ViewExpand, VirtualViewMutExpand,
6    layout::{Coordinates, Coords1d, Layout, VirtualLayoutExpand},
7};
8
9mod layout {
10    use core::{cell::RefCell, fmt::Debug, hash::Hash};
11
12    use alloc::rc::Rc;
13    use cubecl_core::{
14        self as cubecl,
15        format::DebugRaw,
16        hash::{StableHash, StableHasher},
17        zspace::{Shape, Strides, metadata::Metadata},
18    };
19
20    use crate::tensor::layout::LayoutExpand;
21
22    use super::*;
23
24    #[allow(clippy::len_without_is_empty)]
25    pub trait MemoryArg: 'static {
26        fn len(&self) -> usize;
27        fn shape(&self) -> &[usize];
28        fn strides(&self) -> &[usize];
29    }
30
31    impl MemoryArg for TensorArg {
32        fn len(&self) -> usize {
33            self.size()
34        }
35
36        fn shape(&self) -> &[usize] {
37            self.shape()
38        }
39
40        fn strides(&self) -> &[usize] {
41            self.strides()
42        }
43    }
44    impl MemoryArg for BufferArg {
45        fn len(&self) -> usize {
46            self.size()
47        }
48
49        fn shape(&self) -> &[usize] {
50            self.shape()
51        }
52
53        fn strides(&self) -> &[usize] {
54            &[1]
55        }
56    }
57    impl<K: TensorMapKind> MemoryArg for TensorMapArg<K> {
58        fn len(&self) -> usize {
59            self.tensor.size()
60        }
61
62        fn shape(&self) -> &[usize] {
63            self.tensor.shape()
64        }
65
66        fn strides(&self) -> &[usize] {
67            self.tensor.strides()
68        }
69    }
70
71    impl MemoryArg for Metadata {
72        fn len(&self) -> usize {
73            self.shape.num_elements()
74        }
75
76        fn shape(&self) -> &[usize] {
77            &self.shape
78        }
79
80        fn strides(&self) -> &[usize] {
81            &self.strides
82        }
83    }
84
85    /// Special launch arg that gets the handle and types of the view, to allow inferring launch
86    /// state based on type/handle metadata, avoiding duplication. All `LaunchArg`s also implement
87    /// this trait.
88    pub trait ViewLayoutLaunchArg: CubeType + Send + Sync + 'static {
89        /// The runtime argument for the kernel.
90        type RuntimeArg: Send + Sync;
91        /// Compilation argument.
92        type CompilationArg: CompilationArg;
93
94        fn register<B: MemoryArg>(
95            arg: Self::RuntimeArg,
96            buffer: &B,
97            ty: Type,
98            launcher: &mut KernelLauncher,
99        ) -> Self::CompilationArg;
100
101        /// Register an input variable during compilation that fill the [`KernelBuilder`].
102        fn expand(
103            arg: &Self::CompilationArg,
104            ty: Type,
105            builder: &mut KernelBuilder,
106        ) -> <Self as CubeType>::ExpandType;
107
108        /// Register an output variable during compilation that fill the [`KernelBuilder`].
109        fn expand_output(
110            arg: &Self::CompilationArg,
111            ty: Type,
112            builder: &mut KernelBuilder,
113        ) -> <Self as CubeType>::ExpandType {
114            Self::expand(arg, ty, builder)
115        }
116    }
117
118    impl<T: LaunchArg + Send + Sync> ViewLayoutLaunchArg for T {
119        type RuntimeArg = <T as LaunchArg>::RuntimeArg;
120        type CompilationArg = <T as LaunchArg>::CompilationArg;
121
122        fn register<B: MemoryArg>(
123            arg: Self::RuntimeArg,
124            _buffer: &B,
125            _ty: Type,
126            launcher: &mut KernelLauncher,
127        ) -> Self::CompilationArg {
128            <T as LaunchArg>::register(arg, launcher)
129        }
130
131        fn expand(
132            arg: &Self::CompilationArg,
133            _ty: Type,
134            builder: &mut KernelBuilder,
135        ) -> <Self as CubeType>::ExpandType {
136            <T as LaunchArg>::expand(arg, builder)
137        }
138    }
139
140    pub struct VirtualViewLayoutLaunch<C: Coordinates, S: Coordinates, B: MemoryArg> {
141        #[allow(clippy::type_complexity)]
142        register: Box<
143            dyn FnOnce(&B, Type, &mut KernelLauncher) -> VirtualViewLayoutCompilationArg<C, S>
144                + Send
145                + Sync,
146        >,
147    }
148
149    impl<C: Coordinates, S: Coordinates, B: MemoryArg> VirtualViewLayoutLaunch<C, S, B> {
150        pub fn new<L: Layout<Coordinates = C, SourceCoordinates = S> + ViewLayoutLaunchArg>(
151            layout: L::RuntimeArg,
152        ) -> Self {
153            Self {
154                register: Box::new(move |buffer, ty, launcher| {
155                    let comp_arg = L::register::<B>(layout, buffer, ty, launcher);
156                    let comp_arg_2 = comp_arg.clone();
157                    let expand = Rc::new(RefCell::new(
158                        move |ty: Type, builder: &mut KernelBuilder, is_out: bool| {
159                            let expand = match is_out {
160                                true => L::expand_output(&comp_arg_2, ty, builder),
161                                false => L::expand(&comp_arg_2, ty, builder),
162                            };
163                            VirtualLayoutExpand::new(expand)
164                        },
165                    ));
166                    VirtualViewLayoutCompilationArg::new(comp_arg, expand)
167                }),
168            }
169        }
170
171        pub fn register(
172            self,
173            buffer: &B,
174            ty: Type,
175            launcher: &mut KernelLauncher,
176        ) -> VirtualViewLayoutCompilationArg<C, S> {
177            (self.register)(buffer, ty, launcher)
178        }
179    }
180
181    type ExpandFn<C, S> =
182        Rc<RefCell<dyn FnMut(Type, &mut KernelBuilder, bool) -> VirtualLayoutExpand<C, S> + Send>>;
183
184    #[derive(Clone)]
185    pub struct VirtualViewLayoutCompilationArg<C: Coordinates, S: Coordinates> {
186        type_name: String,
187        debug: Rc<dyn core::fmt::Debug>,
188        hash: StableHash,
189        expand: ExpandFn<C, S>,
190    }
191
192    // SAFETY: The struct is readonly, so `Sync` is safe to implement
193    unsafe impl<C: Coordinates, S: Coordinates> Send for VirtualViewLayoutCompilationArg<C, S> {}
194    unsafe impl<C: Coordinates, S: Coordinates> Sync for VirtualViewLayoutCompilationArg<C, S> {}
195
196    impl<C: Coordinates, S: Coordinates> VirtualViewLayoutCompilationArg<C, S> {
197        pub fn new<L: CompilationArg + 'static>(arg: L, expand: ExpandFn<C, S>) -> Self {
198            // Hash ahead of time so we don't need to store the actual data, which would be far
199            // more complex
200            let hash = StableHasher::hash_one(&arg);
201            Self {
202                type_name: core::any::type_name::<L>().to_string(),
203                debug: Rc::new(arg),
204                hash,
205                expand,
206            }
207        }
208
209        pub fn expand(&self, ty: Type, builder: &mut KernelBuilder) -> VirtualLayoutExpand<C, S> {
210            let mut expand = self.expand.borrow_mut();
211            (expand)(ty, builder, false)
212        }
213
214        pub fn expand_output(
215            &self,
216            ty: Type,
217            builder: &mut KernelBuilder,
218        ) -> VirtualLayoutExpand<C, S> {
219            let mut expand = self.expand.borrow_mut();
220            (expand)(ty, builder, true)
221        }
222    }
223
224    impl<C: Coordinates, S: Coordinates> PartialEq for VirtualViewLayoutCompilationArg<C, S> {
225        fn eq(&self, other: &Self) -> bool {
226            self.type_name == other.type_name && self.hash == other.hash
227        }
228    }
229    impl<C: Coordinates, S: Coordinates> Eq for VirtualViewLayoutCompilationArg<C, S> {}
230
231    impl<C: Coordinates, S: Coordinates> core::hash::Hash for VirtualViewLayoutCompilationArg<C, S> {
232        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
233            self.type_name.hash(state);
234            self.hash.hash(state);
235        }
236    }
237
238    impl<C: Coordinates, S: Coordinates> core::fmt::Debug for VirtualViewLayoutCompilationArg<C, S> {
239        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
240            f.debug_struct(stringify!(VirtualLayout))
241                .field("type", &DebugRaw(&self.type_name))
242                .field("value", &self.debug)
243                .finish()
244        }
245    }
246
247    #[derive(CubeType)]
248    pub struct ConcreteLayout<L: Layout + ViewLayoutLaunchArg> {
249        value: L,
250    }
251
252    #[cube]
253    impl<L: Layout + ViewLayoutLaunchArg> Layout for ConcreteLayout<L> {
254        type Coordinates = L::Coordinates;
255        type SourceCoordinates = L::SourceCoordinates;
256
257        fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
258            self.value.to_source_pos(pos)
259        }
260
261        fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
262            self.value.to_source_pos_checked(pos)
263        }
264
265        fn shape(&self) -> Self::Coordinates {
266            self.value.shape()
267        }
268
269        fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
270            self.value.is_in_bounds(pos)
271        }
272    }
273
274    impl<L: Layout + ViewLayoutLaunchArg> Deref for ConcreteLayout<L> {
275        type Target = L;
276
277        fn deref(&self) -> &Self::Target {
278            &self.value
279        }
280    }
281
282    impl<L: Layout + ViewLayoutLaunchArg> Deref for ConcreteLayoutExpand<L> {
283        type Target = <L as CubeType>::ExpandType;
284
285        fn deref(&self) -> &Self::Target {
286            &self.value
287        }
288    }
289
290    pub struct ConcreteLayoutLaunch<L: Layout + ViewLayoutLaunchArg> {
291        meta: Metadata,
292        ty: Type,
293        value: L::RuntimeArg,
294    }
295
296    impl<L: Layout + ViewLayoutLaunchArg> ConcreteLayoutLaunch<L> {
297        pub fn new(meta: Metadata, ty: Type, value: L::RuntimeArg) -> Self {
298            Self { meta, ty, value }
299        }
300
301        pub fn from_handle(handle: &TensorBinding, ty: Type, value: L::RuntimeArg) -> Self {
302            Self {
303                meta: Metadata::new(handle.shape.clone(), handle.strides.clone())
304                    .with_tiling(handle.tiling)
305                    .expect("a binding's tiling describes its own rank"),
306                ty,
307                value,
308            }
309        }
310
311        pub fn from_shape_strides(
312            shape: Shape,
313            strides: Strides,
314            ty: Type,
315            value: L::RuntimeArg,
316        ) -> Self {
317            Self {
318                meta: Metadata::new(shape, strides),
319                ty,
320                value,
321            }
322        }
323    }
324
325    pub struct ConcreteLayoutCompilationArg<L: Layout + ViewLayoutLaunchArg> {
326        ty: Type,
327        value: L::CompilationArg,
328    }
329
330    impl<L: Layout + ViewLayoutLaunchArg> Debug for ConcreteLayoutCompilationArg<L> {
331        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332            f.debug_struct("ConcreteLayoutCompilationArg")
333                .field("ty", &self.ty)
334                .field("value", &self.value)
335                .finish()
336        }
337    }
338
339    impl<L: Layout + ViewLayoutLaunchArg> Hash for ConcreteLayoutCompilationArg<L> {
340        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
341            self.ty.hash(state);
342            self.value.hash(state);
343        }
344    }
345
346    impl<L: Layout + ViewLayoutLaunchArg> Eq for ConcreteLayoutCompilationArg<L> {}
347    impl<L: Layout + ViewLayoutLaunchArg> PartialEq for ConcreteLayoutCompilationArg<L> {
348        fn eq(&self, other: &Self) -> bool {
349            self.ty == other.ty && self.value == other.value
350        }
351    }
352
353    impl<L: Layout + ViewLayoutLaunchArg> Clone for ConcreteLayoutCompilationArg<L> {
354        fn clone(&self) -> Self {
355            Self {
356                ty: self.ty,
357                value: self.value.clone(),
358            }
359        }
360    }
361
362    impl<L: Layout + ViewLayoutLaunchArg> LaunchArg for ConcreteLayout<L> {
363        type RuntimeArg = ConcreteLayoutLaunch<L>;
364        type CompilationArg = ConcreteLayoutCompilationArg<L>;
365
366        fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
367            ConcreteLayoutCompilationArg {
368                value: L::register(arg.value, &arg.meta, arg.ty, launcher),
369                ty: arg.ty,
370            }
371        }
372
373        fn expand(
374            arg: &Self::CompilationArg,
375            builder: &mut KernelBuilder,
376        ) -> <Self as CubeType>::ExpandType {
377            ConcreteLayoutExpand {
378                value: L::expand(&arg.value, arg.ty, builder),
379            }
380        }
381    }
382}
383
384pub use layout::*;
385
386mod dynamic {
387    use cubecl_common::quant::scheme::QuantScheme;
388
389    use crate::{
390        quant::{
391            self,
392            view::{RegisterDynamic, run_with_quant_type},
393        },
394        tensor::{
395            ViewMut, ViewMutExpand, VirtualViewExpand,
396            launch::layout::{ViewLayoutLaunchArg, VirtualViewLayoutLaunch},
397            layout::as_dyn::{IntoDyn, IntoDyn2Layout, IntoDynLayout},
398        },
399    };
400
401    use super::*;
402
403    #[allow(clippy::type_complexity)]
404    pub enum ViewArg<C: Coordinates> {
405        Array(BufferArg, VirtualViewLayoutLaunch<C, Coords1d, BufferArg>),
406        Tensor(TensorArg, VirtualViewLayoutLaunch<C, Coords1d, TensorArg>),
407        TensorMapTiled(
408            TensorMapArg<Tiled>,
409            VirtualViewLayoutLaunch<C, Sequence<i32>, TensorMapArg<Tiled>>,
410        ),
411        TensorMapIm2col(
412            TensorMapArg<Im2col>,
413            VirtualViewLayoutLaunch<C, (Sequence<i32>, Sequence<i32>), TensorMapArg<Im2col>>,
414        ),
415        Quantized {
416            values: Box<ViewArg<C>>,
417            scales: ScaleBindings<C>,
418            scheme: QuantScheme,
419        },
420    }
421
422    /// The scale bindings of a quantized view, one per scheme level.
423    ///
424    /// Only the block scales are addressed per position, so they bind as a view; the per-tensor
425    /// scale of a two-level scheme covers the whole tensor and binds as a buffer holding its one
426    /// scale in the first element, read once per kernel as f32.
427    pub struct ScaleBindings<C: Coordinates> {
428        pub(crate) inner: Box<ViewArg<C>>,
429        pub(crate) global_scale: Option<BufferArg>,
430        /// A lookup scheme's `2^bits`-entry table, present exactly under
431        /// [`QuantMode::Lookup`](cubecl_common::quant::scheme::QuantMode). Not a scale level —
432        /// [`len`](Self::len) never counts it — but it rides here because it is the same kind of
433        /// thing: a side binding the dequantizing read folds in.
434        pub(crate) table: Option<BufferArg>,
435    }
436
437    impl<C: Coordinates> ScaleBindings<C> {
438        /// The binding of a one-level scheme's scales.
439        pub fn one(scales: ViewArg<C>) -> Self {
440            Self {
441                inner: Box::new(scales),
442                global_scale: None,
443                table: None,
444            }
445        }
446
447        /// The bindings of a two-level scheme: the block scales and the per-tensor scale they are
448        /// normalized against.
449        pub fn two(scales: ViewArg<C>, global_scale: BufferArg) -> Self {
450            Self {
451                inner: Box::new(scales),
452                global_scale: Some(global_scale),
453                table: None,
454            }
455        }
456
457        /// The bindings of a one-level lookup scheme
458        /// ([`QuantMode::Lookup`](cubecl_common::quant::scheme::QuantMode)): the scales and the
459        /// table each stored field indexes, so a read reconstructs `table[field] * scale`.
460        ///
461        /// `table` must hold `2^bits` f32 entries — registration checks its length against the
462        /// scheme, and the unpack's mask bounds every index to that range.
463        pub fn lookup(scales: ViewArg<C>, table: BufferArg) -> Self {
464            Self {
465                inner: Box::new(scales),
466                global_scale: None,
467                table: Some(table),
468            }
469        }
470
471        /// The number of bound scale levels, matched against the scheme's at construction: the
472        /// inner binding plus the global scale when bound. The lookup table is not a level and is
473        /// never counted; its presence is checked against the scheme's mode instead
474        /// ([`quant::check_table_bindings`]).
475        #[allow(clippy::len_without_is_empty, reason = "never empty by construction")]
476        pub fn len(&self) -> usize {
477            1 + self.global_scale.iter().count()
478        }
479    }
480
481    /// [`ScaleBindings`] between registration and expansion.
482    #[derive(Clone)]
483    pub struct ScaleBindingsCompilationArg<C: Coordinates> {
484        pub(crate) inner: Box<ViewCompilationArg<C>>,
485        pub(crate) global_scale: Option<BufferCompilationArg>,
486        pub(crate) table: Option<BufferCompilationArg>,
487    }
488
489    impl<C: Coordinates> ScaleBindingsCompilationArg<C> {
490        /// See [`ScaleBindings::len`].
491        #[allow(clippy::len_without_is_empty, reason = "never empty by construction")]
492        pub fn len(&self) -> usize {
493            1 + self.global_scale.iter().count()
494        }
495    }
496
497    impl<C: Coordinates> Eq for ScaleBindingsCompilationArg<C> {}
498    impl<C: Coordinates> PartialEq for ScaleBindingsCompilationArg<C> {
499        fn eq(&self, other: &Self) -> bool {
500            self.inner == other.inner
501                && self.global_scale == other.global_scale
502                && self.table == other.table
503        }
504    }
505    impl<C: Coordinates> core::hash::Hash for ScaleBindingsCompilationArg<C> {
506        fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
507            self.inner.hash(state);
508            self.global_scale.hash(state);
509            self.table.hash(state);
510        }
511    }
512    impl<C: Coordinates> core::fmt::Debug for ScaleBindingsCompilationArg<C> {
513        fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
514            f.debug_struct("ScaleBindings")
515                .field("inner", &self.inner)
516                .field("global_scale", &self.global_scale)
517                .field("table", &self.table)
518                .finish()
519        }
520    }
521
522    impl<C: Coordinates> ViewArg<C> {
523        pub fn new_array<
524            L: Layout<Coordinates = C, SourceCoordinates = Coords1d> + ViewLayoutLaunchArg,
525        >(
526            buffer: BufferArg,
527            layout: L::RuntimeArg,
528        ) -> Self {
529            let layout = VirtualViewLayoutLaunch::new::<L>(layout);
530            ViewArg::Array(buffer, layout)
531        }
532
533        pub fn new_tensor<
534            L: Layout<Coordinates = C, SourceCoordinates = Coords1d> + ViewLayoutLaunchArg,
535        >(
536            buffer: TensorArg,
537            layout: L::RuntimeArg,
538        ) -> Self {
539            let layout = VirtualViewLayoutLaunch::new::<L>(layout);
540            ViewArg::Tensor(buffer, layout)
541        }
542
543        pub fn new_tensor_map_tiled<
544            L: Layout<Coordinates = C, SourceCoordinates: IntoDyn> + ViewLayoutLaunchArg,
545        >(
546            buffer: TensorMapArg<Tiled>,
547            layout: L::RuntimeArg,
548        ) -> ViewArg<C> {
549            let layout = VirtualViewLayoutLaunch::new::<IntoDynLayout<L>>(layout);
550            ViewArg::TensorMapTiled(buffer, layout)
551        }
552
553        pub fn new_tensor_map_im2col<
554            L: Layout<Coordinates = C, SourceCoordinates = (P, O)> + ViewLayoutLaunchArg,
555            P: IntoDyn,
556            O: IntoDyn,
557        >(
558            buffer: TensorMapArg<Im2col>,
559            layout: L::RuntimeArg,
560        ) -> ViewArg<C> {
561            let layout = VirtualViewLayoutLaunch::new::<IntoDyn2Layout<L, P, O>>(layout);
562            ViewArg::TensorMapIm2col(buffer, layout)
563        }
564
565        /// Create a new view arg that dequantizes on read, against one scale binding per scheme
566        /// level. The inner scales layout should take values indices and map them to the
567        /// corresponding scale.
568        ///
569        /// Panics when the bindings and the scheme's levels disagree in count, since a missing
570        /// level would be dropped from the reconstruction, and for an global level this reader
571        /// cannot serve. See [`quant::check_scale_bindings`].
572        pub fn new_quantized(values: Self, scales: ScaleBindings<C>, scheme: QuantScheme) -> Self {
573            quant::check_scale_bindings(&scheme, scales.len());
574            quant::check_table_bindings(&scheme, scales.table.is_some());
575            Self::Quantized {
576                values: Box::new(values),
577                scales,
578                scheme,
579            }
580        }
581    }
582    #[derive(Clone)]
583    pub enum ViewCompilationArg<C: Coordinates> {
584        Array {
585            buffer: BufferCompilationArg,
586            layout: VirtualViewLayoutCompilationArg<C, Coords1d>,
587        },
588        TensorMapTiled {
589            buffer: (),
590            layout: VirtualViewLayoutCompilationArg<C, Sequence<i32>>,
591        },
592        TensorMapIm2col {
593            buffer: (),
594            layout: VirtualViewLayoutCompilationArg<C, (Sequence<i32>, Sequence<i32>)>,
595        },
596        Quantized {
597            values: Box<ViewCompilationArg<C>>,
598            scales: ScaleBindingsCompilationArg<C>,
599            scheme: QuantScheme,
600        },
601    }
602
603    impl<C: Coordinates> Eq for ViewCompilationArg<C> {}
604    impl<C: Coordinates> PartialEq for ViewCompilationArg<C> {
605        fn eq(&self, other: &Self) -> bool {
606            match (self, other) {
607                (
608                    ViewCompilationArg::Array { buffer, layout },
609                    ViewCompilationArg::Array {
610                        buffer: buffer_other,
611                        layout: layout_other,
612                    },
613                ) => buffer == buffer_other && layout == layout_other,
614                (
615                    ViewCompilationArg::TensorMapTiled { buffer, layout },
616                    ViewCompilationArg::TensorMapTiled {
617                        buffer: buffer_other,
618                        layout: layout_other,
619                    },
620                ) => buffer == buffer_other && layout == layout_other,
621                (
622                    ViewCompilationArg::TensorMapIm2col { buffer, layout },
623                    ViewCompilationArg::TensorMapIm2col {
624                        buffer: buffer_other,
625                        layout: layout_other,
626                    },
627                ) => buffer == buffer_other && layout == layout_other,
628                (
629                    ViewCompilationArg::Quantized {
630                        values,
631                        scales,
632                        scheme,
633                    },
634                    ViewCompilationArg::Quantized {
635                        values: values_other,
636                        scales: scales_other,
637                        scheme: scheme_other,
638                    },
639                ) => values == values_other && scales == scales_other && scheme == scheme_other,
640                _ => false,
641            }
642        }
643    }
644    impl<C: Coordinates> core::hash::Hash for ViewCompilationArg<C> {
645        fn hash<H: core::hash::Hasher>(&self, ra_expand_state: &mut H) {
646            match self {
647                ViewCompilationArg::Array { buffer, layout } => {
648                    buffer.hash(ra_expand_state);
649                    layout.hash(ra_expand_state);
650                }
651                ViewCompilationArg::TensorMapTiled { buffer, layout } => {
652                    buffer.hash(ra_expand_state);
653                    layout.hash(ra_expand_state);
654                }
655                ViewCompilationArg::TensorMapIm2col { buffer, layout } => {
656                    buffer.hash(ra_expand_state);
657                    layout.hash(ra_expand_state);
658                }
659                ViewCompilationArg::Quantized {
660                    values,
661                    scales,
662                    scheme,
663                } => {
664                    values.hash(ra_expand_state);
665                    scales.hash(ra_expand_state);
666                    scheme.hash(ra_expand_state);
667                }
668            }
669        }
670    }
671    impl<C: Coordinates> core::fmt::Debug for ViewCompilationArg<C> {
672        fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
673            match self {
674                ViewCompilationArg::Array { buffer, layout } => f
675                    .debug_struct("ArrayView")
676                    .field("buffer", &buffer)
677                    .field("layout", &layout)
678                    .finish(),
679                ViewCompilationArg::TensorMapTiled { buffer, layout } => f
680                    .debug_struct("TensorMapTiledView")
681                    .field("buffer", &buffer)
682                    .field("layout", &layout)
683                    .finish(),
684                ViewCompilationArg::TensorMapIm2col { buffer, layout } => f
685                    .debug_struct("TensorMapIm2colView")
686                    .field("buffer", &buffer)
687                    .field("layout", &layout)
688                    .finish(),
689                ViewCompilationArg::Quantized {
690                    values,
691                    scales,
692                    scheme,
693                } => f
694                    .debug_struct("QuantizedView")
695                    .field("values", &values)
696                    .field("scales", &scales)
697                    .field("scheme", &scheme)
698                    .finish(),
699            }
700        }
701    }
702
703    impl<E: CubePrimitive, C: Coordinates + 'static> LaunchArg for View<'static, E, C> {
704        type RuntimeArg = ViewArg<C>;
705        type CompilationArg = ViewCompilationArg<C>;
706
707        fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
708            let ty = launcher.with_scope(|scope| {
709                let storage = E::Scalar::elem_type(scope);
710                let vector_size = E::__expand_vector_size(scope);
711                Type::new(storage).with_vector_size(vector_size)
712            });
713            match arg {
714                ViewArg::Array(buffer, layout) => ViewCompilationArg::Array {
715                    layout: layout.register(&buffer, ty, launcher),
716                    buffer: <[E] as LaunchArg>::register(buffer, launcher),
717                },
718                ViewArg::Tensor(buffer, layout) => ViewCompilationArg::Array {
719                    layout: layout.register(&buffer, ty, launcher),
720                    buffer: <[E] as LaunchArg>::register(buffer.into_buffer_arg(), launcher),
721                },
722                ViewArg::TensorMapTiled(buffer, layout) => ViewCompilationArg::TensorMapTiled {
723                    layout: layout.register(&buffer, ty, launcher),
724                    buffer: <TensorMap<E, Tiled> as LaunchArg>::register(buffer, launcher),
725                },
726                ViewArg::TensorMapIm2col(buffer, layout) => ViewCompilationArg::TensorMapIm2col {
727                    layout: layout.register(&buffer, ty, launcher),
728                    buffer: <TensorMap<E, Im2col> as LaunchArg>::register(buffer, launcher),
729                },
730                ViewArg::Quantized {
731                    values,
732                    scales,
733                    scheme,
734                } => {
735                    let register = RegisterDynamic {
736                        values: *values,
737                        scales,
738                        scheme,
739                        launcher,
740                        _ty: PhantomData::<E>,
741                    };
742                    run_with_quant_type(register, scheme)
743                }
744            }
745        }
746        fn expand(
747            arg: &Self::CompilationArg,
748            builder: &mut KernelBuilder,
749        ) -> <Self as CubeType>::ExpandType {
750            let storage = E::Scalar::elem_type(builder);
751            let vector_size = E::__expand_vector_size(builder);
752            let ty = Type::new(storage).with_vector_size(vector_size);
753            match arg {
754                ViewCompilationArg::Array { buffer, layout } => {
755                    let layout = layout.expand(ty, builder);
756                    let buffer = <Box<[E]> as LaunchArg>::expand(buffer, builder);
757                    let view =
758                        VirtualViewMutExpand::<E, C, Coords1d, Box<[E]>>::new(buffer, layout);
759                    ViewExpand::new(&builder.scope, view)
760                }
761                ViewCompilationArg::TensorMapTiled { buffer, layout } => {
762                    let layout = layout.expand(ty, builder);
763                    let buffer = <TensorMap<E, Tiled> as LaunchArg>::expand(buffer, builder);
764                    let view =
765                        VirtualViewMutExpand::<E, C, Sequence<i32>, TensorMap<E, Tiled>>::new(
766                            buffer, layout,
767                        );
768                    ViewExpand::new(&builder.scope, view)
769                }
770                ViewCompilationArg::TensorMapIm2col { buffer, layout } => {
771                    let layout = layout.expand(ty, builder);
772                    let buffer = <TensorMap<E, Im2col> as LaunchArg>::expand(buffer, builder);
773                    let view = VirtualViewExpand::<
774                        E,
775                        C,
776                        (Sequence<i32>, Sequence<i32>),
777                        TensorMap<E, Im2col>,
778                    >::new(buffer, layout);
779                    ViewExpand::new(&builder.scope, view)
780                }
781                ViewCompilationArg::Quantized {
782                    values,
783                    scales,
784                    scheme,
785                } => quant::view::expand_dynamic(values, scales, *scheme, builder),
786            }
787        }
788    }
789
790    impl<E: CubePrimitive, C: Coordinates + 'static> LaunchArg for ViewMut<'static, E, C> {
791        type RuntimeArg = ViewArg<C>;
792        type CompilationArg = ViewCompilationArg<C>;
793
794        fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
795            <View<'static, E, C> as LaunchArg>::register(arg, launcher)
796        }
797
798        fn expand(
799            arg: &Self::CompilationArg,
800            builder: &mut KernelBuilder,
801        ) -> <Self as CubeType>::ExpandType {
802            let storage = E::Scalar::elem_type(builder);
803            let vector_size = E::__expand_vector_size(builder);
804            let ty = Type::new(storage).with_vector_size(vector_size);
805            match arg {
806                ViewCompilationArg::Array { buffer, layout } => {
807                    let layout = layout.expand(ty, builder);
808                    let buffer = <Box<[E]> as LaunchArg>::expand(buffer, builder);
809                    let view =
810                        VirtualViewMutExpand::<E, C, Coords1d, Box<[E]>>::new(buffer, layout);
811                    ViewMutExpand::new(&builder.scope, view)
812                }
813                ViewCompilationArg::TensorMapTiled { buffer, layout } => {
814                    let layout = layout.expand(ty, builder);
815                    let buffer = <TensorMap<E, Tiled> as LaunchArg>::expand(buffer, builder);
816                    let view =
817                        VirtualViewMutExpand::<E, C, Sequence<i32>, TensorMap<E, Tiled>>::new(
818                            buffer, layout,
819                        );
820                    ViewMutExpand::new(&builder.scope, view)
821                }
822                ViewCompilationArg::TensorMapIm2col { .. } => {
823                    unimplemented!("im2col not supported for writing")
824                }
825                ViewCompilationArg::Quantized { .. } => {
826                    unimplemented!("quantized views not supported for writing")
827                }
828            }
829        }
830    }
831}
832
833pub use dynamic::*;
cubecl_std/tensor/view/launch.rs

cubecl_std/tensor/view/
launch.rs
