cubek-std 0.3.0-pre.2

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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

use cubecl::prelude::*;

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

#[cube]
impl<E: Float> Tile<E, Plane> {
    pub fn row_max(&self, acc: &mut RowWise<E>, base: &RowWise<E>) {
        match &self.kind {
            TileKind::Unit(t) => t.row_max(acc, base),
            TileKind::WhiteboxFragment(t) => t.row_max(acc, base),
            TileKind::Bounce(b) => b.row_max(acc, base),
            TileKind::Register(t) => t.row_max(acc, base),
            _ => panic!("row_max: unsupported tile variant"),
        }
    }

    pub fn row_sum(&self, acc: &mut RowWise<E>) {
        match &self.kind {
            TileKind::Unit(t) => t.row_sum(acc),
            TileKind::WhiteboxFragment(t) => t.row_sum(acc),
            TileKind::Bounce(b) => b.row_sum(acc),
            TileKind::Register(t) => t.row_sum(acc),
            _ => panic!("row_sum: unsupported tile variant"),
        }
    }

    pub fn exp_diff(&mut self, rowwise: &RowWise<E>) {
        match &mut self.kind {
            TileKind::Unit(t) => t.exp_diff(rowwise),
            TileKind::WhiteboxFragment(t) => t.exp_diff(rowwise),
            TileKind::Bounce(b) => b.exp_diff(rowwise),
            TileKind::Register(t) => t.exp_diff(rowwise),
            _ => panic!("exp_diff: unsupported tile variant"),
        }
    }

    pub fn rowwise_scale(&mut self, scale: &RowWise<E>) {
        match &mut self.kind {
            TileKind::Unit(t) => t.rowwise_scale(scale),
            TileKind::WhiteboxFragment(t) => t.rowwise_scale(scale),
            TileKind::Bounce(b) => b.rowwise_scale(scale),
            TileKind::Register(t) => t.rowwise_scale(scale),
            _ => panic!("rowwise_scale: unsupported tile variant"),
        }
    }

    /// Multiply each row of `self` by `scale[r]`. The `Bounce` arm
    /// round-trips through smem to keep the cmma fragment current.
    pub fn scale_mul<SM: Float>(&mut self, scale: &RowWise<SM>) {
        let scale_e = RowWise::<SM>::cast_from::<E>(scale);
        match &mut self.kind {
            TileKind::Bounce(b) => {
                b.cmma_to_fragment();
                b.rowwise_scale(&scale_e);
                b.fragment_to_cmma();
            }
            TileKind::WhiteboxFragment(t) => t.rowwise_scale(&scale_e),
            TileKind::Unit(t) => t.rowwise_scale(&scale_e),
            TileKind::Register(t) => t.rowwise_scale(&scale_e),
            _ => panic!("scale_mul: unsupported tile variant"),
        }
    }

    /// Divide each row by `running_state_l[r]`; fully-masked rows stay zero.
    pub fn scale_div<SM: Float>(&mut self, running_state_l: &RowWise<SM>) {
        let mut scale = RowWise::<SM>::cast_from::<E>(running_state_l);
        scale.recip_inplace();
        match &mut self.kind {
            TileKind::Bounce(b) => {
                b.cmma_to_fragment();
                b.rowwise_scale(&scale);
                b.fragment_to_cmma();
            }
            TileKind::WhiteboxFragment(t) => t.rowwise_scale(&scale),
            TileKind::Unit(t) => t.rowwise_scale(&scale),
            TileKind::Register(t) => t.rowwise_scale(&scale),
            _ => panic!("scale_div: unsupported tile variant"),
        }
    }
}