cubek_std/launch/
tma.rs

1//! Helpers for building TMA (Tensor Memory Accelerator) descriptors.
2
3use std::marker::PhantomData;
4
5use cubecl::prelude::*;
6use cubecl::server::TensorMapMeta;
7pub use cubecl::zspace::metadata::Metadata;
8use cubecl::zspace::{Shape, Strides, shape, strides};
9
10use crate::MatrixLayout;
11
12/// Defines the non-contiguous stride alignment in terms of powers of two
13pub fn stride_align_bits(strides: &[usize], layout: &MatrixLayout, dtype: &ElemType) -> u32 {
14    let exclude_dim = match layout {
15        MatrixLayout::RowMajor => strides.len() - 1,
16        MatrixLayout::ColMajor => strides.len() - 2,
17    };
18    strides
19        .iter()
20        .enumerate()
21        .filter(|(i, _)| *i != exclude_dim)
22        .map(|(_, it)| (*it * dtype.size_bits()) / 8)
23        .map(|it| it.trailing_zeros())
24        .min()
25        .unwrap_or(31)
26}
27
28/// CUDA's TMA loads f32 as tf32 internally; remap explicitly so the descriptor matches.
29pub fn remap_storage_for_tma(ty: ElemType) -> ElemType {
30    if ty == f32::elem_type_native() {
31        tf32::elem_type_native()
32    } else {
33        ty
34    }
35}
36
37/// TMA assumes the last stride is contiguous and discards it. For ColMajor inputs we therefore
38/// swap the inner two dims so the contiguous one ends up last. The tensor's own metadata stays
39/// in its original layout: only the TMA descriptor sees the transposed form.
40///
41/// `shape` and `strides` may have different ranks (the matmul builder constructs them
42/// transiently mismatched and aligns them afterwards). Each is swapped on its own inner pair.
43///
44/// Returns `true` if a swap occurred.
45pub fn transpose_inner_for_tma(
46    shape: &mut Shape,
47    strides: &mut Strides,
48    layout: MatrixLayout,
49) -> bool {
50    if matches!(layout, MatrixLayout::ColMajor) {
51        let s_rank = shape.num_dims();
52        let t_rank = strides.rank();
53        shape.swap(s_rank - 1, s_rank - 2);
54        strides.swap(t_rank - 1, t_rank - 2);
55        true
56    } else {
57        false
58    }
59}
60
61/// One matmul-style operand's tiled tensor map. Collapses the binding to the 3-D
62/// `(batch, rows, cols)` the descriptor expects, swaps the inner pair for a col-major
63/// operand (TMA discards the last stride) and keeps the batch stride outermost.
64/// `box_shape` is the box one bulk copy moves, in logical `(rows, cols)`; it rides the
65/// same swap. Returns the arg plus whether the swap occurred.
66pub fn tma_operand(
67    binding: TensorBinding,
68    batches: usize,
69    layout: MatrixLayout,
70    box_shape: (usize, usize),
71    storage_ty: ElemType,
72    swizzle: TensorMapSwizzle,
73) -> (TensorMapArg<Tiled>, bool) {
74    let rank = binding.shape.len();
75    let mut shape = shape![batches, binding.shape[rank - 2], binding.shape[rank - 1]];
76    let mut strides: Strides = if rank > 2 {
77        binding.strides[rank - 3..].into()
78    } else {
79        strides![binding.strides[0], binding.strides[1]]
80    };
81    let transposed = transpose_inner_for_tma(&mut shape, &mut strides, layout);
82    // Re-insert the batch stride after the (possible) swap so it stays outermost.
83    if strides.len() == 2 {
84        let stride = strides[0];
85        strides.insert(0, stride);
86    }
87
88    let (box_rows, box_cols) = box_shape;
89    let tile_size = match transposed {
90        true => shape![1, box_cols, box_rows],
91        false => shape![1, box_rows, box_cols],
92    };
93    let meta = tma_meta_tiled(
94        Metadata::new(shape, strides),
95        tile_size,
96        remap_storage_for_tma(storage_ty),
97        swizzle,
98    );
99    let arg = TensorMapArg {
100        tensor: binding.into_tensor_arg(),
101        metadata: meta,
102        _kind: PhantomData,
103    };
104    (arg, transposed)
105}
106
107/// A storage-tiled operand's tensor map: the descriptor is the physical
108/// `[.., R/tr, C/tc, tr, tc]` the data is already stored in, and the box is one whole storage
109/// tile. Where [`tma_operand`] collapses a binding to the descriptor's `(batch, rows, cols)`
110/// and may swap the inner pair, this one takes the binding as it is: the storage tile is the
111/// innermost pair and therefore the fastest-varying, so one box is one contiguous run and there
112/// is nothing to transpose.
113///
114/// # Panics
115///
116/// A binding whose innermost two dims are not `tile`, or whose tile is not stored a row at a
117/// time: the descriptor takes the innermost dim to step by one and each row to follow the last,
118/// so any other tile would load in silence as the wrong values. The routine's own checks refuse
119/// both first.
120pub fn tma_operand_tiled(
121    binding: TensorBinding,
122    tile: (usize, usize),
123    storage_ty: ElemType,
124    swizzle: TensorMapSwizzle,
125) -> TensorMapArg<Tiled> {
126    let rank = binding.shape.len();
127    let stored = (binding.shape[rank - 2], binding.shape[rank - 1]);
128    assert_eq!(
129        stored, tile,
130        "tma_operand_tiled: the binding is stored in {stored:?} tiles, not the {tile:?} asked for"
131    );
132    let inner = (binding.strides[rank - 2], binding.strides[rank - 1]);
133    assert_eq!(
134        inner,
135        (tile.1, 1),
136        "tma_operand_tiled: a box is one contiguous run only for a tile stored a row at a time, \
137         which steps ({}, 1); this one steps {inner:?}",
138        tile.1
139    );
140    // One box per storage tile: unit on every outer dim, the whole tile on the inner pair.
141    let mut dims = vec![1usize; rank];
142    dims[rank - 2] = tile.0;
143    dims[rank - 1] = tile.1;
144    let box_shape: Shape = dims.into();
145    let meta = tma_meta_tiled(
146        Metadata::new(binding.shape.clone(), binding.strides.clone()),
147        box_shape,
148        remap_storage_for_tma(storage_ty),
149        swizzle,
150    );
151    TensorMapArg {
152        tensor: binding.into_tensor_arg(),
153        metadata: meta,
154        _kind: PhantomData,
155    }
156}
157
158/// Build a tiled [`TensorMapMeta`] with the defaults shared by every current call site
159/// (no interleave, no prefetch, OOB-fill = zero, elem_stride = `[1; rank]`).
160pub fn tma_meta_tiled(
161    metadata: Metadata,
162    tile_size: Shape,
163    elem_ty: ElemType,
164    swizzle: TensorMapSwizzle,
165) -> TensorMapMeta {
166    let rank = metadata.rank();
167    TensorMapMeta {
168        format: TensorMapFormat::Tiled(TiledArgs { tile_size }),
169        metadata,
170        elem_stride: Strides::new(&vec![1; rank]),
171        interleave: TensorMapInterleave::None,
172        swizzle,
173        prefetch: TensorMapPrefetch::None,
174        oob_fill: OobFill::Zero,
175        elem_ty,
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use cubecl::zspace::Tiling;
182
183    use super::*;
184
185    /// A weight tiled to a 32x64 stage: `[256, 512]` stored as `[8, 8, 32, 64]`.
186    fn tiled() -> TensorBinding {
187        let client = cubecl::test_device().client();
188        let shape = shape![8, 8, 32, 64];
189        let strides = strides![64 * 32 * 8, 64 * 32, 64, 1];
190        TensorBinding {
191            handle: client.empty(256 * 512 * 4).binding(),
192            strides,
193            shape,
194            tiling: Tiling::new(&[2, 2]).unwrap(),
195        }
196    }
197
198    #[test]
199    fn a_stored_descriptor_keeps_the_rank_and_boxes_one_storage_tile() {
200        let arg = tma_operand_tiled(
201            tiled(),
202            (32, 64),
203            f32::elem_type_native(),
204            TensorMapSwizzle::None,
205        );
206        // The descriptor is the buffer as it is stored, not collapsed to (batch, row, col).
207        assert_eq!(&arg.metadata.metadata.shape()[..], &[8, 8, 32, 64]);
208        assert_eq!(
209            &arg.metadata.metadata.strides()[..],
210            &[64 * 32 * 8, 64 * 32, 64, 1]
211        );
212        // One box is one storage tile: unit outside, the whole tile on the inner pair, which is
213        // the fastest-varying pair, so the box is one contiguous run.
214        let TensorMapFormat::Tiled(TiledArgs { tile_size }) = &arg.metadata.format else {
215            panic!("expected a tiled tensor map");
216        };
217        assert_eq!(&tile_size[..], &[1, 1, 32, 64]);
218    }
219
220    #[test]
221    #[should_panic(expected = "stored in (32, 64) tiles, not the (64, 64) asked for")]
222    fn a_stored_descriptor_refuses_a_tile_the_buffer_does_not_hold() {
223        tma_operand_tiled(
224            tiled(),
225            (64, 64),
226            f32::elem_type_native(),
227            TensorMapSwizzle::None,
228        );
229    }
230}
cubek_std/launch/tma.rs

cubek_std/launch/
tma.rs
