rudnn/attention/kernel_ir/components/tile/
mask.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::library::tensor::layout::Coordinates;
3use ruda_kernel::dsl::prelude::*;
4use ruda_kernel::library::tensor::layout::Coords2d;
5use ruda_kernel::tiling::tile::{
6    Mask, MaskExpand, MaskLayout, Plane, StridedTile, Tile, allocate_unit_tile,
7    allocate_whitebox_fragment, mask_layout_absolute_pos,
8};
9
10/// Comptime configuration that drives [`MaskTile::new`] and
11/// `MaskTile::should_mask`. Built attention-side from a `TileSoftmax`.
12#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
13pub struct MaskConfig {
14    pub layout: MaskLayout,
15    pub causal: bool,
16    pub materialized: bool,
17}
18
19#[derive(RudaType)]
20/// Mask tile for Tile Attention. It is an additive mask: the result of `apply`
21/// should be added, not multiplied.
22pub enum MaskTile<F: Float> {
23    /// When a mask tensor is supplied. Also contains a logical part.
24    Materialized(MaterializedTileMask<F>),
25    /// When no mask tensor is supplied. Used for out-of-bounds and causal mask.
26    Logical(LogicalTileMask),
27}
28
29#[ruda]
30impl<F: Float> MaskTile<F> {
31    pub fn new(
32        out_of_bounds: ComptimeOption<Coords2d>,
33        #[comptime] config: MaskConfig,
34    ) -> MaskTile<F> {
35        let logical_mask = LogicalTileMask::new(config, out_of_bounds);
36
37        if comptime!(config.materialized) {
38            let fragment: Tile<F, Plane, ReadWrite> = match comptime!(config.layout) {
39                MaskLayout::Unit(l) => allocate_unit_tile::<F, Plane>(comptime!(l)),
40                MaskLayout::WhiteboxFragment(l) => {
41                    allocate_whitebox_fragment::<F, Plane>(comptime!(l))
42                }
43            };
44            MaskTile::new_Materialized(MaterializedTileMask::<F> {
45                fragment,
46                logical_mask,
47            })
48        } else {
49            MaskTile::new_Logical(logical_mask)
50        }
51    }
52
53    /// Loads the mask data into the fragment if a tile is given; otherwise
54    /// only updates the logical mask's origin.
55    pub fn update<E: Numeric, ES: Size>(
56        &mut self,
57        new_origin: Coords2d,
58        tile: ComptimeOption<StridedTile<E, ES>>,
59    ) {
60        match self {
61            MaskTile::Materialized(m) => {
62                m.logical_mask.update_origin(new_origin);
63                m.update_tile(tile.unwrap());
64            }
65            MaskTile::Logical(l) => l.update_origin(new_origin),
66        }
67    }
68}
69
70#[ruda]
71impl<F: Float> Mask for MaskTile<F> {
72    fn should_mask(&self, local_pos: Coords2d) -> bool {
73        match self {
74            MaskTile::Materialized(m) => m.should_mask(local_pos),
75            MaskTile::Logical(l) => l.should_mask(local_pos),
76        }
77    }
78}
79
80#[derive(RudaType)]
81/// Origin of the logical mask, updated when changing partition or tile within
82/// partition.
83pub struct LogicalIterOrigin {
84    row: RuntimeCell<u32>,
85    col: RuntimeCell<u32>,
86}
87
88#[ruda]
89impl LogicalIterOrigin {
90    fn init() -> LogicalIterOrigin {
91        LogicalIterOrigin {
92            row: RuntimeCell::new(0),
93            col: RuntimeCell::new(0),
94        }
95    }
96
97    fn read(&self) -> Coords2d {
98        (self.row.read(), self.col.read())
99    }
100
101    fn update(&mut self, new: Coords2d) {
102        self.row.store(new.0);
103        self.col.store(new.1);
104    }
105}
106
107#[derive(RudaType)]
108pub struct LogicalTileMask {
109    logical_iter_origin: LogicalIterOrigin,
110    #[ruda(comptime)]
111    causal: bool,
112    out_of_bounds: ComptimeOption<Coords2d>,
113    #[ruda(comptime)]
114    fragment_layout: MaskLayout,
115}
116
117#[ruda]
118impl LogicalTileMask {
119    pub fn new(
120        #[comptime] config: MaskConfig,
121        out_of_bounds: ComptimeOption<Coords2d>,
122    ) -> LogicalTileMask {
123        LogicalTileMask {
124            logical_iter_origin: LogicalIterOrigin::init(),
125            causal: comptime!(config.causal),
126            out_of_bounds,
127            fragment_layout: comptime!(config.layout),
128        }
129    }
130
131    pub fn should_mask(&self, local_pos: Coords2d) -> bool {
132        let pos_in_tile = mask_layout_absolute_pos(self.fragment_layout, local_pos);
133
134        let pos = Coords2d::add(self.logical_iter_origin.read(), pos_in_tile);
135
136        let causal_masked = self.causal && pos.0 < pos.1;
137
138        #[comptime]
139        let oob_masked = match self.out_of_bounds {
140            ComptimeOption::Some(bounds) => !Coords2d::is_in_bounds(&pos, &bounds),
141            ComptimeOption::None => false,
142        };
143
144        causal_masked || oob_masked
145    }
146
147    pub fn update_origin(&mut self, new_origin: Coords2d) {
148        self.logical_iter_origin.update(new_origin);
149    }
150}
151
152#[derive(RudaType)]
153pub struct MaterializedTileMask<F: Float> {
154    fragment: Tile<F, Plane, ReadWrite>,
155    logical_mask: LogicalTileMask,
156}
157
158#[ruda]
159impl<F: Float> MaterializedTileMask<F> {
160    pub fn should_mask(&self, local_pos: Coords2d) -> bool {
161        let logical_masked = self.logical_mask.should_mask(local_pos);
162        let materialized_masked = self.fragment.should_mask(local_pos);
163
164        logical_masked || materialized_masked
165    }
166
167    pub fn update_tile<MSK: Numeric, MSKS: Size>(&mut self, tile: StridedTile<MSK, MSKS>) {
168        self.fragment
169            .load_mask_from_strided_tile::<MSK, MSKS>(&tile);
170    }
171}
rudnn/attention/kernel_ir/components/tile/mask.rs

rudnn/attention/kernel_ir/components/tile/
mask.rs
