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
81
82
83
84
85
86
87
88
89
90
//! Tile-level sparsity.
//!
//! Systolic hardware pays the same cost for a tile of zeros as for a tile of real
//! work, so the highest-value sparsity optimisation is to *skip whole tiles* that
//! are entirely zero. This module finds those tiles.
use crate::bf16::Bf16;
use crate::lattice::PaddedTileLattice;
/// A type that knows what its additive-identity "zero" is.
pub trait IsZero {
/// True if this value is the additive identity.
fn is_zero(&self) -> bool;
}
impl IsZero for f32 {
#[inline]
fn is_zero(&self) -> bool {
*self == 0.0
}
}
impl IsZero for f64 {
#[inline]
fn is_zero(&self) -> bool {
*self == 0.0
}
}
impl IsZero for i8 {
#[inline]
fn is_zero(&self) -> bool {
*self == 0
}
}
impl IsZero for i32 {
#[inline]
fn is_zero(&self) -> bool {
*self == 0
}
}
impl IsZero for Bf16 {
#[inline]
fn is_zero(&self) -> bool {
Bf16::is_zero(*self)
}
}
impl<T: IsZero> PaddedTileLattice<T> {
/// True if every element of tile `(tile_row, tile_col)` is zero.
pub fn is_tile_zero(&self, tile_row: usize, tile_col: usize) -> bool {
self.tile_slice(tile_row, tile_col)
.iter()
.all(|x| x.is_zero())
}
/// Count how many storage tiles are entirely zero.
pub fn count_zero_tiles(&self) -> usize {
self.iter_tiles()
.filter(|(_, _, slice)| slice.iter().all(|x| x.is_zero()))
.count()
}
/// Collect the `(tile_row, tile_col)` coordinates of every non-zero tile.
///
/// These are exactly the tiles a sparsity-aware kernel needs to feed through
/// the systolic array; the rest can be skipped.
pub fn nonzero_tile_coords(&self) -> Vec<(usize, usize)> {
self.iter_tiles()
.filter(|(_, _, slice)| !slice.iter().all(|x| x.is_zero()))
.map(|(r, c, _)| (r, c))
.collect()
}
/// Fraction of storage tiles that are entirely zero, in `0.0..=1.0`.
pub fn tile_sparsity(&self) -> f64 {
let total = self.num_tiles();
if total == 0 {
return 0.0;
}
self.count_zero_tiles() as f64 / total as f64
}
/// Fraction of storage tiles that contain at least one non-zero element.
pub fn tile_density(&self) -> f64 {
1.0 - self.tile_sparsity()
}
}