#![cfg_attr(
all(simd_popcnt_have_sve, any(target_feature = "sve", feature = "std")),
feature(stdarch_aarch64_sve)
)]
#![cfg_attr(
not(any(
test,
all(
feature = "std",
any(target_arch = "x86", target_arch = "x86_64"),
not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")),
),
all(
feature = "std",
target_arch = "aarch64",
simd_popcnt_have_sve,
not(target_feature = "sve"),
),
)),
no_std
)]
#[cfg(target_arch = "aarch64")]
use core::arch::aarch64::*;
#[cfg(target_arch = "x86")]
#[allow(unused_imports)]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
#[allow(unused_imports)]
use core::arch::x86_64::*;
#[cfg(all(
target_arch = "aarch64",
simd_popcnt_have_sve,
feature = "std",
not(target_feature = "sve")
))]
use std::arch::is_aarch64_feature_detected;
#[must_use]
#[inline]
pub fn popcnt(bytes: &[u8]) -> u64 {
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
popcnt_x86(bytes)
}
#[cfg(target_arch = "aarch64")]
{
popcnt_aarch64(bytes)
}
#[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
{
popcnt_scalar(bytes)
}
}
pub trait PopcntExt {
#[must_use]
fn popcnt(&self) -> u64;
}
macro_rules! impl_popcnt_ext {
($($t:ty),+ $(,)?) => {$(
impl PopcntExt for [$t] {
#[inline]
fn popcnt(&self) -> u64 {
let bytes = unsafe {
core::slice::from_raw_parts(
self.as_ptr().cast::<u8>(),
core::mem::size_of_val(self),
)
};
popcnt(bytes)
}
}
)+};
}
impl_popcnt_ext!(
u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize
);
#[inline]
fn tail_u64(rem: &[u8]) -> u64 {
let mut buf = [0u8; 8];
buf[..rem.len()].copy_from_slice(rem);
u64::from_ne_bytes(buf)
}
macro_rules! popcnt_scalar_loop {
($bytes:expr) => {{
let mut cnt = 0u64;
let (chunks, rem) = $bytes.as_chunks::<8>();
for chunk in chunks {
cnt += u64::from_ne_bytes(*chunk).count_ones() as u64;
}
if !rem.is_empty() {
cnt += tail_u64(rem).count_ones() as u64;
}
cnt
}};
}
#[allow(dead_code)]
#[inline]
fn popcnt_scalar(bytes: &[u8]) -> u64 {
popcnt_scalar_loop!(bytes)
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[inline]
fn popcnt_x86(bytes: &[u8]) -> u64 {
#[cfg(target_feature = "avx512vpopcntdq")]
{
if bytes.len() >= 40 {
unsafe { popcnt_avx512(bytes) }
} else {
popcnt_scalar_static(bytes)
}
}
#[cfg(all(target_feature = "avx2", not(target_feature = "avx512vpopcntdq")))]
{
let mut cnt = 0u64;
let mut rest = bytes;
if bytes.len() >= 512 {
let n = bytes.len() / 32 * 32;
cnt += unsafe { popcnt_avx2(&bytes[..n]) };
rest = &bytes[n..];
}
cnt + popcnt_scalar_static(rest)
}
#[cfg(all(
not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")),
feature = "std"
))]
{
popcnt_x86_runtime(bytes)
}
#[cfg(all(
not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")),
not(feature = "std")
))]
{
popcnt_scalar_static(bytes)
}
}
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
any(
target_feature = "avx2",
target_feature = "avx512vpopcntdq",
not(feature = "std")
)
))]
#[inline]
fn popcnt_scalar_static(bytes: &[u8]) -> u64 {
#[cfg(target_feature = "popcnt")]
{
unsafe { popcnt_scalar_hw(bytes) }
}
#[cfg(not(target_feature = "popcnt"))]
{
popcnt_scalar(bytes)
}
}
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")),
feature = "std"
))]
#[inline]
fn has_avx512() -> bool {
use core::sync::atomic::{AtomicI32, Ordering};
static HAS_AVX512: AtomicI32 = AtomicI32::new(-1);
let cached = HAS_AVX512.load(Ordering::Relaxed);
if cached != -1 {
return cached != 0;
}
let v = (is_x86_feature_detected!("avx512f")
&& is_x86_feature_detected!("avx512bw")
&& is_x86_feature_detected!("avx512vpopcntdq")) as i32;
HAS_AVX512.store(v, Ordering::Relaxed);
v != 0
}
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")),
feature = "std"
))]
#[inline]
fn popcnt_x86_runtime(bytes: &[u8]) -> u64 {
if bytes.len() >= 40 && has_avx512() {
return unsafe { popcnt_avx512(bytes) };
}
let mut cnt = 0u64;
let mut rest = bytes;
if bytes.len() >= 512 && is_x86_feature_detected!("avx2") {
let n = bytes.len() / 32 * 32;
cnt += unsafe { popcnt_avx2(&bytes[..n]) };
rest = &bytes[n..];
}
cnt += if is_x86_feature_detected!("popcnt") {
unsafe { popcnt_scalar_hw(rest) }
} else {
popcnt_scalar(rest)
};
cnt
}
#[allow(dead_code)]
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "popcnt")]
#[inline]
fn popcnt_scalar_hw(bytes: &[u8]) -> u64 {
popcnt_scalar_loop!(bytes) }
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
not(target_feature = "avx512vpopcntdq"),
any(target_feature = "avx2", feature = "std")
))]
#[target_feature(enable = "avx2")]
#[inline]
fn csa256(a: __m256i, b: __m256i, c: __m256i) -> (__m256i, __m256i) {
let u = _mm256_xor_si256(a, b);
let h = _mm256_or_si256(_mm256_and_si256(a, b), _mm256_and_si256(u, c));
let l = _mm256_xor_si256(u, c);
(h, l)
}
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
not(target_feature = "avx512vpopcntdq"),
any(target_feature = "avx2", feature = "std")
))]
#[target_feature(enable = "avx2")]
#[inline]
fn popcnt256(v: __m256i) -> __m256i {
let lookup1 = _mm256_setr_epi8(
4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8, 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7,
7, 8,
);
let lookup2 = _mm256_setr_epi8(
4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0, 4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1,
1, 0,
);
let low_mask = _mm256_set1_epi8(0x0f);
let lo = _mm256_and_si256(v, low_mask);
let hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask);
let popcnt1 = _mm256_shuffle_epi8(lookup1, lo);
let popcnt2 = _mm256_shuffle_epi8(lookup2, hi);
_mm256_sad_epu8(popcnt1, popcnt2)
}
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
not(target_feature = "avx512vpopcntdq"),
any(target_feature = "avx2", feature = "std")
))]
#[target_feature(enable = "avx2")]
#[inline]
#[rustfmt::skip]
fn popcnt_avx2(bytes: &[u8]) -> u64 {
let zero = _mm256_setzero_si256();
let mut cnt = zero;
let mut ones = zero;
let mut twos = zero;
let mut fours = zero;
let mut eights = zero;
let mut twos_a;
let mut twos_b;
let mut fours_a;
let mut fours_b;
let mut eights_a;
let mut eights_b;
let mut sixteens;
let (blocks, tail) = bytes.as_chunks::<512>();
for chunk in blocks {
let p = chunk.as_ptr().cast::<__m256i>();
unsafe {
(twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(0)), _mm256_loadu_si256(p.add(1)));
(twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(2)), _mm256_loadu_si256(p.add(3)));
(fours_a, twos) = csa256(twos, twos_a, twos_b);
(twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(4)), _mm256_loadu_si256(p.add(5)));
(twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(6)), _mm256_loadu_si256(p.add(7)));
(fours_b, twos) = csa256(twos, twos_a, twos_b);
(eights_a, fours) = csa256(fours, fours_a, fours_b);
(twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(8)), _mm256_loadu_si256(p.add(9)));
(twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(10)), _mm256_loadu_si256(p.add(11)));
(fours_a, twos) = csa256(twos, twos_a, twos_b);
(twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(12)), _mm256_loadu_si256(p.add(13)));
(twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(14)), _mm256_loadu_si256(p.add(15)));
(fours_b, twos) = csa256(twos, twos_a, twos_b);
(eights_b, fours) = csa256(fours, fours_a, fours_b);
(sixteens, eights) = csa256(eights, eights_a, eights_b);
cnt = _mm256_add_epi64(cnt, popcnt256(sixteens));
}
}
cnt = _mm256_slli_epi64(cnt, 4);
cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(eights), 3));
cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(fours), 2));
cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(twos), 1));
cnt = _mm256_add_epi64(cnt, popcnt256(ones));
let (vecs, _) = tail.as_chunks::<32>();
for chunk in vecs {
let v = unsafe { _mm256_loadu_si256(chunk.as_ptr().cast::<__m256i>()) };
cnt = _mm256_add_epi64(cnt, popcnt256(v));
}
let lanes: [u64; 4] = unsafe { core::mem::transmute(cnt) };
lanes[0] + lanes[1] + lanes[2] + lanes[3]
}
#[cfg(all(
any(target_arch = "x86", target_arch = "x86_64"),
any(
all(not(target_feature = "avx2"), feature = "std"),
target_feature = "avx512vpopcntdq"
)
))]
#[target_feature(enable = "avx512f,avx512bw,avx512vpopcntdq")]
#[inline]
fn popcnt_avx512(bytes: &[u8]) -> u64 {
let mut cnt0 = _mm512_setzero_si512();
let (blocks, tail256) = bytes.as_chunks::<256>();
if !blocks.is_empty() {
let mut cnt1 = _mm512_setzero_si512();
let mut cnt2 = _mm512_setzero_si512();
let mut cnt3 = _mm512_setzero_si512();
for chunk in blocks {
let p = chunk.as_ptr();
unsafe {
let v0 = _mm512_loadu_si512(p.add(0).cast());
let v1 = _mm512_loadu_si512(p.add(64).cast());
let v2 = _mm512_loadu_si512(p.add(128).cast());
let v3 = _mm512_loadu_si512(p.add(192).cast());
cnt0 = _mm512_add_epi64(cnt0, _mm512_popcnt_epi64(v0));
cnt1 = _mm512_add_epi64(cnt1, _mm512_popcnt_epi64(v1));
cnt2 = _mm512_add_epi64(cnt2, _mm512_popcnt_epi64(v2));
cnt3 = _mm512_add_epi64(cnt3, _mm512_popcnt_epi64(v3));
}
}
cnt0 = _mm512_add_epi64(cnt0, cnt1);
cnt2 = _mm512_add_epi64(cnt2, cnt3);
cnt0 = _mm512_add_epi64(cnt0, cnt2);
}
let (vecs, tail64) = tail256.as_chunks::<64>();
for chunk in vecs {
let v = unsafe { _mm512_loadu_si512(chunk.as_ptr().cast()) };
cnt0 = _mm512_add_epi64(cnt0, _mm512_popcnt_epi64(v));
}
if !tail64.is_empty() {
let len = tail64.len();
let mask = (u64::MAX >> (64 - len)) as __mmask64;
unsafe {
let v = _mm512_maskz_loadu_epi8(mask, tail64.as_ptr().cast());
cnt0 = _mm512_add_epi64(cnt0, _mm512_popcnt_epi64(v));
}
}
_mm512_reduce_add_epi64(cnt0) as u64
}
#[cfg(target_arch = "aarch64")]
#[inline]
fn popcnt_aarch64(bytes: &[u8]) -> u64 {
#[cfg(all(target_feature = "sve", simd_popcnt_have_sve))]
{
unsafe { popcnt_arm_sve(bytes) }
}
#[cfg(not(all(target_feature = "sve", simd_popcnt_have_sve)))]
{
popcnt_neon(bytes)
}
}
#[cfg(all(
target_arch = "aarch64",
not(all(target_feature = "sve", simd_popcnt_have_sve))
))]
#[inline]
fn vpadalq(sum: uint64x2_t, t: uint8x16_t) -> uint64x2_t {
unsafe { vpadalq_u32(sum, vpaddlq_u16(vpaddlq_u8(t))) }
}
#[cfg(all(
target_arch = "aarch64",
not(all(target_feature = "sve", simd_popcnt_have_sve))
))]
#[inline]
fn popcnt_neon(bytes: &[u8]) -> u64 {
#[cfg(all(simd_popcnt_have_sve, feature = "std"))]
if is_aarch64_feature_detected!("sve") {
return unsafe { popcnt_arm_sve(bytes) };
}
const CHUNK: usize = 64;
let mut cnt = 0u64;
let iters = bytes.len() / CHUNK;
let ptr = bytes.as_ptr();
if iters > 0 {
unsafe {
let mut sum = vdupq_n_u64(0);
let zero = vdupq_n_u8(0);
let mut i = 0usize;
while i < iters {
let mut t0 = zero;
let mut t1 = zero;
let mut t2 = zero;
let mut t3 = zero;
let limit = (i + 31).min(iters);
while i < limit {
let input = vld1q_u8_x4(ptr.add(i * CHUNK));
t0 = vaddq_u8(t0, vcntq_u8(input.0));
t1 = vaddq_u8(t1, vcntq_u8(input.1));
t2 = vaddq_u8(t2, vcntq_u8(input.2));
t3 = vaddq_u8(t3, vcntq_u8(input.3));
i += 1;
}
sum = vpadalq(sum, t0);
sum = vpadalq(sum, t1);
sum = vpadalq(sum, t2);
sum = vpadalq(sum, t3);
}
let mut tmp = [0u64; 2];
vst1q_u64(tmp.as_mut_ptr(), sum);
cnt += tmp[0] + tmp[1];
}
}
let rest = &bytes[iters * CHUNK..];
cnt += popcnt_scalar_loop!(rest);
cnt
}
#[cfg(all(
target_arch = "aarch64",
simd_popcnt_have_sve,
any(target_feature = "sve", feature = "std")
))]
#[target_feature(enable = "sve")]
#[inline]
fn popcnt_arm_sve(bytes: &[u8]) -> u64 {
unsafe {
let mut i = 0usize;
let mut vcnt0 = svdup_n_u64(0);
let vl = svcntb() as usize; let ptr = bytes.as_ptr();
let len = bytes.len();
if i + vl * 4 <= len {
let mut vcnt1 = svdup_n_u64(0);
let mut vcnt2 = svdup_n_u64(0);
let mut vcnt3 = svdup_n_u64(0);
loop {
let v0 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i)));
let v1 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl)));
let v2 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl * 2)));
let v3 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl * 3)));
vcnt0 = svadd_u64_x(svptrue_b64(), vcnt0, svcnt_u64_x(svptrue_b64(), v0));
vcnt1 = svadd_u64_x(svptrue_b64(), vcnt1, svcnt_u64_x(svptrue_b64(), v1));
vcnt2 = svadd_u64_x(svptrue_b64(), vcnt2, svcnt_u64_x(svptrue_b64(), v2));
vcnt3 = svadd_u64_x(svptrue_b64(), vcnt3, svcnt_u64_x(svptrue_b64(), v3));
i += vl * 4;
if i + vl * 4 > len {
break;
}
}
vcnt0 = svadd_u64_x(svptrue_b64(), vcnt0, vcnt1);
vcnt2 = svadd_u64_x(svptrue_b64(), vcnt2, vcnt3);
vcnt0 = svadd_u64_x(svptrue_b64(), vcnt0, vcnt2);
}
let mut pg = svwhilelt_b8_u64(i as u64, len as u64);
while svptest_any(svptrue_b8(), pg) {
let v = svreinterpret_u64_u8(svld1_u8(pg, ptr.add(i)));
vcnt0 = svadd_u64_x(svptrue_b64(), vcnt0, svcnt_u64_x(svptrue_b64(), v));
i += vl;
pg = svwhilelt_b8_u64(i as u64, len as u64);
}
svaddv_u64(svptrue_b64(), vcnt0)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn reference(bytes: &[u8]) -> u64 {
bytes.iter().map(|b| b.count_ones() as u64).sum()
}
fn popcnt64_bitwise(x: u64) -> u64 {
const M1: u64 = 0x5555555555555555;
const M2: u64 = 0x3333333333333333;
const M4: u64 = 0x0F0F0F0F0F0F0F0F;
const H01: u64 = 0x0101010101010101;
let x = x - ((x >> 1) & M1);
let x = (x & M2) + ((x >> 2) & M2);
let x = (x + (x >> 4)) & M4;
x.wrapping_mul(H01) >> 56
}
#[test]
fn empty() {
assert_eq!(popcnt(&[]), 0);
}
#[test]
fn all_ones() {
for &size in &[
0, 1, 7, 8, 31, 32, 39, 40, 63, 64, 255, 256, 511, 512, 4095, 4096, 65537,
] {
let bytes = vec![0xFFu8; size];
assert_eq!(popcnt(&bytes), size as u64 * 8, "size={size}");
}
}
#[test]
fn all_zeros() {
let bytes = vec![0u8; 65536];
assert_eq!(popcnt(&bytes), 0);
}
#[test]
fn single_bits() {
for bit in 0u64..64 {
let val = 1u64 << bit;
assert_eq!(popcnt(&val.to_le_bytes()), 1, "bit={bit}");
}
}
#[test]
fn ext_trait_widths() {
let u8s: &[u8] = &[0xFF, 0x0F, 0x00, 0xAB, 0x01];
assert_eq!(
u8s.popcnt(),
u8s.iter().map(|x| x.count_ones() as u64).sum()
);
let u16s: &[u16] = &[0xFFFF, 0x0F0F, 0x1234, 0];
assert_eq!(
u16s.popcnt(),
u16s.iter().map(|x| x.count_ones() as u64).sum()
);
let u32s: &[u32] = &[u32::MAX, 0, 0x8000_0001];
assert_eq!(
u32s.popcnt(),
u32s.iter().map(|x| x.count_ones() as u64).sum()
);
let u64s: &[u64] = &[u64::MAX, 0x0F0F_0F0F_0F0F_0F0F, 0];
assert_eq!(
u64s.popcnt(),
u64s.iter().map(|x| x.count_ones() as u64).sum()
);
let i32s = [-1i32, 0, 1, i32::MIN];
assert_eq!(
i32s.popcnt(),
i32s.iter().map(|x| x.count_ones() as u64).sum()
);
assert_eq!([u128::MAX, 0].popcnt(), 128);
}
#[test]
fn pseudorandom_all_sizes() {
let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
let mut next = || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
state
};
const MAX_SIZE: usize = 4695;
const MAX_OFF: usize = 7;
let mut bytes = vec![0u8; MAX_SIZE + MAX_OFF + 1];
for b in bytes.iter_mut() {
*b = (next() & 0xFF) as u8;
}
let sizes =
(0usize..=600).chain([1023, 1024, 1025, 2048, 4095, 4096, 4097, 4608, MAX_SIZE]);
for size in sizes {
for &off in &[0usize, 1, 3, MAX_OFF] {
let slice = &bytes[off..off + size];
assert_eq!(popcnt(slice), reference(slice), "size={size} off={off}");
}
}
}
#[test]
fn suffix_sweep() {
let size = std::env::var("SIMD_POPCNT_TEST_SIZE")
.ok()
.and_then(|s| s.parse::<usize>().ok())
.unwrap_or(20_000);
let ones = vec![0xFFu8; size];
check_all_suffixes(&ones);
let mut state: u64 = 0x2545_F491_4F6C_DD1D;
let mut bytes = vec![0u8; size];
for b in bytes.iter_mut() {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
*b = state as u8;
}
check_all_suffixes(&bytes);
}
fn check_all_suffixes(bytes: &[u8]) {
let total: u64 = bytes.iter().map(|&b| popcnt64_bitwise(b as u64)).sum();
let mut prefix = 0u64; for (i, &byte) in bytes.iter().enumerate() {
assert_eq!(popcnt(&bytes[i..]), total - prefix, "suffix at offset {i}");
prefix += popcnt64_bitwise(byte as u64);
}
assert_eq!(popcnt(&bytes[bytes.len()..]), 0);
}
}