cubek-std 0.3.0-pre.1

CubeK: Standard Library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

use cubecl::prelude::*;

use crate::tile::mask::Mask;
use crate::tile::{Plane, Tile, TileExpand, TileKind, TileKindExpand};

#[cube]
impl<E: Float> Tile<E, Plane> {
    /// Multiply by `scale` and add `-inf` at masked positions.
    pub fn scale_and_mask<M: Mask>(&mut self, scale: E, mask: &M) {
        match &mut self.kind {
            TileKind::Unit(t) => t.scale_and_mask::<M>(scale, mask),
            TileKind::WhiteboxFragment(t) => t.scale_and_mask::<M>(scale, mask),
            TileKind::Bounce(b) => b.scale_and_mask::<M>(scale, mask),
            TileKind::Register(t) => t.scale_and_mask::<M>(scale, mask),
            _ => panic!("scale_and_mask: unsupported tile variant"),
        }
    }

    pub fn fill_zero(&mut self) {
        match &mut self.kind {
            TileKind::Unit(t) => t.fill_zero(),
            TileKind::WhiteboxFragment(t) => t.zero(),
            TileKind::Bounce(b) => b.fill_zero(),
            TileKind::Cmma(t) => t.fill_zero(),
            TileKind::Register(t) => t.fill_zero(),
            TileKind::RowWise(t) => t.init_zero(),
            _ => panic!("fill_zero: unsupported tile variant"),
        }
    }
}