cubecl_tile/tile/kind/memory/view/
masked.rs

1use cubecl::prelude::*;
2use cubecl_std::tensor::{View, ViewMut, layout::Coordinates};
3
4/// A [`View`] plus a comptime `check` flag: reads past the overhang zero, writes skip.
5#[derive(CubeType)]
6pub struct Masked<'a, T: CubePrimitive, C: Coordinates + 'a> {
7    view: View<'a, T, C>,
8    #[cube(comptime)]
9    pub(crate) check: bool,
10}
11
12#[cube]
13impl<'a, T: CubePrimitive, C: Coordinates + 'a> Masked<'a, T, C> {
14    pub fn new(view: View<'a, T, C>, #[comptime] check: bool) -> Self {
15        Masked::<'a, T, C> { view, check }
16    }
17
18    pub fn read(&self, pos: C) -> T {
19        if comptime!(self.check) {
20            self.view.read_checked(pos)
21        } else {
22            // `check == false`: the launch proved this access in-bounds.
23            self.view.read_unchecked(pos)
24        }
25    }
26
27    /// The value at `pos`, untested: the reader proved it in bounds, so reading past them is out
28    /// of bounds, not masked.
29    pub fn read_unchecked(&self, pos: C) -> T {
30        self.view.read_unchecked(pos)
31    }
32
33    /// Whether `pos` lands on real data; always `true` when `check` is `false`.
34    pub fn is_in_bounds(&self, pos: C) -> bool {
35        if comptime!(self.check) {
36            self.view.is_in_bounds(pos)
37        } else {
38            true.runtime()
39        }
40    }
41
42    /// Whether the non-empty box at `pos` with `extent` is wholly in bounds.
43    pub(crate) fn block_in_bounds(&self, pos: C, extent: C) -> bool {
44        if comptime!(self.check) {
45            let one = C::from_int(pos.clone(), 1i64);
46            let far = C::sub(C::add(pos, extent), one);
47            self.view.is_in_bounds(far)
48        } else {
49            true.runtime()
50        }
51    }
52
53    pub fn shape(&self) -> C {
54        self.view.shape()
55    }
56}
57
58#[cube]
59impl<'a, T: CubePrimitive, C: Coordinates + 'static> Masked<'a, T, C> {
60    /// The buffer the stage's arrangement places at `pos`, spanning `size`: the address `ldmatrix`
61    /// reads. Panics where the view masks.
62    pub(crate) fn line_slice(&self, pos: C, size: C) -> &[T] {
63        comptime!(assert!(
64            !self.check,
65            "Masked::line_slice: an address reads past the overhang a masked read would zero; \
66             serve the operand from a stage whose accesses the launch proved in bounds"
67        ));
68        self.view
69            .clone()
70            .slice_unchecked(pos, size)
71            .as_linear_slice()
72    }
73
74    /// The one item at `pos` as a run of the buffer, what an async copy reads by address: empty
75    /// where the view masks and `pos` falls past the bound, so the copy reads nothing and the
76    /// engine zero-fills the item a masked read would have zeroed.
77    pub(crate) fn item_run(&self, pos: C) -> &[T] {
78        let one = C::from_int(pos.clone(), 1i64);
79        let run = self
80            .view
81            .clone()
82            .slice_unchecked(pos.clone(), one)
83            .as_linear_slice();
84        if comptime!(self.check) {
85            let len = select(self.view.is_in_bounds(pos), 1usize, 0usize);
86            &run[0..len]
87        } else {
88            run
89        }
90    }
91}
92
93/// The mutable twin of [`Masked`].
94#[derive(CubeType)]
95pub struct MaskedMut<'a, T: CubePrimitive, C: Coordinates + 'a> {
96    view: ViewMut<'a, T, C>,
97    #[cube(comptime)]
98    pub(crate) check: bool,
99}
100
101#[cube]
102impl<'a, T: CubePrimitive, C: Coordinates + 'a> MaskedMut<'a, T, C> {
103    pub fn new(view: ViewMut<'a, T, C>, #[comptime] check: bool) -> Self {
104        MaskedMut::<'a, T, C> { view, check }
105    }
106
107    pub fn read(&self, pos: C) -> T {
108        if comptime!(self.check) {
109            self.view.read_checked(pos)
110        } else {
111            self.view.read_unchecked(pos)
112        }
113    }
114
115    pub fn write(&mut self, pos: C, value: T) {
116        if comptime!(self.check) {
117            self.view.write_checked(pos, value);
118        } else {
119            self.view.write(pos, value);
120        }
121    }
122
123    /// Mutable counterpart to [`Masked::block_in_bounds`].
124    pub(crate) fn block_in_bounds(&self, pos: C, extent: C) -> bool {
125        if comptime!(self.check) {
126            let one = C::from_int(pos.clone(), 1i64);
127            let far = C::sub(C::add(pos, extent), one);
128            self.view.is_in_bounds(far)
129        } else {
130            true.runtime()
131        }
132    }
133
134    pub fn shape(&self) -> C {
135        self.view.shape()
136    }
137}
cubecl_tile/tile/kind/memory/view/masked.rs

cubecl_tile/tile/kind/memory/view/
masked.rs
