cubek_std/tile/variants/
bounce.rs

1use cubecl;
2use cubecl::prelude::*;
3
4use crate::StageIdent;
5use crate::tile::{
6    Plane, RowWise, Tile, TileKind, TileKindExpand,
7    mask::Mask,
8    scope::{TileScope, assert_plane_scope},
9    variants::{
10        instruction::cmma::CmmaTile,
11        whitebox_fragment::{InnerLayout, WhiteboxFragment, WhiteboxFragmentLayout},
12    },
13};
14
15/// Comptime configuration for [`BounceTile`].
16///
17/// A bounce tile bundles an opaque cmma fragment together with a shared-memory
18/// scratch slice and a [`WhiteboxFragment`] view, so row-wise operations can be
19/// expressed as `copy_from` between the inner pieces. From the caller's point
20/// of view it is a single [`Tile`] variant — only valid when the tile's
21/// scope is `Plane`.
22#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
23pub struct BounceConfig {
24    pub tile_shape: (u32, u32),
25    pub num_planes: u32,
26    pub plane_dim: u32,
27    pub inner_layout: InnerLayout,
28}
29
30#[derive(CubeType)]
31pub struct BounceTile<N: Numeric> {
32    pub cmma: CmmaTile<N>,
33    pub smem: Shared<[N]>,
34    pub fragment: WhiteboxFragment<N>,
35}
36
37#[cube]
38impl<N: Numeric> BounceTile<N> {
39    pub fn new(cmma: CmmaTile<N>, #[comptime] cfg: BounceConfig) -> BounceTile<N> {
40        let total_tile_size = comptime!((cfg.tile_shape.0 * cfg.tile_shape.1) as usize);
41        let smem_size = comptime!(total_tile_size * cfg.num_planes as usize);
42        let start = UNIT_POS_Y as usize * total_tile_size;
43        let end = start + total_tile_size;
44        let smem = Shared::new_slice(smem_size).map(|smem| &smem[start..end]);
45
46        let layout = comptime!(WhiteboxFragmentLayout::new(
47            cfg.tile_shape,
48            cfg.plane_dim,
49            cfg.inner_layout
50        ));
51        let fragment = WhiteboxFragment::new(layout);
52
53        BounceTile::<N> {
54            cmma,
55            smem,
56            fragment,
57        }
58    }
59}
60
61#[cube]
62impl<E: Float> BounceTile<E> {
63    /// Synchronizes the fragment view from the cmma fragment via smem.
64    /// Call before any rowwise/elementwise op so the fragment reflects the
65    /// current cmma state.
66    pub fn cmma_to_fragment(&mut self) {
67        let stride = comptime!(self.cmma.tile_size.n());
68        cubecl::cmma::store(
69            &mut self.smem,
70            &self.cmma.matrix,
71            stride,
72            cubecl::cmma::MatrixLayout::RowMajor,
73        );
74        sync_cube();
75        self.fragment.load_from_slice(&self.smem);
76        sync_cube();
77    }
78
79    /// Synchronizes the cmma fragment from the fragment view via smem. Call
80    /// after rowwise/elementwise edits to make the cmma side current for the
81    /// next mma.
82    pub fn fragment_to_cmma(&mut self) {
83        let stride = comptime!(self.cmma.tile_size.n());
84        self.fragment.store_to(&mut self.smem);
85        sync_cube();
86        cubecl::cmma::load_with_layout(
87            &mut self.cmma.matrix,
88            &self.smem,
89            stride,
90            cubecl::cmma::MatrixLayout::RowMajor,
91        );
92    }
93
94    pub fn row_max(&self, acc: &mut RowWise<E>, base: &RowWise<E>) {
95        self.fragment.row_max(acc, base);
96    }
97
98    pub fn row_sum(&self, acc: &mut RowWise<E>) {
99        self.fragment.row_sum(acc);
100    }
101
102    pub fn exp_diff(&mut self, rowwise: &RowWise<E>) {
103        self.fragment.exp_diff(rowwise);
104    }
105
106    pub fn rowwise_scale(&mut self, scale: &RowWise<E>) {
107        self.fragment.rowwise_scale(scale);
108    }
109
110    pub fn scale_and_mask<M: Mask>(&mut self, scale: E, mask: &M) {
111        self.fragment.scale_and_mask::<M>(scale, mask);
112    }
113
114    /// Zeros the cmma fragment. The fragment view is not the live storage at
115    /// fill_zero call sites (always invoked before any cmma_to_fragment), so
116    /// only cmma needs clearing.
117    pub fn fill_zero(&mut self) {
118        cubecl::cmma::fill(&mut self.cmma.matrix, E::from_int(0));
119    }
120
121    /// Writes the (already-softmaxed) fragment view of this bounce tile into
122    /// `softmaxed`. The source fragment is plane-fragmented; for a `Bounce`
123    /// destination this routes through the destination's smem into its cmma
124    /// fragment.
125    pub fn write_fragment_to<Lhs: Float, Sc: TileScope>(&self, softmaxed: &mut Tile<Lhs, Sc>) {
126        write_fragment_into::<E, Lhs, Sc>(&self.fragment, softmaxed);
127    }
128}
129
130#[cube]
131fn write_fragment_into<Acc: Float, Lhs: Float, Sc: TileScope>(
132    src: &WhiteboxFragment<Acc>,
133    softmaxed: &mut Tile<Lhs, Sc>,
134) {
135    match &mut softmaxed.kind {
136        TileKind::Bounce(d) => {
137            let stride = comptime!(d.cmma.tile_size.n());
138            src.store_to(&mut d.smem);
139            sync_cube();
140            cubecl::cmma::load(&mut d.cmma.matrix, &d.smem, stride);
141        }
142        TileKind::WhiteboxFragment(d) => {
143            let total = comptime!(src.layout.unit_size.0 * src.layout.unit_size.1);
144            for i in 0..total {
145                d.array[i as usize] = Lhs::cast_from(src.array[i as usize]);
146            }
147        }
148        _ => panic!("write_fragment_to: unsupported softmaxed variant"),
149    }
150}
151
152#[cube]
153/// Wraps a freshly built `CmmaTile` in a `Tile::Bounce`. Panics at expansion
154/// time unless `Sc = Plane`.
155pub fn allocate_bounce_tile<E: Numeric, Sc: TileScope>(
156    cmma: CmmaTile<E>,
157    #[comptime] cfg: BounceConfig,
158) -> Tile<E, Sc> {
159    comptime!(assert_plane_scope(Sc::KIND));
160    Tile::from_kind(TileKind::new_Bounce(BounceTile::<E>::new(cmma, cfg)))
161}
162
163#[cube]
164impl<N: Numeric> BounceTile<N> {
165    /// Copies into the bounce tile's cmma fragment from `source`. Bounce
166    /// always loads through its CMMA representation (the WhiteboxFragment
167    /// view is synced lazily on demand by softmax/scale ops); supported
168    /// sources mirror [`CmmaTile::copy_from`].
169    pub fn copy_from<SE: Numeric, SS: Size, Sc: TileScope>(
170        &mut self,
171        source: &Tile<SE, Sc>,
172        #[comptime] ident: StageIdent,
173    ) {
174        self.cmma.copy_from::<SE, SS, Sc>(source, ident);
175    }
176
177    /// Zero-init the bounce tile (clears its cmma fragment).
178    pub fn init_zero(&mut self) {
179        self.cmma.init_zero();
180    }
181}
182
183#[cube]
184impl<Acc: Float> BounceTile<Acc> {
185    /// Online softmax for the Bounce variant. cmma → fragment once at entry
186    /// so all subsequent rowwise ops read/write the fragment view; the post-
187    /// exp values are still in the fragment at the end (we skip
188    /// `fragment_to_cmma` on `score` because its cmma is cleared next
189    /// iteration), and we stream straight into `softmaxed` via
190    /// `write_fragment_to`.
191    pub fn softmax<Lhs: Float, M: Mask>(
192        &mut self,
193        mask: &M,
194        softmaxed: &mut Tile<Lhs, Plane>,
195        state: &mut (RowWise<Acc>, RowWise<Acc>),
196        head_dim_factor: Acc,
197    ) -> RowWise<Acc> {
198        let num_rows = comptime!(state.0.num_rows);
199        let mut max_buf = RowWise::<Acc>::new_min_value(num_rows);
200        let mut sum_buf = RowWise::<Acc>::new_zero(num_rows);
201
202        self.cmma_to_fragment();
203
204        self.scale_and_mask::<M>(head_dim_factor, mask);
205        self.row_max(&mut max_buf, &state.0);
206        self.exp_diff(&max_buf);
207        self.row_sum(&mut sum_buf);
208
209        let exp_m_diff = state.0.exp_diff(&max_buf);
210        let new_l = exp_m_diff.mul(&state.1).add(&sum_buf);
211
212        self.write_fragment_to::<Lhs, Plane>(softmaxed);
213
214        RowWise::copy_from(&mut state.0, &max_buf);
215        RowWise::copy_from(&mut state.1, &new_l);
216
217        exp_m_diff
218    }
219}
cubek_std/tile/variants/bounce.rs

cubek_std/tile/variants/
bounce.rs
