logicaffeine_compile/spatial/decode_engine/
lane.rs

1//! The block-scale lane: a weight block's integer sums as the array and the activation side deliver them, turned into
2//! the block's f32 value and added into its row, exactly as [`super::numerics::dot`] does under
3//! [`ActScaling::PerBlock`](super::numerics::ActScaling::PerBlock).
4//!
5//! The array never sees a signed code. A weight code `q` (a decoder's unsigned code, one or two 4-bit digits) meets
6//! an activation word in excess form: a one-word code `a` as `a + 128`, a two-word code's signed high byte as
7//! `hi + 128` and its unsigned low byte as itself. Per 16-element unit the array delivers `P′ = Σ q·a′` for each word,
8//! the weight image carries `Q = Σ q`, and the activation side carries `S′ = Σ a′`. The unit's sums against the signed
9//! codes follow exactly:
10//!
11//! `Σ (q − z)·a = P′ − 128·Q − z·S′ + 128·16·z` and `Σ a = S′ − 128·16` for an excess word, `P′ − z·S′` and `S′`
12//! for the low byte, the high word weighted by 256;
13//!
14//! the coefficient sum is `Σ_u c_u·Σ(q − z)·a` and the offset sum `Σ_u n_u·Σ a`. The lane multiplies them by the
15//! block's scales times the activation scale, sums the two terms exactly, rounds once and adds into the row.
16
17use super::numerics::{exact_sum, fadd, ActRow, ActScale, Scale, Words};
18use super::rtl::weights::{Units, UNIT};
19
20/// What the array and the activation side hand a lane for one unit of one word.
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub struct UnitSums {
23    /// `Σ q·a′` over the unit, the digits recombined.
24    pub products: i64,
25    /// `Σ a′` over the unit.
26    pub act: i64,
27}
28
29/// One word's unit sums: whether the word crossed in excess form, and its weight in the code (1, or 256 for a
30/// two-word code's high byte).
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct WordSums {
33    pub units: Vec<UnitSums>,
34    pub excess: bool,
35    pub weight: i64,
36}
37
38/// One weight block as a lane reads it.
39#[derive(Clone, Debug, PartialEq, Eq)]
40pub struct LaneBlock {
41    /// Per unit: the sub-block multiplier, the offset multiplier, and `Σ q` from the weight image.
42    pub c: Vec<i64>,
43    pub n: Vec<i64>,
44    pub q: Vec<i64>,
45    /// The format's code offset `z`.
46    pub offset: i64,
47    pub a: Scale,
48    pub b: Scale,
49    pub act: ActScale,
50    pub words: Vec<WordSums>,
51}
52
53/// The block's two exact sums: `Σ coef·a` and `Σ offs·a`.
54pub fn block_sums(blk: &LaneBlock) -> (i128, i128) {
55    let unit = UNIT as i64;
56    let (mut ia, mut ib) = (0i128, 0i128);
57    for w in &blk.words {
58        let excess = if w.excess { 128 } else { 0 };
59        for (u, s) in w.units.iter().enumerate() {
60            let products = s.products - excess * blk.q[u] - blk.offset * s.act + excess * unit * blk.offset;
61            let act = s.act - excess * unit;
62            ia += (w.weight * blk.c[u]) as i128 * products as i128;
63            ib += (w.weight * blk.n[u]) as i128 * act as i128;
64        }
65    }
66    (ia, ib)
67}
68
69/// The block's value rounded once and added into `acc`.
70pub fn lane(blk: &LaneBlock, acc: f32) -> f32 {
71    let (ia, ib) = block_sums(blk);
72    let mut terms = [(0i128, 0i32); 2];
73    let mut count = 0;
74    for (s, sum) in [(blk.a, ia), (blk.b, ib)] {
75        if s.sig != 0 && sum != 0 {
76            terms[count] = (s.sig as i128 * blk.act.m as i128 * sum, s.exp + blk.act.e);
77            count += 1;
78        }
79    }
80    fadd(acc, exact_sum(&terms[..count]))
81}
82
83/// What the array, the weight image and the activation side deliver a lane for block `b` of a row: the row's decoded
84/// units `units` (the block's own), its block scales `a`, `b_scale`, and the block's activations of `act` in `words`.
85pub fn deliver(units: &Units, scales: (Scale, Scale), offset: u32, act: &ActRow, words: Words, b: usize) -> LaneBlock {
86    let len = units.code.len();
87    let start = b * len;
88    let count = len / UNIT;
89    let unit_sums = |a: &dyn Fn(usize) -> i64| -> Vec<UnitSums> {
90        (0..count)
91            .map(|u| {
92                let range = u * UNIT..(u + 1) * UNIT;
93                UnitSums {
94                    products: range.clone().map(|i| units.code[i] as i64 * a(start + i)).sum(),
95                    act: range.map(|i| a(start + i)).sum(),
96                }
97            })
98            .collect()
99    };
100    let code = |i: usize| act.codes[i] as i64;
101    let word_sums = match words {
102        Words::One => vec![WordSums { units: unit_sums(&|i| code(i) + 128), excess: true, weight: 1 }],
103        Words::Two => vec![
104            WordSums { units: unit_sums(&|i| code(i).div_euclid(256) + 128), excess: true, weight: 256 },
105            WordSums { units: unit_sums(&|i| code(i).rem_euclid(256)), excess: false, weight: 1 },
106        ],
107    };
108    LaneBlock {
109        c: units.c.iter().map(|&v| v as i64).collect(),
110        n: units.n.iter().map(|&v| v as i64).collect(),
111        q: (0..count).map(|u| units.code[u * UNIT..(u + 1) * UNIT].iter().map(|&q| q as i64).sum()).collect(),
112        offset: offset as i64,
113        a: scales.0,
114        b: scales.1,
115        act: act.scales[b],
116        words: word_sums,
117    }
118}
119
120// ── The per-32 lane ─────────────────────────────────────────────────────────────────────────────────
121
122/// Elements a per-32 activation scale covers: a sub-block of the lane.
123pub const SUB_BLOCK: usize = 32;
124
125/// What the array, the weight image and the activation side deliver the per-32 lane for the 32-element sub-block `s`
126/// of a row's weight block `b`: the block's decoded units `units` (the sub-block's two among them), the block's scales,
127/// and the sub-block's activations of `act`, scaled per 32 elements, in `words`.
128pub fn deliver32(units: &Units, scales: (Scale, Scale), offset: u32, act: &ActRow, words: Words, b: usize, s: usize) -> LaneBlock {
129    let block = deliver(units, scales, offset, act, words, b);
130    let first = s * SUB_BLOCK / UNIT;
131    let count = SUB_BLOCK / UNIT;
132    let start = b * units.code.len() + s * SUB_BLOCK;
133    LaneBlock {
134        c: block.c[first..first + count].to_vec(),
135        n: block.n[first..first + count].to_vec(),
136        q: block.q[first..first + count].to_vec(),
137        offset: block.offset,
138        a: block.a,
139        b: block.b,
140        act: act.scales[start / SUB_BLOCK],
141        words: block.words.into_iter().map(|w| WordSums { units: w.units[first..first + count].to_vec(), ..w }).collect(),
142    }
143}
144
145/// A group of the row — [`super::numerics::GROUP`] elements, its sub-blocks `subs` — rounded once and added into
146/// `acc`: every sub-block's two exact sums times its block's scales and its own activation scale, the group's terms
147/// summed exactly, as [`super::numerics::dot`] does under [`ActScaling::Sub32`](super::numerics::ActScaling::Sub32).
148pub fn group(subs: &[LaneBlock], acc: f32) -> f32 {
149    let mut terms = Vec::with_capacity(2 * subs.len());
150    for blk in subs {
151        let (ia, ib) = block_sums(blk);
152        for (s, sum) in [(blk.a, ia), (blk.b, ib)] {
153            if s.sig != 0 && sum != 0 {
154                terms.push((s.sig as i128 * blk.act.m as i128 * sum, s.exp + blk.act.e));
155            }
156        }
157    }
158    fadd(acc, exact_sum(&terms))
159}
logicaffeine_compile/spatial/decode_engine/lane.rs

logicaffeine_compile/spatial/decode_engine/
lane.rs
