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

cubecl_std/tensor/view/
launch.rs
