substrate_data/wire/frame/lanes/
packing.rs

1//! Packing and unpacking: groups through the processor's kernels, the rest one by one.
2
3use super::Backend;
4
5/// Lanes of a group: a 256-bit vector of 32-bit words.
6pub const LANES: usize = 8;
7/// Numbers a lane holds in a group: as many as its word has bits.
8pub(super) const ROWS: usize = u32::BITS as usize;
9/// Numbers in a group.
10pub const GROUP: usize = LANES * ROWS;
11/// The widest numbers packed in groups.
12pub const GROUP_MAX_BITS: u32 = u32::BITS;
13pub(super) const WORD_BYTES: usize = 4;
14const BYTE_BITS: usize = 8;
15
16/// A group's bytes, its numbers `width` bits each.
17#[must_use]
18pub const fn group_bytes(width: u32) -> usize {
19    LANES * width as usize * WORD_BYTES
20}
21
22/// The bytes `n` numbers `width` bits each take.
23#[must_use]
24pub fn packed_bytes(n: usize, width: u32) -> Option<usize> {
25    let (groups, rest) = split(n, width);
26    groups
27        .checked_mul(group_bytes(width))?
28        .checked_add(rest.checked_mul(width as usize)?.div_ceil(BYTE_BITS))
29}
30
31/// How many whole groups `n` numbers make, and how many go after them.
32fn split(n: usize, width: u32) -> (usize, usize) {
33    if width > GROUP_MAX_BITS {
34        (0, n)
35    } else {
36        (n / GROUP, n % GROUP)
37    }
38}
39
40/// The low 32 bits of `v` (a grouped number fits them: groups are only
41/// packed up to 32 bits wide).
42fn low_word(v: u64) -> u32 {
43    u32::try_from(v & u64::from(u32::MAX)).unwrap_or(u32::MAX)
44}
45
46/// The low 64 bits of `v`.
47fn low_u64(v: u128) -> u64 {
48    u64::try_from(v & u128::from(u64::MAX)).unwrap_or(u64::MAX)
49}
50
51/// A word's mask of `width` low bits.
52pub(super) fn mask(width: u32) -> u32 {
53    u32::MAX >> (u32::BITS - width)
54}
55
56/// Groups a kernel call is handed at a time: enough for the widest backend's
57/// stride, small enough to stay in the cache.
58const BATCH: usize = 4;
59
60/// Packs `items` (`n` of them, each fitting `width` bits) onto `out`.
61pub fn pack(out: &mut Vec<u8>, n: usize, width: u32, items: impl Iterator<Item = u64>) {
62    let (groups, rest) = split(n, width);
63    let start = out.len();
64    out.resize(start + packed_bytes(n, width).unwrap_or(0), 0);
65    let mut items = items;
66    let backend = Backend::active();
67    let size = group_bytes(width);
68    let mut batch = [[0u32; GROUP]; BATCH];
69    let mut at = start;
70    let mut left = groups;
71    while left > 0 {
72        let take = left.min(BATCH);
73        for group in &mut batch[..take] {
74            for (slot, v) in group.iter_mut().zip(items.by_ref()) {
75                *slot = low_word(v);
76            }
77        }
78        backend.pack_groups(&batch[..take], width, &mut out[at..at + take * size]);
79        at += take * size;
80        left -= take;
81    }
82    pack_sequential(&mut out[at..], rest, width, items);
83}
84
85/// Numbers one after another, little-endian bits.
86fn pack_sequential(out: &mut [u8], n: usize, width: u32, items: impl Iterator<Item = u64>) {
87    let mut acc: u128 = 0;
88    let mut have = 0u32;
89    let mut at = 0usize;
90    for v in items.take(n) {
91        acc |= u128::from(v) << have;
92        have += width;
93        while have >= u8::BITS {
94            out[at] = acc.to_le_bytes()[0];
95            at += 1;
96            acc >>= u8::BITS;
97            have -= u8::BITS;
98        }
99    }
100    if have > 0 {
101        out[at] = acc.to_le_bytes()[0];
102    }
103}
104
105/// Each of the `n` numbers in `bytes` (`width` bits each), in order, to `f`.
106pub fn unpack(bytes: &[u8], n: usize, width: u32, mut f: impl FnMut(u64)) {
107    let (groups, rest) = split(n, width);
108    let backend = Backend::active();
109    let size = group_bytes(width);
110    let mut batch = [[0u32; GROUP]; BATCH];
111    let mut at = 0;
112    let mut left = groups;
113    while left > 0 {
114        let take = left.min(BATCH);
115        backend.unpack_groups(&bytes[at..at + take * size], width, &mut batch[..take]);
116        for &v in batch[..take].as_flattened() {
117            f(u64::from(v));
118        }
119        at += take * size;
120        left -= take;
121    }
122    let tail = &bytes[groups * size..];
123    for i in 0..rest {
124        f(sequential(tail, i, width));
125    }
126}
127
128/// Number `i` of `bytes` (`n` numbers, `width` bits each) — `i < n`.
129#[must_use]
130pub fn get(bytes: &[u8], n: usize, width: u32, i: usize) -> u64 {
131    let (groups, _) = split(n, width);
132    if i >= groups * GROUP {
133        let tail = &bytes[groups * group_bytes(width)..];
134        return sequential(tail, i - groups * GROUP, width);
135    }
136    let (group, j) = (i / GROUP, i % GROUP);
137    let (lane, row) = (j % LANES, j / LANES);
138    let bit = row * width as usize;
139    let (word, shift) = (bit / u32::BITS as usize, bit % u32::BITS as usize);
140    let at = group * group_bytes(width);
141    let word_at = |k: usize| {
142        let o = at + (k * LANES + lane) * WORD_BYTES;
143        bytes.get(o..o + WORD_BYTES).map_or(0, |b| {
144            u64::from(u32::from_le_bytes(b.try_into().unwrap_or_default()))
145        })
146    };
147    let two = word_at(word)
148        | if shift + width as usize > u32::BITS as usize {
149            word_at(word + 1) << u32::BITS
150        } else {
151            0
152        };
153    (two >> shift) & (u64::MAX >> (u64::BITS - width))
154}
155
156/// Number `i` of numbers packed one after another.
157fn sequential(bytes: &[u8], i: usize, width: u32) -> u64 {
158    // A number spans at most 9 bytes (64 bits starting anywhere in a byte).
159    const SPAN: usize = 9;
160    let bit = i * width as usize;
161    let (byte, shift) = (bit / BYTE_BITS, bit % BYTE_BITS);
162    let mut window = [0u8; 16];
163    let end = (byte + SPAN).min(bytes.len());
164    if byte < end {
165        window[..end - byte].copy_from_slice(&bytes[byte..end]);
166    }
167    let mask = u128::from(u64::MAX) >> (u64::BITS - width);
168    low_u64((u128::from_le_bytes(window) >> shift) & mask)
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    fn roundtrip(n: usize, width: u32) {
176        let max = u64::MAX >> (u64::BITS - width);
177        let items: Vec<u64> = (0..n as u64)
178            .map(|i| i.wrapping_mul(0x9E37_79B9_7F4A_7C15) & max)
179            .collect();
180        let mut out = Vec::new();
181        pack(&mut out, n, width, items.iter().copied());
182        assert_eq!(Some(out.len()), packed_bytes(n, width));
183        let mut back = Vec::new();
184        unpack(&out, n, width, |v| back.push(v));
185        assert_eq!(back, items, "n {n}, width {width}");
186        for (i, v) in items.iter().enumerate() {
187            assert_eq!(
188                get(&out, n, width, i),
189                *v,
190                "get {i} of n {n}, width {width}"
191            );
192        }
193    }
194
195    #[test]
196    fn every_width_round_trips_in_groups_and_tail() {
197        for width in 1..=u64::BITS {
198            for n in [0, 1, 7, GROUP - 1, GROUP, GROUP + 3, 3 * GROUP + 100] {
199                roundtrip(n, width);
200            }
201        }
202    }
203}
substrate_data/wire/frame/lanes/packing.rs

substrate_data/wire/frame/lanes/
packing.rs
