cubecl_runtime/
tma.rs

1use alloc::vec::Vec;
2#[cfg(any(target_os = "windows", target_os = "linux", target_os = "macos"))]
3use serde::{Deserialize, Serialize};
4
5/// Format of [`TensorMap`]
6#[cfg_attr(
7    any(target_os = "windows", target_os = "linux", target_os = "macos"),
8    derive(Serialize, Deserialize)
9)]
10#[derive(Hash, PartialEq, Eq, Clone, Debug)]
11pub enum TensorMapFormat {
12    /// Simple tiling
13    Tiled {
14        /// Tile size that's loaded from memory in each copy operation. Must have `rank` elements.
15        /// In matmul, for example, this might be `batch x m x k`, or whatever the stage size is.
16        /// If a dimension isn't present in the tile, it should just be set to `1`.
17        ///
18        /// For CUDA, this must be a power of two and `<= 256` on each dimension.
19        tile_size: Vec<u32>,
20    },
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    Im2col {
27        /// Pixel box lower corner. This is the logical upper left corner in the input tensor,
28        /// when offset is 0. The length of this value should equal the *spatial* dimensions of
29        /// the input tensor (i.e. `h, w` for an NHWC tensor). Should normally be set to `-padding`.
30        pixel_box_lower_corner: Vec<i32>,
31        /// Pixel box top corner. This is the logical lower right corner in the input tensor,
32        /// when offset is 0. The length of this value should equal the *spatial* dimensions of
33        /// the input tensor (i.e. `h, w` for an NHWC tensor). Should normally be set to
34        /// `padding - kernel_size - 1` (where `kernel_size` accounts for dilation). This is not
35        /// equal to padding, it's equal to the bounding box for the *top left corner of the kernel*.
36        pixel_box_upper_corner: Vec<i32>,
37        /// Channels to load per pixel, should be a multiple or divisor of the matmul tile size.
38        /// This is not the total number of channels in the tensor, but only the number loaded in
39        /// each load. Must be <= 256 and aligned to 16 bytes.
40        channels_per_pixel: u32,
41        /// Pixels per column, equivalent to the `m`/`n` dimension of each tile in the matrix
42        /// multiplication. i.e. `NHW` for a 4D tensor.
43        /// Must be <= 256 and aligned to 16 bytes
44        pixels_per_column: u32,
45    },
46    /// Wide im2col
47    Im2colWide {
48        /// Pixel box lower corner width. TODO: How does this work?
49        pixel_box_lower_corner_width: i32,
50        /// Pixel box upper corner width. TODO: How does this work?
51        pixel_box_upper_corner_width: i32,
52        /// Channels per pixel
53        channels_per_pixel: u32,
54        /// Pixels per column
55        pixels_per_column: u32,
56    },
57}
58
59/// Interleave setting for [`TensorMap`]
60#[cfg_attr(
61    any(target_os = "windows", target_os = "linux", target_os = "macos"),
62    derive(Serialize, Deserialize)
63)]
64#[derive(Default, Hash, PartialEq, Eq, Clone, Debug, Copy)]
65pub enum TensorMapInterleave {
66    /// No interleaving
67    #[default]
68    None,
69    /// Interleaved with 16 bytes chunks in the last dim.
70    /// i.e. NC/8HWC8 with f16
71    B16,
72    /// Interleaved with 32 bytes chunks in the last dim.
73    /// i.e. NC/16HWC16 with f16
74    B32,
75}
76
77/// Data are organized in a specific order in global memory; however, this may not match the order
78/// in which the application accesses data in shared memory. This difference in data organization
79/// may cause bank conflicts when shared memory is accessed. In order to avoid this problem, data
80/// can be loaded to shared memory with shuffling across shared memory banks. When interleave is
81/// [`TensorMapInterleave::B32`], swizzle must be [`TensorMapSwizzle::B32`].
82/// Other interleave modes can have any swizzling pattern.
83#[cfg_attr(
84    any(target_os = "windows", target_os = "linux", target_os = "macos"),
85    derive(Serialize, Deserialize)
86)]
87#[derive(Default, Hash, PartialEq, Eq, Clone, Debug, Copy)]
88pub enum TensorMapSwizzle {
89    /// No swizzling
90    #[default]
91    None,
92    /// Swizzle 16B chunks within 32B span
93    B32,
94    /// Swizzle 16B chunks within 64B span
95    B64,
96    /// Swizzle 16B chunks within 128B span
97    B128,
98    /// Swizzle 32B chunks within 128B span
99    B128Atom32B,
100    /// Swizzle 32B chunks within 128B span, additionally swap lower 8B with upper 8B within each 16B for every alternate row
101    B128Atom32BFlip8B,
102    /// Swizzle 64B chunks within 128B span
103    B128Atom64B,
104}
105
106/// Additional prefetching to perform during load
107/// Specifies L2 fetch size which indicates the byte granularity at which L2 requests are filled from DRAM
108#[cfg_attr(
109    any(target_os = "windows", target_os = "linux", target_os = "macos"),
110    derive(Serialize, Deserialize)
111)]
112#[derive(Default, Hash, PartialEq, Eq, Clone, Debug, Copy)]
113pub enum TensorMapPrefetch {
114    /// No extra prefetch
115    #[default]
116    None,
117    /// Prefetch 64 bytes
118    B64,
119    /// Prefetch 128 bytes
120    B128,
121    /// Prefetch 256 bytes
122    B256,
123}
124
125/// What value to use when filling out of bounds values
126#[cfg_attr(
127    any(target_os = "windows", target_os = "linux", target_os = "macos"),
128    derive(Serialize, Deserialize)
129)]
130#[derive(Default, Hash, PartialEq, Eq, Clone, Debug, Copy)]
131pub enum OobFill {
132    /// Fill zeroes
133    #[default]
134    Zero,
135    /// Fill NaN
136    NaN,
137}
cubecl_runtime/tma.rs

cubecl_runtime/
tma.rs
