cubecl_tile/stage/
base.rs

1//! Where a stage lives and how it lays its cells out ([`StageStorage`]).
2
3use crate::{Axis, ChunkSwizzle, LineBytes, Refusal, Space, StageForm};
4use cubecl::prelude::TensorMapSwizzle;
5
6/// Where a stage lives and how it lays its cells out, stated at
7/// [`Stages::smem`](crate::Stages::smem).
8#[derive(Clone, PartialEq, Eq, Hash, Debug)]
9pub enum StageStorage {
10    /// Shared memory grouped into `block`-sized tiles, one edge per axis of the operation.
11    Tiled {
12        block: Vec<(Axis, usize)>,
13        chunks: RowChunks,
14    },
15    Strided,
16    /// The plane's units, one line per unit.
17    Lines {
18        read: UnitRead,
19    },
20}
21
22/// How a [`Tiled`](StageStorage::Tiled) block lays each row down, in 16-byte chunks.
23#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, serde::Serialize, serde::Deserialize)]
24pub enum RowChunks {
25    /// Each row one contiguous run, in order; readable by `cmma` and TMA.
26    InOrder,
27    /// Each row's chunks XOR-permuted by a row key; only per-unit-addressed readers can read it.
28    Swizzled,
29    /// Each row in order followed by one unused chunk, so rows start on distinct banks.
30    Padded,
31}
32
33impl RowChunks {
34    /// Bytes one chunk holds: the widest read one unit issues.
35    pub const CHUNK_BYTES: usize = 16;
36}
37
38/// How a value reaches the unit asking for it, out of the lines the plane holds.
39#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, serde::Serialize, serde::Deserialize)]
40pub enum UnitRead {
41    /// By plane shuffle: one shuffle per word of the line plus a dynamic pick.
42    Shuffle,
43    /// Through a plane-owned shared-memory window written once per load. Serialized as `Window`.
44    #[serde(rename = "Window")]
45    PlaneShared,
46}
47
48impl StageStorage {
49    /// The swizzle a TMA descriptor moving the whole of a stage over `space` in one box lands its
50    /// rows with, so they lie as this storage keeps them: the stage's lines `vector_size` values of
51    /// `elem_bits` bits, which a line holds whole bytes of (a packed operand's line is its words).
52    /// What the launch builds the descriptor with, and what the stage fill holds the stage to.
53    ///
54    /// # Errors
55    ///
56    /// Blocks splitting a row, which a box lands whole; blocks stacking a number of rows the
57    /// swizzle's period of rows does not divide, whose rows a box keys off their place in the
58    /// stage rather than in the block; padded rows, which a box lands dense, and swizzled rows
59    /// longer than a descriptor swizzle spans.
60    pub fn tma_swizzle(
61        &self,
62        space: &Space,
63        vector_size: usize,
64        elem_bits: usize,
65    ) -> Result<TensorMapSwizzle, Refusal> {
66        assert!(
67            (vector_size * elem_bits).is_multiple_of(8),
68            "StageStorage::tma_swizzle: a line of {vector_size} {elem_bits}-bit values is not whole \
69             bytes"
70        );
71        let form = StageForm::dense(
72            space,
73            vector_size,
74            self.clone(),
75            LineBytes(vector_size * elem_bits / 8),
76        );
77        let swizzle = form
78            .rows
79            .tma_swizzle()
80            .map_err(|why| Refusal::RowsNoDescriptorLands { why })?;
81        let rank = space.rank();
82        let row = space.axis_at(rank - 1);
83        let period = match swizzle {
84            TensorMapSwizzle::None => 1,
85            _ => ChunkSwizzle::ROWS_PER_PERIOD,
86        };
87        for block in self.nesting(space) {
88            let rows = match rank {
89                1 => 1,
90                _ => block.extent_at(rank - 2),
91            };
92            if block.extent_at(rank - 1) < space.extent_at(rank - 1) || !rows.is_multiple_of(period)
93            {
94                return Err(Refusal::BoxSplitsStageBlocks {
95                    axis: row,
96                    rows,
97                    period,
98                });
99            }
100        }
101        Ok(swizzle)
102    }
103
104    /// Bytes a stage over `space` holds laid out this way, its lines `vector_size` values of
105    /// `elem_bits` bits, which a line holds whole bytes of: every line of its cells, and the
106    /// padding after each row. What [`Stages::smem`](crate::Stages::smem) allocates for it.
107    pub fn bytes(&self, space: &Space, vector_size: usize, elem_bits: usize) -> usize {
108        let line = vector_size * elem_bits / 8;
109        StageForm::dense(space, vector_size, self.clone(), LineBytes(line)).cells() * line
110    }
111
112    /// The storage-tiling nesting a stage over `space` gets, coarse to fine; empty is row-major.
113    pub(crate) fn nesting(&self, space: &Space) -> Vec<Space> {
114        match self {
115            StageStorage::Lines { .. } => {
116                panic!("StageStorage::Lines: the plane's units are not shared memory")
117            }
118            StageStorage::Tiled { block, .. } => {
119                let nested = Space::new(
120                    &space
121                        .axes()
122                        .map(|axis| {
123                            let edge = block
124                                .iter()
125                                .find(|&&(a, _)| a == axis)
126                                .unwrap_or_else(|| {
127                                    panic!(
128                                        "StageStorage::Tiled: the block states no edge for {axis:?}"
129                                    )
130                                })
131                                .1;
132                            (axis, edge)
133                        })
134                        .collect::<Vec<_>>(),
135                );
136                if &nested == space {
137                    return Vec::new();
138                }
139                vec![nested]
140            }
141            StageStorage::Strided => Vec::new(),
142        }
143    }
144}
cubecl_tile/stage/base.rs

cubecl_tile/stage/
base.rs
