systile 0.1.0

A TPU-native tiled tensor data structure: the Systolic Tile Lattice. Padding-aware, sublane/lane laid out, bf16/int8 first, with a CPU reference simulator of systolic dataflow.
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
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()
    }
}