cubek_std/tile/variants/
row_wise.rs

1use cubecl;
2use cubecl::prelude::*;
3
4/// Below this value a row is treated as numerically zero. Above f16's
5/// smallest normal (~6.1e-5).
6pub const FULLY_MASKED_ROW_THRESHOLD: f32 = 1e-4;
7
8#[derive(CubeType)]
9/// Contains one value per row of a fragment for which the unit contributes
10///
11/// Example: For a 8x8 tile shared by a plane of 32 units,
12/// every unit holds 8 values in the tile.
13///
14/// In the following layout, values are held contiguously, and num_rows=1 because
15/// every two occurrences of the same plane id are in the same row
16///  0,  0,  1,  1,  2,  2,  3,  3,
17///  4,  4,  5,  5,  6,  6,  7,  7,
18///  8,  8,  9,  9, 10, 10, 11, 11,
19/// 12, 12, 13, 13, 14, 14, 15, 15,
20/// 16, 16, 17, 17, 18, 18, 19, 19,
21/// 20, 20, 21, 21, 22, 22, 23, 23,
22/// 24, 24, 25, 25, 26, 26, 27, 27,
23/// 28, 28, 29, 29, 30, 30, 31, 31,
24///
25/// In the following layout, values are held disjointly, and num_rows=2 because
26/// the two occurrences of the same plane id are not in the same row
27///  0,  1,  2,  3,  4,  5,  6,  7,
28///  8,  9, 10, 11, 12, 13, 14, 15,
29/// 16, 17, 18, 19, 20, 21, 22, 23,
30/// 24, 25, 26, 27, 28, 29, 30, 31,
31///  0,  1,  2,  3,  4,  5,  6,  7,
32///  8,  9, 10, 11, 12, 13, 14, 15,
33/// 16, 17, 18, 19, 20, 21, 22, 23,
34/// 24, 25, 26, 27, 28, 29, 30, 31,
35pub struct RowWise<E: Numeric> {
36    pub vals: Array<E>,
37    #[cube(comptime)]
38    pub num_rows: usize,
39}
40
41#[cube]
42impl<E: Numeric> RowWise<E> {
43    pub fn new_filled(#[comptime] num_rows: usize, val: E) -> RowWise<E> {
44        let mut vals = Array::new(num_rows);
45        for i in 0..num_rows {
46            vals[i] = val;
47        }
48        RowWise::<E> { vals, num_rows }
49    }
50
51    pub fn fill(&mut self, val: E) {
52        for i in 0..self.num_rows {
53            self.vals[i] = val;
54        }
55    }
56
57    pub fn init_zero(&mut self) {
58        self.fill(E::from_int(0));
59    }
60
61    pub fn new_min_value(#[comptime] num_rows: usize) -> RowWise<E> {
62        Self::new_filled(num_rows, E::min_value())
63    }
64
65    pub fn new_zero(#[comptime] num_rows: usize) -> RowWise<E> {
66        Self::new_filled(num_rows, E::from_int(0))
67    }
68
69    pub fn copy_from(&mut self, other: &RowWise<E>) {
70        for i in 0..self.num_rows {
71            self.vals[i] = other.vals[i]
72        }
73    }
74
75    pub fn add(&self, other: &RowWise<E>) -> RowWise<E> {
76        let mut result = Array::new(self.num_rows);
77        for i in 0..self.num_rows {
78            result[i] = self.vals[i] + other.vals[i];
79        }
80        RowWise::<E> {
81            vals: result,
82            num_rows: self.num_rows,
83        }
84    }
85
86    pub fn add_inplace(&mut self, other: &RowWise<E>) {
87        for i in 0..self.num_rows {
88            self.vals[i] += other.vals[i];
89        }
90    }
91
92    pub fn mul(&self, other: &RowWise<E>) -> RowWise<E> {
93        let mut result = Array::new(self.num_rows);
94        for i in 0..self.num_rows {
95            result[i] = self.vals[i] * other.vals[i];
96        }
97        RowWise::<E> {
98            vals: result,
99            num_rows: self.num_rows,
100        }
101    }
102
103    pub fn mul_inplace(&mut self, other: &RowWise<E>) {
104        for i in 0..self.num_rows {
105            self.vals[i] *= other.vals[i];
106        }
107    }
108
109    pub fn max_inplace(&mut self, other: &RowWise<E>) {
110        for i in 0..self.num_rows {
111            self.vals[i] = max(self.vals[i], other.vals[i]);
112        }
113    }
114
115    pub fn replace_at(&mut self, i: usize, new_val: E) {
116        self.vals[i] = new_val;
117    }
118
119    pub fn cast_from<E2: Float>(row_wise: &RowWise<E>) -> RowWise<E2> {
120        let num_rows = row_wise.num_rows;
121        let mut vals = Array::new(num_rows);
122
123        for i in 0..num_rows {
124            vals[i] = E2::cast_from(row_wise.vals[i]);
125        }
126
127        RowWise::<E2> { vals, num_rows }
128    }
129}
130
131#[cube]
132impl<E: Float> RowWise<E> {
133    /// Per-row `e^(self - other)`.
134    pub fn exp_diff(&self, other: &RowWise<E>) -> RowWise<E> {
135        let mut vals = Array::new(self.num_rows);
136
137        for i in 0..self.num_rows {
138            vals[i] = (self.vals[i] - other.vals[i]).exp();
139        }
140
141        RowWise::<E> {
142            vals,
143            num_rows: self.num_rows,
144        }
145    }
146
147    /// `v -> 1/v` per row, with `v == 0` (fully-masked row) staying zero.
148    pub fn recip_inplace(&mut self) {
149        for i in 0..self.num_rows {
150            let row_val = self.vals[i];
151
152            let epsilon = E::new(FULLY_MASKED_ROW_THRESHOLD);
153            let not_masked = E::cast_from(row_val >= epsilon);
154            let safe_val = clamp_min(row_val, epsilon);
155            let recip = safe_val.recip();
156            self.vals[i] = not_masked * recip;
157        }
158    }
159}
cubek_std/tile/variants/row_wise.rs

cubek_std/tile/variants/
row_wise.rs
