Skip to content

Instantly share code, notes, and snippets.

@m1el
Last active May 13, 2024 10:51
Show Gist options
  • Save m1el/0bea9ee6717b5fcd5ffab768054f45b3 to your computer and use it in GitHub Desktop.
Save m1el/0bea9ee6717b5fcd5ffab768054f45b3 to your computer and use it in GitHub Desktop.
#![feature(portable_simd)]
/// A multiplier constant which shifts each bit of a byte into the lowest
/// bit of the corresponding byte in u64. Constant calculation:
/// `(0..8).fold(0, |acc, x| acc | (1 << ((8 + 1) * x)))`
const PACK_BYTES: u64 = 0x8040201008040201;
/// Retain only the lowest bit in each byte of u64, making it a valid [bool; 8].
/// Constant calculation: `(0..8).fold(0, |acc, x| acc | (1 << (8 * x)))`
const BOOL_MASK: u64 = 0x0101010101010101;
pub fn bitpack(value: [bool; 8]) -> u8 {
let bytes = unsafe { core::mem::transmute::<[bool; 8], [u8; 8]>(value) };
let inter = u64::from_le_bytes(bytes).wrapping_mul(PACK_BYTES);
(inter >> 56) as u8
}
pub fn unbitpack(value: u8) -> [bool; 8] {
let unpacked = (value as u64).wrapping_mul(PACK_BYTES) >> 7;
let bytes = (unpacked & BOOL_MASK).to_le_bytes();
unsafe { core::mem::transmute::<[u8; 8], [bool; 8]>(bytes) }
}
use std::simd::{LaneCount, Simd, SupportedLaneCount};
pub fn bitpack_simd<const L: usize>(value: [[bool; 8]; L]) -> [u8; L]
where LaneCount<L>: SupportedLaneCount
{
let arr = value.map(|x| {
let bytes = unsafe { core::mem::transmute::<[bool; 8], [u8; 8]>(x) };
u64::from_le_bytes(bytes)
});
let lanes = Simd::from_array(arr);
let bytes = ((lanes * Simd::splat(PACK_BYTES)) >> Simd::splat(56)).cast::<u8>();
bytes.to_array()
}
pub fn unbitpack_simd<const L: usize>(value: [u8; L]) -> [[bool; 8]; L]
where LaneCount<L>: SupportedLaneCount
{
let bytes = Simd::from_array(value);
let value = bytes.cast::<u64>() * Simd::splat(PACK_BYTES);
let value = (value >> Simd::splat(7)) & Simd::splat(BOOL_MASK);
value.to_array().map(|x| {
unsafe { core::mem::transmute::<[u8; 8], [bool; 8]>(x.to_le_bytes()) }
})
}
pub fn bitpack_64(value: [bool; 64]) -> u64 {
let bytes = unsafe { core::mem::transmute::<[bool; 64], [[u8; 8]; 8]>(value) };
let lanes = Simd::from_array(bytes.map(|x| u64::from_le_bytes(x)));
let bytes = ((lanes * Simd::splat(PACK_BYTES)) >> Simd::splat(56)).cast::<u8>();
u64::from_le_bytes(bytes.to_array())
}
pub fn unbitpack_64(value: u64) -> [bool; 64] {
let bytes = Simd::from_array(value.to_le_bytes());
let value = bytes.cast::<u64>() * Simd::splat(PACK_BYTES);
let value = (value >> Simd::splat(7)) & Simd::splat(BOOL_MASK);
let bytes = value.to_array().map(|p| p.to_le_bytes());
unsafe { core::mem::transmute::<[[u8; 8]; 8], [bool; 64]>(bytes) }
}
pub fn main() {
for b in 0..=255_u8 {
assert!(b == bitpack(unbitpack(b)));
let big = u64::from_le_bytes([b; 8]);
assert!(big == bitpack_64(unbitpack_64(big)));
assert!(big.to_le_bytes() == bitpack_simd(unbitpack_simd(big.to_le_bytes())));
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment