cubek_std/tile/ops/
rowwise.rs

1use cubecl::prelude::*;
2
3use crate::tile::{Plane, RowWise, Tile, TileExpand, TileKind, TileKindExpand};
4
5#[cube]
6impl<E: Float> Tile<E, Plane> {
7    pub fn row_max(&self, acc: &mut RowWise<E>, base: &RowWise<E>) {
8        match &self.kind {
9            TileKind::Unit(t) => t.row_max(acc, base),
10            TileKind::WhiteboxFragment(t) => t.row_max(acc, base),
11            TileKind::Bounce(b) => b.row_max(acc, base),
12            TileKind::Register(t) => t.row_max(acc, base),
13            _ => panic!("row_max: unsupported tile variant"),
14        }
15    }
16
17    pub fn row_sum(&self, acc: &mut RowWise<E>) {
18        match &self.kind {
19            TileKind::Unit(t) => t.row_sum(acc),
20            TileKind::WhiteboxFragment(t) => t.row_sum(acc),
21            TileKind::Bounce(b) => b.row_sum(acc),
22            TileKind::Register(t) => t.row_sum(acc),
23            _ => panic!("row_sum: unsupported tile variant"),
24        }
25    }
26
27    pub fn exp_diff(&mut self, rowwise: &RowWise<E>) {
28        match &mut self.kind {
29            TileKind::Unit(t) => t.exp_diff(rowwise),
30            TileKind::WhiteboxFragment(t) => t.exp_diff(rowwise),
31            TileKind::Bounce(b) => b.exp_diff(rowwise),
32            TileKind::Register(t) => t.exp_diff(rowwise),
33            _ => panic!("exp_diff: unsupported tile variant"),
34        }
35    }
36
37    pub fn rowwise_scale(&mut self, scale: &RowWise<E>) {
38        match &mut self.kind {
39            TileKind::Unit(t) => t.rowwise_scale(scale),
40            TileKind::WhiteboxFragment(t) => t.rowwise_scale(scale),
41            TileKind::Bounce(b) => b.rowwise_scale(scale),
42            TileKind::Register(t) => t.rowwise_scale(scale),
43            _ => panic!("rowwise_scale: unsupported tile variant"),
44        }
45    }
46
47    /// Multiply each row of `self` by `scale[r]`. The `Bounce` arm
48    /// round-trips through smem to keep the cmma fragment current.
49    pub fn scale_mul<SM: Float>(&mut self, scale: &RowWise<SM>) {
50        let scale_e = RowWise::<SM>::cast_from::<E>(scale);
51        match &mut self.kind {
52            TileKind::Bounce(b) => {
53                b.cmma_to_fragment();
54                b.rowwise_scale(&scale_e);
55                b.fragment_to_cmma();
56            }
57            TileKind::WhiteboxFragment(t) => t.rowwise_scale(&scale_e),
58            TileKind::Unit(t) => t.rowwise_scale(&scale_e),
59            TileKind::Register(t) => t.rowwise_scale(&scale_e),
60            _ => panic!("scale_mul: unsupported tile variant"),
61        }
62    }
63
64    /// Divide each row by `running_state_l[r]`; fully-masked rows stay zero.
65    pub fn scale_div<SM: Float>(&mut self, running_state_l: &RowWise<SM>) {
66        let mut scale = RowWise::<SM>::cast_from::<E>(running_state_l);
67        scale.recip_inplace();
68        match &mut self.kind {
69            TileKind::Bounce(b) => {
70                b.cmma_to_fragment();
71                b.rowwise_scale(&scale);
72                b.fragment_to_cmma();
73            }
74            TileKind::WhiteboxFragment(t) => t.rowwise_scale(&scale),
75            TileKind::Unit(t) => t.rowwise_scale(&scale),
76            TileKind::Register(t) => t.rowwise_scale(&scale),
77            _ => panic!("scale_div: unsupported tile variant"),
78        }
79    }
80}
cubek_std/tile/ops/rowwise.rs

cubek_std/tile/ops/
rowwise.rs
