substrate_data/wire/frame/lanes/
packing.rs

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

substrate_data/wire/frame/lanes/
packing.rs
