Skip to main content

cubek_std/tile/ops/
mask.rs

1//! `Tile::should_mask` and `Tile::load_mask_from_strided_tile` dispatchers.
2
3use cubecl;
4use cubecl::{prelude::*, std::tensor::layout::Coords2d};
5
6use crate::tile::{
7    StridedTile, Tile, TileExpand, TileKind, TileKindExpand, TileScope,
8    mask::{Mask, MaskExpand},
9};
10
11#[cube]
12impl<E: Numeric, Sc: TileScope> Mask for Tile<E, Sc> {
13    fn should_mask(&self, local_pos: Coords2d) -> bool {
14        match &self.kind {
15            TileKind::Unit(t) => t.should_mask(local_pos),
16            TileKind::WhiteboxFragment(t) => t.should_mask(local_pos),
17            _ => panic!(
18                "Mask::should_mask is only defined for Tile::Unit and Tile::WhiteboxFragment"
19            ),
20        }
21    }
22}
23
24#[cube]
25impl<N: Numeric, Sc: TileScope> Tile<N, Sc> {
26    /// Materialize a mask fragment from a `StridedTile` into `Unit` or
27    /// `WhiteboxFragment`.
28    pub fn load_mask_from_strided_tile<E: Numeric, ES: Size>(&mut self, tile: &StridedTile<E, ES>) {
29        match &mut self.kind {
30            TileKind::Unit(t) => t.load_from_strided_tile::<E, ES>(tile),
31            TileKind::WhiteboxFragment(t) => t.load_from_strided_tile::<E, ES>(tile),
32            _ => panic!("load_mask_from_strided_tile: unsupported tile variant"),
33        }
34    }
35}