cubek_std/tile/
mask.rs

1use cubecl;
2use cubecl::{prelude::*, std::tensor::layout::Coords2d};
3
4use crate::tile::variants::{InnerLayout, UnitTileLayout, WhiteboxFragmentLayout};
5
6/// Minimal mask abstraction used by row-wise tile operations.
7/// Returns `true` when the element at `local_pos` should be treated as masked
8/// (i.e. driven to -inf by `Tile::scale_and_mask`).
9#[cube]
10pub trait Mask: CubeType {
11    fn should_mask(&self, local_pos: Coords2d) -> bool;
12}
13
14/// Layout of an attention-style mask fragment across the units of a plane.
15/// Purely comptime — all variants carry only comptime data.
16#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
17pub enum MaskLayout {
18    /// Each unit owns a full row-major copy of the tile.
19    Unit(UnitTileLayout),
20    /// The tile is fragmented across plane units, with the layout described by
21    /// [`WhiteboxFragmentLayout`].
22    WhiteboxFragment(WhiteboxFragmentLayout),
23}
24
25impl MaskLayout {
26    pub const fn unit(num_rows: u32, num_cols: u32) -> MaskLayout {
27        MaskLayout::Unit(UnitTileLayout {
28            num_rows,
29            num_cols,
30            transposed_load: false,
31        })
32    }
33
34    pub const fn whitebox_fragment(
35        tile_shape: Coords2d,
36        plane_dim: u32,
37        inner_layout: InnerLayout,
38    ) -> MaskLayout {
39        let total_elements = tile_shape.0 * tile_shape.1;
40        let elements_per_unit = total_elements.div_ceil(plane_dim);
41        let (num_rows_per_unit, num_cols_per_unit) = match inner_layout {
42            InnerLayout::Contiguous => (1u32, elements_per_unit),
43            InnerLayout::SplitRows => (2u32, elements_per_unit / 2u32),
44        };
45        let unit_size = (num_rows_per_unit, num_cols_per_unit);
46        let num_units_per_row = tile_shape.1 / unit_size.1;
47
48        MaskLayout::WhiteboxFragment(WhiteboxFragmentLayout {
49            total_size: tile_shape,
50            unit_size,
51            num_units_per_row,
52            plane_dim,
53        })
54    }
55}
56
57#[cube]
58/// Returns how many units in a plane participate in the same row.
59pub fn mask_layout_num_units_per_row(#[comptime] layout: MaskLayout) -> comptime_type!(u32) {
60    match layout {
61        MaskLayout::Unit(_) => 1u32,
62        MaskLayout::WhiteboxFragment(l) => comptime!(l.total_size.1 / l.unit_size.1),
63    }
64}
65
66#[cube]
67/// Maps a per-unit `(row, col)` to its absolute position within the tile.
68pub fn mask_layout_absolute_pos(#[comptime] layout: MaskLayout, local_pos: Coords2d) -> Coords2d {
69    match layout {
70        MaskLayout::Unit(_) => local_pos,
71        MaskLayout::WhiteboxFragment(l) => {
72            let abs_row_index = {
73                let row_0 = UNIT_POS_X / l.num_units_per_row;
74                let row_jump = comptime!(l.plane_dim / l.num_units_per_row);
75                local_pos.0 * row_jump + row_0
76            };
77            let abs_col_index = l.unit_size.1 * (UNIT_POS_X % l.num_units_per_row) + local_pos.1;
78            (abs_row_index, abs_col_index)
79        }
80    }
81}
cubek_std/tile/mask.rs

cubek_std/tile/
mask.rs
