cubecl_tile/launch/arg/
tma.rs

1//! A TMA tensor map as an operand's values ([`TileArgLaunch::tensor_map`]) and its in-kernel
2//! layouts.
3
4use cubecl::prelude::*;
5use cubecl_std::tensor::{
6    layout::{CoordsDyn, Layout, LayoutExpand, as_dyn::IntoDynLayout},
7    view::launch::VirtualViewLayoutLaunch,
8};
9
10use crate::*;
11
12/// The box a TMA descriptor moves: logical `(rows, cols)`, optional batch, and whether the
13/// descriptor is transposed (column-major).
14#[derive(Clone, Copy, PartialEq, Eq, Debug)]
15pub struct TmaBox {
16    pub rows: u32,
17    pub cols: u32,
18    pub batch: Option<u32>,
19    pub transposed: bool,
20}
21
22/// What the TMA family launches from: the tensor map, the axes it spans, and its box.
23pub struct TmaOperand {
24    pub map: TensorMapArg<Tiled>,
25    pub axes: Vec<Axis>,
26    pub shape: TmaBox,
27}
28
29impl TmaOperand {
30    /// The tensor map as an operand's values, served at the element it moves, under no scales.
31    pub fn arg(self) -> TileArgLaunch {
32        TileArgLaunch::tensor_map(self.map, &self.axes, self.shape)
33    }
34}
35
36impl TileArgLaunch {
37    /// A TMA tensor map as a tile argument over `axes`: two, or three with `shape.batch` set.
38    pub fn tensor_map(tensor_map: TensorMapArg<Tiled>, axes: &[Axis], shape: TmaBox) -> Self {
39        let layout = Self::dyn_layout(axes, shape);
40        Self::of_tensor_map(
41            tensor_map,
42            VirtualViewLayoutLaunch::new::<IntoDynLayout<TmaDynLayout>>(layout),
43            TileSpec::direct(axes),
44        )
45    }
46
47    /// The layout of a descriptor over `axes`.
48    fn dyn_layout(axes: &[Axis], shape: TmaBox) -> TmaDynLayoutLaunch {
49        let batched = match (axes.len(), shape.batch) {
50            (2, None) => false,
51            (3, Some(_)) => true,
52            (r, batch) => panic!(
53                "TileArgLaunch::tensor_map: the descriptor is (batch, row, col); {r} axes with \
54                 batch {batch:?} is not a matrix nor a batched one"
55            ),
56        };
57        let dims = (shape.batch.unwrap_or(1), shape.rows, shape.cols);
58        TmaDynLayoutLaunch::new(dims, batched, shape.transposed)
59    }
60}
61
62/// Where a box of a storage-tiled operand lands its values: the order the descriptor states the
63/// operand in.
64#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
65pub enum TmaLanding {
66    /// One storage tile a box, landing as the contiguous run it is stored as: the descriptor is
67    /// the operand as stored, `[.., R/tr, C/tc, tr, tc]`.
68    Tiles,
69    /// Whole tiles along the rows, landing the box's rows one after another, each a run of its
70    /// tiles' rows, as a row-major operand's box would land them: the descriptor is the operand as
71    /// `[R/tr, tr, C/tc, tc]`, the tiles' rows nested inside the tiles' columns. Along the columns
72    /// a box is whole tiles, or a part of one tile's run that divides it, starting where the box's
73    /// own columns do.
74    Rows,
75}
76
77/// In-kernel tensor-map layout aligning logical [`CoordsDyn`] to the descriptor's
78/// `(batch, row, col)`.
79#[derive(CubeType, CubeLaunch, Clone)]
80pub(crate) struct TmaDynLayout {
81    /// Logical `(batch, rows, cols)` of the operand.
82    dims: (u32, u32, u32),
83    #[cube(comptime)]
84    batched: bool,
85    #[cube(comptime)]
86    transposed: bool,
87}
88
89#[cube]
90impl Layout for TmaDynLayout {
91    type Coordinates = CoordsDyn;
92    type SourceCoordinates = CoordsDyn;
93
94    fn to_source_pos(&self, pos: Self::Coordinates) -> Self::SourceCoordinates {
95        let (batch, _rows, _cols) = self.dims;
96        let mut src = CoordsDyn::new();
97        if comptime!(self.batched) {
98            // A unit-batch descriptor is a broadcast: always read batch 0.
99            src.push(select(batch == 1, 0u32, pos[0]));
100        } else {
101            src.push(0u32);
102        }
103        let (r, c) = comptime!(if self.batched { (1, 2) } else { (0, 1) });
104        // TMA discards the last stride, so a col-major descriptor is transposed; swap back.
105        if comptime!(self.transposed) {
106            src.push(pos[c]);
107            src.push(pos[r]);
108        } else {
109            src.push(pos[r]);
110            src.push(pos[c]);
111        }
112        src
113    }
114
115    fn to_source_pos_checked(&self, pos: Self::Coordinates) -> (Self::SourceCoordinates, bool) {
116        // TMA loads are clamped by the descriptor; no in-kernel bounds check.
117        (self.to_source_pos(pos), true)
118    }
119
120    fn shape(&self) -> Self::Coordinates {
121        let (batch, rows, cols) = self.dims;
122        let mut s = CoordsDyn::new();
123        if comptime!(self.batched) {
124            s.push(batch);
125        }
126        s.push(rows);
127        s.push(cols);
128        s
129    }
130
131    fn is_in_bounds(&self, _pos: Self::Coordinates) -> bool {
132        true
133    }
134}
cubecl_tile/launch/arg/tma.rs

cubecl_tile/launch/arg/
tma.rs
