ruda_kernel/dsl/frontend/container/tensor/
tensormap.rs

1use alloc::vec;
2use core::marker::PhantomData;
3
4use crate::dsl::{prelude::*, unexpanded};
5use ruda_core::ir::{Type, VectorSize};
6use ruda::runtime::server::TensorMapMeta;
7use ruda_core::tensor::{Strides, metadata::Metadata, strides};
8use paste::paste;
9
10pub use ruda::runtime::tma::*;
11
12pub trait TensorMapKind: RudaType + Clone + Copy + Send + Sync + 'static {
13    type Args: Clone;
14
15    fn as_format(args: Self::Args) -> TensorMapFormat;
16}
17
18/// Regular tiled tensor map
19#[derive(RudaType, RudaLaunch, Clone, Copy)]
20pub struct Tiled {}
21/// Im2col indexing. Loads a "column" (not the same column as im2col) of pixels into shared
22/// memory, with a certain offset (kernel position). The corners are the bounds to load pixels
23/// from *at offset 0*, so the top left corner of the kernel. The offset is added to the
24/// corner offsets, so a `(-1, -1)` corner will stop the bounding box at `(1, 1)` for kernel
25/// offset `(2, 2)`.
26#[derive(RudaType, RudaLaunch, Clone, Copy)]
27pub struct Im2col;
28/// 1D im2col, not properly supported yet
29#[derive(RudaType, RudaLaunch, Clone, Copy)]
30pub struct Im2colWide;
31
32impl TensorMapKind for Tiled {
33    type Args = TiledArgs;
34
35    fn as_format(args: Self::Args) -> TensorMapFormat {
36        TensorMapFormat::Tiled(args)
37    }
38}
39
40impl TensorMapKind for Im2col {
41    type Args = Im2colArgs;
42
43    fn as_format(args: Self::Args) -> TensorMapFormat {
44        TensorMapFormat::Im2col(args)
45    }
46}
47
48impl TensorMapKind for Im2colWide {
49    type Args = Im2colWideArgs;
50
51    fn as_format(args: Self::Args) -> TensorMapFormat {
52        TensorMapFormat::Im2colWide(args)
53    }
54}
55
56/// Grid constant tensor map, currently only maps to CUDA tensormap. May be interleaved or swizzled,
57/// but last dimension must be contiguous (since strides don't include the last dimension).
58///
59/// The tensormap is treated as an opaque type at runtime.
60///
61pub struct TensorMapArg<R: Runtime, K: TensorMapKind> {
62    pub tensor: TensorArg<R>,
63    pub metadata: TensorMapMeta,
64    pub _kind: PhantomData<K>,
65}
66
67impl<R: Runtime, K: TensorMapKind> TensorMapArg<R, K> {
68    pub fn new(args: K::Args, tensor: TensorArg<R>, ty: impl Into<Type>) -> Self {
69        let ty = ty.into();
70        let TensorArg::Handle { handle, .. } = &tensor else {
71            panic!("Can't use alias for TensorMap")
72        };
73        let rank = handle.shape.len();
74        Self {
75            metadata: TensorMapMeta {
76                format: K::as_format(args),
77                metadata: Metadata::new(handle.shape.clone(), handle.strides.clone()),
78                elem_stride: strides![1; rank],
79                interleave: TensorMapInterleave::None,
80                swizzle: TensorMapSwizzle::None,
81                prefetch: TensorMapPrefetch::None,
82                oob_fill: OobFill::Zero,
83                storage_ty: ty.storage_type(),
84            },
85            tensor,
86            _kind: PhantomData,
87        }
88    }
89
90    pub fn with_elem_stride(mut self, elem_stride: Strides) -> Self {
91        self.metadata.elem_stride = elem_stride;
92        self
93    }
94
95    pub fn with_interleave(mut self, interleave: TensorMapInterleave) -> Self {
96        self.metadata.interleave = interleave;
97        self
98    }
99
100    pub fn with_swizzle(mut self, swizzle: TensorMapSwizzle) -> Self {
101        self.metadata.swizzle = swizzle;
102        self
103    }
104
105    pub fn with_prefetch(mut self, prefetch: TensorMapPrefetch) -> Self {
106        self.metadata.prefetch = prefetch;
107        self
108    }
109
110    pub fn with_nan_fill(mut self) -> Self {
111        self.metadata.oob_fill = OobFill::NaN;
112        self
113    }
114}
115
116/// A CUDA `CUtensorMap` object. Represents a tensor encoded with a lot of metadata, and is an
117/// opaque packed object at runtime. Does not support retrieving any shapes or strides, nor does
118/// it give access to the pointer. So these need to be passed separately in an aliased `Tensor` if needed.
119///
120/// Also see [`ruda::runtime::tma`].
121#[derive(Clone)]
122pub struct TensorMap<E: RudaPrimitive, K: TensorMapKind> {
123    _ty: PhantomData<E>,
124    _kind: PhantomData<K>,
125}
126
127impl<E: RudaPrimitive, K: TensorMapKind> Copy for TensorMap<E, K> {}
128
129impl<E: RudaPrimitive, K: TensorMapKind> TensorMap<E, K> {}
130
131impl<E: RudaPrimitive, K: TensorMapKind> IntoMut for NativeExpand<TensorMap<E, K>> {
132    fn into_mut(self, _scope: &mut Scope) -> Self {
133        self
134    }
135}
136
137impl<E: RudaPrimitive, K: TensorMapKind> RudaType for TensorMap<E, K> {
138    type ExpandType = NativeExpand<TensorMap<E, K>>;
139}
140
141impl<E: RudaPrimitive, K: TensorMapKind> RudaType for *const TensorMap<E, K> {
142    type ExpandType = NativeExpand<TensorMap<E, K>>;
143}
144
145impl<E: RudaPrimitive, K: TensorMapKind> RudaType for *mut TensorMap<E, K> {
146    type ExpandType = NativeExpand<TensorMap<E, K>>;
147}
148
149impl<E: RudaPrimitive, K: TensorMapKind> Vectorized for TensorMap<E, K> {}
150impl<E: RudaPrimitive, K: TensorMapKind> VectorizedExpand for NativeExpand<TensorMap<E, K>> {
151    fn vector_size(&self) -> VectorSize {
152        1
153    }
154}
155
156impl<E: RudaPrimitive, K: TensorMapKind> LaunchArg for TensorMap<E, K> {
157    type RuntimeArg<R: Runtime> = TensorMapArg<R, K>;
158    type CompilationArg = ();
159
160    fn register<R: Runtime>(
161        arg: Self::RuntimeArg<R>,
162        launcher: &mut KernelLauncher<R>,
163    ) -> Self::CompilationArg {
164        let ty = launcher.with_scope(|scope| E::as_type(scope));
165        launcher.register_tensor_map(arg, ty);
166    }
167
168    fn expand(
169        _arg: &Self::CompilationArg,
170        builder: &mut KernelBuilder,
171    ) -> NativeExpand<TensorMap<E, K>> {
172        let tensor = builder.input_tensor_map(E::as_type(&builder.scope));
173        tensor.into()
174    }
175    fn expand_output(
176        _arg: &Self::CompilationArg,
177        builder: &mut KernelBuilder,
178    ) -> NativeExpand<TensorMap<E, K>> {
179        let tensor = builder.output_tensor_map(E::as_type(&builder.scope));
180        tensor.into()
181    }
182}
183
184/// Commit an async tensor operation. Not sure how this works, poor docs. But you need to call it
185/// after a write, but not after reads.
186pub fn tma_group_commit() {
187    unexpanded!()
188}
189
190pub mod tma_group_commit {
191    use ruda_core::ir::TmaOps;
192
193    use super::*;
194
195    pub fn expand(scope: &mut Scope) {
196        scope.register(TmaOps::CommitGroup)
197    }
198}
199
200/// Wait until at most `max_pending` TMA copy operations are in flight.
201pub fn tma_group_wait(_max_pending: u32) {
202    unexpanded!()
203}
204
205pub mod tma_group_wait {
206    use ruda_core::ir::TmaOps;
207
208    use super::*;
209
210    pub fn expand(scope: &mut Scope, max_pending: u32) {
211        scope.register(TmaOps::WaitGroup { max_pending })
212    }
213}
214
215/// Wait TMA copy operations have finished reading from shared memory, with at most `max_pending`
216/// operations being unfinished.
217///
218/// # Example
219///
220/// I believe you may use `max_pending` like this.
221///
222/// ```ignore
223/// copy_data(smem1);
224/// copy_data(smem2);
225/// copy_data(smem3);
226/// copy_data(smem4);
227/// tma_wait_read(2);
228/// // reuse smem1 & smem2 while 3 and 4 are still pending
229/// ```
230pub fn tma_group_wait_read(_max_pending: u32) {
231    unexpanded!()
232}
233
234pub mod tma_group_wait_read {
235    use ruda_core::ir::TmaOps;
236
237    use super::*;
238
239    pub fn expand(scope: &mut Scope, max_pending: u32) {
240        scope.register(TmaOps::WaitGroupRead { max_pending })
241    }
242}
243
244macro_rules! tma_store {
245    ($dim: literal, $($arg: expr),*) => {
246        paste! {
247            /// Copy a tile from a shared memory `src` to a global memory `dst`, with the provided
248            /// offsets. Should be combined with ``memcpy_async_tensor_commit`` and
249            /// ``memcpy_async_tensor_wait_read``.
250            #[allow(unused)]
251            pub fn [<tma_store_ $dim d>]<T: RudaPrimitive, T2: RudaPrimitive<Scalar = T::Scalar>>(
252                src: &Slice<T2>,
253                dst: &mut TensorMap<T, Tiled>,
254                $($arg: i32),*
255            ) {
256                unexpanded!()
257            }
258
259            pub mod [<tma_store_ $dim d>] {
260                use ruda_core::ir::{Instruction, TmaOps};
261
262                use super::*;
263
264                #[allow(clippy::too_many_arguments)]
265                pub fn expand<T: RudaPrimitive, T2: RudaPrimitive<Scalar = T::Scalar>>(
266                    scope: &mut Scope,
267                    src: SliceExpand<T2, ReadOnly>,
268                    dst: NativeExpand<TensorMap<T, Tiled>>,
269                    $($arg: NativeExpand<i32>),*
270                ) {
271                    let (source, source_offset) = src.__to_raw_parts();
272                    let dst = *dst.expand;
273                    let coordinates = vec![$(*$arg.expand),*];
274                    scope.register(Instruction::new(
275                        TmaOps::TmaStore {
276                            source,
277                            coordinates,
278                            offset_source: source_offset,
279                        },
280                        dst,
281                    ))
282                }
283            }
284        }
285    };
286}
287
288tma_store!(1, x);
289tma_store!(2, y, x);
290tma_store!(3, z, y, x);
291tma_store!(4, w, z, y, x);
292tma_store!(5, v, w, z, y, x);
293
294/// Module that contains the implementation details of the metadata functions.
295mod metadata {
296    use ruda_core::ir::{ManagedVariable, Metadata, VariableKind};
297
298    use super::*;
299    use crate::dsl::{
300        ir::{Arithmetic, BinaryOperator, Instruction},
301        prelude::Array,
302    };
303
304    impl<T: Scalar, K: TensorMapKind> TensorMap<T, K> {
305        /// Get a reference to the underlying buffer for the tensor map.
306        pub fn buffer<N: Size>(&self) -> Tensor<Vector<T, N>> {
307            unexpanded!()
308        }
309
310        /// Obtain the stride of input at dimension dim
311        pub fn stride(&self, _dim: usize) -> usize {
312            unexpanded!()
313        }
314
315        /// Obtain the shape of input at dimension dim
316        pub fn shape(&self, _dim: usize) -> usize {
317            unexpanded!()
318        }
319
320        /// Obtain the coordinate corresponding to the given `index` of the tensor at dimension `dim`.
321        ///
322        /// A coordinate is a list of indices corresponding to the multi-dimensional position of an element in the tensor.
323        /// The `dim` element in a coordinate is the position along the `dim` dimension of the tensor.
324        pub fn coordinate(&self, _index: usize, _dim: usize) -> usize {
325            unexpanded!()
326        }
327
328        /// The number of vectorized elements in the tensor.
329        ///
330        /// # Warning
331        ///
332        /// The length will be affected by the vectorization factor. To obtain the number of elements,
333        /// you should multiply the length by the vectorization factor.
334        #[allow(clippy::len_without_is_empty)]
335        pub fn len(&self) -> usize {
336            unexpanded!()
337        }
338
339        /// The length of the buffer representing the tensor in terms of vectorized elements.
340        ///
341        /// # Warning
342        ///
343        /// The buffer length will be affected by the vectorization factor. To obtain the number of
344        /// elements, you should multiply the length by the vectorization factor.
345        #[allow(clippy::len_without_is_empty)]
346        pub fn buffer_len(&self) -> usize {
347            unexpanded!()
348        }
349
350        /// Returns the rank of the tensor.
351        pub fn rank(&self) -> usize {
352            unexpanded!()
353        }
354
355        /// Downcast the tensormap to the given type and panic if the type isn't the same.
356        ///
357        /// This function should only be used to satisfy the Rust type system, when two generic
358        /// types are supposed to be the same.
359        pub fn downcast<E: RudaPrimitive>(&self) -> TensorMap<E, K> {
360            unexpanded!()
361        }
362
363        // Expand function of [buffer](TensorMap::buffer).
364        pub fn __expand_buffer(
365            scope: &mut Scope,
366            expand: NativeExpand<TensorMap<T, K>>,
367        ) -> NativeExpand<Tensor<T>> {
368            expand.__expand_buffer_method(scope)
369        }
370
371        // Expand function of [stride](TensorMap::stride).
372        pub fn __expand_stride(
373            scope: &mut Scope,
374            expand: NativeExpand<TensorMap<T, K>>,
375            dim: NativeExpand<usize>,
376        ) -> NativeExpand<usize> {
377            expand.__expand_stride_method(scope, dim)
378        }
379
380        // Expand function of [shape](TensorMap::shape).
381        pub fn __expand_shape(
382            scope: &mut Scope,
383            expand: NativeExpand<TensorMap<T, K>>,
384            dim: NativeExpand<usize>,
385        ) -> NativeExpand<usize> {
386            expand.__expand_shape_method(scope, dim)
387        }
388
389        // Expand function of [coordinate](TensorMap::coordinate).
390        pub fn __expand_coordinate(
391            scope: &mut Scope,
392            expand: NativeExpand<TensorMap<T, K>>,
393            index: NativeExpand<usize>,
394            dim: NativeExpand<usize>,
395        ) -> NativeExpand<usize> {
396            expand.__expand_coordinate_method(scope, index, dim)
397        }
398
399        // Expand function of [len](TensorMap::len).
400        pub fn __expand_len(
401            scope: &mut Scope,
402            expand: NativeExpand<TensorMap<T, K>>,
403        ) -> NativeExpand<usize> {
404            expand.__expand_len_method(scope)
405        }
406
407        // Expand function of [buffer_len](TensorMap::buffer_len).
408        pub fn __expand_buffer_len(
409            scope: &mut Scope,
410            expand: NativeExpand<TensorMap<T, K>>,
411        ) -> NativeExpand<usize> {
412            expand.__expand_buffer_len_method(scope)
413        }
414
415        // Expand function of [rank](TensorMap::rank).
416        pub fn __expand_rank(
417            scope: &mut Scope,
418            expand: NativeExpand<TensorMap<T, K>>,
419        ) -> NativeExpand<usize> {
420            expand.__expand_rank_method(scope)
421        }
422    }
423
424    impl<T: RudaPrimitive, K: TensorMapKind> NativeExpand<TensorMap<T, K>> {
425        // Expand method of [buffer](TensorMap::buffer).
426        pub fn __expand_buffer_method(self, scope: &mut Scope) -> NativeExpand<Tensor<T>> {
427            let tensor = match self.expand.kind {
428                VariableKind::TensorMapInput(id) => scope.input(id, self.expand.ty),
429                VariableKind::TensorMapOutput(id) => scope.output(id, self.expand.ty),
430                _ => unreachable!(),
431            };
432            tensor.into()
433        }
434
435        // Expand method of [stride](Tensor::stride).
436        pub fn __expand_stride_method(
437            self,
438            scope: &mut Scope,
439            dim: NativeExpand<usize>,
440        ) -> NativeExpand<usize> {
441            let dim: ManagedVariable = dim.into();
442            let out = scope.create_local(usize::as_type(scope));
443            scope.register(Instruction::new(
444                Metadata::Stride {
445                    dim: *dim,
446                    var: self.expand.into(),
447                },
448                out.clone().into(),
449            ));
450            out.into()
451        }
452
453        // Expand method of [shape](Tensor::shape).
454        pub fn __expand_shape_method(
455            self,
456            scope: &mut Scope,
457            dim: NativeExpand<usize>,
458        ) -> NativeExpand<usize> {
459            let dim: ManagedVariable = dim.into();
460            let out = scope.create_local(usize::as_type(scope));
461            scope.register(Instruction::new(
462                Metadata::Shape {
463                    dim: *dim,
464                    var: self.expand.into(),
465                },
466                out.clone().into(),
467            ));
468            out.into()
469        }
470
471        // Expand method of [coordinate](Tensor::coordinate).
472        pub fn __expand_coordinate_method(
473            self,
474            scope: &mut Scope,
475            index: NativeExpand<usize>,
476            dim: NativeExpand<usize>,
477        ) -> NativeExpand<usize> {
478            let index: ManagedVariable = index.into();
479            let stride = self.clone().__expand_stride_method(scope, dim.clone());
480            let shape = self.clone().__expand_shape_method(scope, dim.clone());
481
482            // Compute `num_strides = index / stride`.
483            let num_strides = scope.create_local(usize::as_type(scope));
484            scope.register(Instruction::new(
485                Arithmetic::Div(BinaryOperator {
486                    lhs: *index,
487                    rhs: stride.expand.into(),
488                }),
489                num_strides.clone().into(),
490            ));
491
492            // Compute `coordinate = num_strides % shape `.
493            let coordinate = scope.create_local(usize::as_type(scope));
494            scope.register(Instruction::new(
495                Arithmetic::Modulo(BinaryOperator {
496                    lhs: *num_strides,
497                    rhs: shape.expand.into(),
498                }),
499                coordinate.clone().into(),
500            ));
501
502            coordinate.into()
503        }
504
505        // Expand method of [len](Tensor::len).
506        pub fn __expand_len_method(self, scope: &mut Scope) -> NativeExpand<usize> {
507            let elem: NativeExpand<Array<u32>> = self.expand.into();
508            elem.__expand_len_method(scope)
509        }
510
511        // Expand method of [buffer_len](Tensor::buffer_len).
512        pub fn __expand_buffer_len_method(self, scope: &mut Scope) -> NativeExpand<usize> {
513            let elem: NativeExpand<Array<u32>> = self.expand.into();
514            elem.__expand_buffer_len_method(scope)
515        }
516
517        // Expand method of [rank](Tensor::rank).
518        pub fn __expand_rank_method(self, scope: &mut Scope) -> NativeExpand<usize> {
519            let out = scope.create_local(usize::as_type(scope));
520            scope.register(Instruction::new(Metadata::Rank { var: *self.expand }, *out));
521            out.into()
522        }
523
524        /// Expand method of [`TensorMap::downcast`].
525        pub fn __expand_downcast_method<E: RudaPrimitive>(
526            self,
527            scope: &mut Scope,
528        ) -> NativeExpand<TensorMap<E, K>> {
529            if T::as_type(scope) != E::as_type(scope) && !is_tf32::<E, T>(scope) {
530                panic!("Downcast should only be used to satisfy the Rust type system.")
531            }
532
533            self.expand.into()
534        }
535    }
536}
ruda_kernel/dsl/frontend/container/tensor/tensormap.rs

ruda_kernel/dsl/frontend/container/tensor/
tensormap.rs
