cubek_std/tile/ops/
elementwise.rs

1use cubecl::prelude::*;
2
3use crate::tile::mask::Mask;
4use crate::tile::{Plane, Tile, TileExpand, TileKind, TileKindExpand};
5
6#[cube]
7impl<E: Float> Tile<E, Plane> {
8    /// Multiply by `scale` and add `-inf` at masked positions.
9    pub fn scale_and_mask<M: Mask>(&mut self, scale: E, mask: &M) {
10        match &mut self.kind {
11            TileKind::Unit(t) => t.scale_and_mask::<M>(scale, mask),
12            TileKind::WhiteboxFragment(t) => t.scale_and_mask::<M>(scale, mask),
13            TileKind::Bounce(b) => b.scale_and_mask::<M>(scale, mask),
14            TileKind::Register(t) => t.scale_and_mask::<M>(scale, mask),
15            _ => panic!("scale_and_mask: unsupported tile variant"),
16        }
17    }
18
19    pub fn fill_zero(&mut self) {
20        match &mut self.kind {
21            TileKind::Unit(t) => t.fill_zero(),
22            TileKind::WhiteboxFragment(t) => t.zero(),
23            TileKind::Bounce(b) => b.fill_zero(),
24            TileKind::Cmma(t) => t.fill_zero(),
25            TileKind::Register(t) => t.fill_zero(),
26            TileKind::RowWise(t) => t.init_zero(),
27            _ => panic!("fill_zero: unsupported tile variant"),
28        }
29    }
30}
cubek_std/tile/ops/elementwise.rs

cubek_std/tile/ops/
elementwise.rs
