1#![cfg_attr(simd_popcnt_have_sve, feature(stdarch_aarch64_sve))]
26
27#[cfg(target_arch = "aarch64")]
28use core::arch::aarch64::*;
29#[cfg(target_arch = "x86")]
30use core::arch::x86::*;
31#[cfg(target_arch = "x86_64")]
32use core::arch::x86_64::*;
33#[cfg(all(target_arch = "aarch64", simd_popcnt_have_sve))]
34use std::arch::is_aarch64_feature_detected;
35
36#[must_use]
51#[inline]
52pub fn popcnt(bytes: &[u8]) -> u64 {
53 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
54 {
55 popcnt_x86(bytes)
56 }
57
58 #[cfg(target_arch = "aarch64")]
59 {
60 popcnt_aarch64(bytes)
61 }
62
63 #[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
64 {
65 popcnt_scalar(bytes)
66 }
67}
68
69pub trait PopcntExt {
86 #[must_use]
88 fn popcnt(&self) -> u64;
89}
90
91macro_rules! impl_popcnt_ext {
95 ($($t:ty),+ $(,)?) => {$(
96 impl PopcntExt for [$t] {
97 #[inline]
98 fn popcnt(&self) -> u64 {
99 let bytes = unsafe {
103 core::slice::from_raw_parts(
104 self.as_ptr().cast::<u8>(),
105 core::mem::size_of_val(self),
106 )
107 };
108 popcnt(bytes)
109 }
110 }
111 )+};
112}
113
114impl_popcnt_ext!(
115 u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize
116);
117
118#[inline]
126fn tail_u64(rem: &[u8]) -> u64 {
127 let mut buf = [0u8; 8];
128 buf[..rem.len()].copy_from_slice(rem);
129 u64::from_ne_bytes(buf)
130}
131
132macro_rules! popcnt_scalar_loop {
138 ($bytes:expr) => {{
139 let mut cnt = 0u64;
140 let (chunks, rem) = $bytes.as_chunks::<8>();
141 for chunk in chunks {
142 cnt += u64::from_ne_bytes(*chunk).count_ones() as u64;
143 }
144 if !rem.is_empty() {
145 cnt += tail_u64(rem).count_ones() as u64;
146 }
147 cnt
148 }};
149}
150
151#[allow(dead_code)]
153fn popcnt_scalar(bytes: &[u8]) -> u64 {
154 popcnt_scalar_loop!(bytes)
155}
156
157#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
162fn popcnt_x86(bytes: &[u8]) -> u64 {
163 #[cfg(target_feature = "avx512vpopcntdq")]
165 {
166 if bytes.len() >= 40 {
168 unsafe { popcnt_avx512(bytes) }
169 } else {
170 popcnt_scalar_static(bytes)
171 }
172 }
173
174 #[cfg(all(target_feature = "avx2", not(target_feature = "avx512vpopcntdq")))]
176 {
177 let mut cnt = 0u64;
178 let mut rest = bytes;
179 if bytes.len() >= 512 {
181 let n = bytes.len() / 32 * 32;
182 cnt += unsafe { popcnt_avx2(&bytes[..n]) };
183 rest = &bytes[n..];
184 }
185 cnt + popcnt_scalar_static(rest)
186 }
187
188 #[cfg(not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")))]
190 {
191 popcnt_x86_runtime(bytes)
192 }
193}
194
195#[cfg(all(
198 any(target_arch = "x86", target_arch = "x86_64"),
199 any(target_feature = "avx2", target_feature = "avx512vpopcntdq")
200))]
201#[inline]
202fn popcnt_scalar_static(bytes: &[u8]) -> u64 {
203 #[cfg(target_feature = "popcnt")]
204 {
205 unsafe { popcnt_scalar_hw(bytes) }
207 }
208 #[cfg(not(target_feature = "popcnt"))]
209 {
210 popcnt_scalar(bytes)
211 }
212}
213
214#[cfg(all(
218 any(target_arch = "x86", target_arch = "x86_64"),
219 not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq"))
220))]
221fn has_avx512() -> bool {
222 use core::sync::atomic::{AtomicI32, Ordering};
223 static HAS_AVX512: AtomicI32 = AtomicI32::new(-1);
224 let cached = HAS_AVX512.load(Ordering::Relaxed);
225 if cached != -1 {
226 return cached != 0;
227 }
228 let v = (is_x86_feature_detected!("avx512f")
229 && is_x86_feature_detected!("avx512bw")
230 && is_x86_feature_detected!("avx512vpopcntdq")) as i32;
231 HAS_AVX512.store(v, Ordering::Relaxed);
232 v != 0
233}
234
235#[cfg(all(
238 any(target_arch = "x86", target_arch = "x86_64"),
239 not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq"))
240))]
241fn popcnt_x86_runtime(bytes: &[u8]) -> u64 {
242 if bytes.len() >= 40 && has_avx512() {
244 return unsafe { popcnt_avx512(bytes) };
245 }
246
247 let mut cnt = 0u64;
248 let mut rest = bytes;
249
250 if bytes.len() >= 512 && is_x86_feature_detected!("avx2") {
252 let n = bytes.len() / 32 * 32;
253 cnt += unsafe { popcnt_avx2(&bytes[..n]) };
254 rest = &bytes[n..];
255 }
256
257 cnt += if is_x86_feature_detected!("popcnt") {
262 unsafe { popcnt_scalar_hw(rest) }
263 } else {
264 popcnt_scalar(rest)
265 };
266
267 cnt
268}
269
270#[allow(dead_code)]
274#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
275#[target_feature(enable = "popcnt")]
276fn popcnt_scalar_hw(bytes: &[u8]) -> u64 {
277 popcnt_scalar_loop!(bytes) }
279
280#[cfg(all(
285 any(target_arch = "x86", target_arch = "x86_64"),
286 not(target_feature = "avx512vpopcntdq")
287))]
288#[target_feature(enable = "avx2")]
289#[inline]
290fn csa256(a: __m256i, b: __m256i, c: __m256i) -> (__m256i, __m256i) {
291 let u = _mm256_xor_si256(a, b);
292 let h = _mm256_or_si256(_mm256_and_si256(a, b), _mm256_and_si256(u, c));
293 let l = _mm256_xor_si256(u, c);
294 (h, l)
295}
296
297#[cfg(all(
300 any(target_arch = "x86", target_arch = "x86_64"),
301 not(target_feature = "avx512vpopcntdq")
302))]
303#[target_feature(enable = "avx2")]
304#[inline]
305fn popcnt256(v: __m256i) -> __m256i {
306 let lookup1 = _mm256_setr_epi8(
307 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,
308 7, 8,
309 );
310 let lookup2 = _mm256_setr_epi8(
311 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,
312 1, 0,
313 );
314 let low_mask = _mm256_set1_epi8(0x0f);
315 let lo = _mm256_and_si256(v, low_mask);
316 let hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask);
317 let popcnt1 = _mm256_shuffle_epi8(lookup1, lo);
318 let popcnt2 = _mm256_shuffle_epi8(lookup2, hi);
319 _mm256_sad_epu8(popcnt1, popcnt2)
320}
321
322#[cfg(all(
328 any(target_arch = "x86", target_arch = "x86_64"),
329 not(target_feature = "avx512vpopcntdq")
330))]
331#[target_feature(enable = "avx2")]
332#[rustfmt::skip]
334fn popcnt_avx2(bytes: &[u8]) -> u64 {
335 let zero = _mm256_setzero_si256();
336 let mut cnt = zero;
337 let mut ones = zero;
338 let mut twos = zero;
339 let mut fours = zero;
340 let mut eights = zero;
341 let mut twos_a;
342 let mut twos_b;
343 let mut fours_a;
344 let mut fours_b;
345 let mut eights_a;
346 let mut eights_b;
347 let mut sixteens;
348
349 let (blocks, tail) = bytes.as_chunks::<512>();
351 for chunk in blocks {
352 let p = chunk.as_ptr() as *const __m256i;
353 unsafe {
355 (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(0)), _mm256_loadu_si256(p.add(1)));
356 (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(2)), _mm256_loadu_si256(p.add(3)));
357 (fours_a, twos) = csa256(twos, twos_a, twos_b);
358 (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(4)), _mm256_loadu_si256(p.add(5)));
359 (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(6)), _mm256_loadu_si256(p.add(7)));
360 (fours_b, twos) = csa256(twos, twos_a, twos_b);
361 (eights_a, fours) = csa256(fours, fours_a, fours_b);
362 (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(8)), _mm256_loadu_si256(p.add(9)));
363 (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(10)), _mm256_loadu_si256(p.add(11)));
364 (fours_a, twos) = csa256(twos, twos_a, twos_b);
365 (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(12)), _mm256_loadu_si256(p.add(13)));
366 (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(14)), _mm256_loadu_si256(p.add(15)));
367 (fours_b, twos) = csa256(twos, twos_a, twos_b);
368 (eights_b, fours) = csa256(fours, fours_a, fours_b);
369 (sixteens, eights) = csa256(eights, eights_a, eights_b);
370 cnt = _mm256_add_epi64(cnt, popcnt256(sixteens));
371 }
372 }
373
374 cnt = _mm256_slli_epi64(cnt, 4);
375 cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(eights), 3));
376 cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(fours), 2));
377 cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(twos), 1));
378 cnt = _mm256_add_epi64(cnt, popcnt256(ones));
379
380 let (vecs, _) = tail.as_chunks::<32>();
382 for chunk in vecs {
383 let v = unsafe { _mm256_loadu_si256(chunk.as_ptr() as *const __m256i) };
384 cnt = _mm256_add_epi64(cnt, popcnt256(v));
385 }
386
387 let lanes: [u64; 4] = unsafe { core::mem::transmute(cnt) };
390 lanes[0] + lanes[1] + lanes[2] + lanes[3]
391}
392
393#[cfg(all(
399 any(target_arch = "x86", target_arch = "x86_64"),
400 any(not(target_feature = "avx2"), target_feature = "avx512vpopcntdq")
401))]
402#[target_feature(enable = "avx512f,avx512bw,avx512vpopcntdq")]
403fn popcnt_avx512(bytes: &[u8]) -> u64 {
404 let mut cnt = _mm512_setzero_si512();
405
406 let (blocks, tail256) = bytes.as_chunks::<256>();
408 for chunk in blocks {
409 let p = chunk.as_ptr();
410 unsafe {
412 let v0 = _mm512_loadu_si512(p.add(0) as *const _);
413 let v1 = _mm512_loadu_si512(p.add(64) as *const _);
414 let v2 = _mm512_loadu_si512(p.add(128) as *const _);
415 let v3 = _mm512_loadu_si512(p.add(192) as *const _);
416 cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v0));
417 cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v1));
418 cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v2));
419 cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v3));
420 }
421 }
422
423 let (vecs, tail64) = tail256.as_chunks::<64>();
425 for chunk in vecs {
426 let v = unsafe { _mm512_loadu_si512(chunk.as_ptr() as *const _) };
427 cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v));
428 }
429
430 if !tail64.is_empty() {
432 let len = tail64.len();
433 let mask = (u64::MAX >> (64 - len)) as __mmask64;
435 unsafe {
438 let v = _mm512_maskz_loadu_epi8(mask, tail64.as_ptr() as *const _);
439 cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v));
440 }
441 }
442
443 _mm512_reduce_add_epi64(cnt) as u64
444}
445
446#[cfg(target_arch = "aarch64")]
451fn popcnt_aarch64(bytes: &[u8]) -> u64 {
452 #[cfg(all(target_feature = "sve", simd_popcnt_have_sve))]
454 {
455 unsafe { popcnt_arm_sve(bytes) }
456 }
457
458 #[cfg(not(all(target_feature = "sve", simd_popcnt_have_sve)))]
460 {
461 popcnt_neon(bytes)
462 }
463}
464
465#[cfg(target_arch = "aarch64")]
466#[inline]
467fn vpadalq(sum: uint64x2_t, t: uint8x16_t) -> uint64x2_t {
468 unsafe { vpadalq_u32(sum, vpaddlq_u16(vpaddlq_u8(t))) }
469}
470
471#[cfg(target_arch = "aarch64")]
472fn popcnt_neon(bytes: &[u8]) -> u64 {
473 #[cfg(simd_popcnt_have_sve)]
475 if is_aarch64_feature_detected!("sve") {
476 return unsafe { popcnt_arm_sve(bytes) };
477 }
478
479 const CHUNK: usize = 64;
481 let mut cnt = 0u64;
482 let iters = bytes.len() / CHUNK;
483 let ptr = bytes.as_ptr();
484
485 if iters > 0 {
486 unsafe {
489 let mut sum = vdupq_n_u64(0);
490 let zero = vdupq_n_u8(0);
491 let mut i = 0usize;
492
493 while i < iters {
494 let mut t0 = zero;
495 let mut t1 = zero;
496 let mut t2 = zero;
497 let mut t3 = zero;
498
499 let limit = (i + 31).min(iters);
502 while i < limit {
503 let input = vld4q_u8(ptr.add(i * CHUNK));
504 t0 = vaddq_u8(t0, vcntq_u8(input.0));
505 t1 = vaddq_u8(t1, vcntq_u8(input.1));
506 t2 = vaddq_u8(t2, vcntq_u8(input.2));
507 t3 = vaddq_u8(t3, vcntq_u8(input.3));
508 i += 1;
509 }
510
511 sum = vpadalq(sum, t0);
512 sum = vpadalq(sum, t1);
513 sum = vpadalq(sum, t2);
514 sum = vpadalq(sum, t3);
515 }
516
517 let mut tmp = [0u64; 2];
518 vst1q_u64(tmp.as_mut_ptr(), sum);
519 cnt += tmp[0] + tmp[1];
520 }
521 }
522
523 let rest = &bytes[iters * CHUNK..];
526 cnt += popcnt_scalar_loop!(rest);
527 cnt
528}
529
530#[cfg(all(target_arch = "aarch64", simd_popcnt_have_sve))]
535#[target_feature(enable = "sve")]
536fn popcnt_arm_sve(bytes: &[u8]) -> u64 {
537 unsafe {
540 let mut i = 0usize;
541 let mut vcnt = svdup_n_u64(0);
542 let vl = svcntb() as usize; let ptr = bytes.as_ptr();
544 let len = bytes.len();
545
546 while i + vl * 4 <= len {
548 let v0 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i)));
549 let v1 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl)));
550 let v2 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl * 2)));
551 let v3 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl * 3)));
552 vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v0));
553 vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v1));
554 vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v2));
555 vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v3));
556 i += vl * 4;
557 }
558
559 let mut pg = svwhilelt_b8_u64(i as u64, len as u64);
562 while svptest_any(svptrue_b8(), pg) {
563 let v = svreinterpret_u64_u8(svld1_u8(pg, ptr.add(i)));
564 vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v));
565 i += vl;
566 pg = svwhilelt_b8_u64(i as u64, len as u64);
567 }
568
569 svaddv_u64(svptrue_b64(), vcnt)
570 }
571}
572
573#[cfg(test)]
578mod tests {
579 use super::*;
580
581 fn reference(bytes: &[u8]) -> u64 {
583 bytes.iter().map(|b| b.count_ones() as u64).sum()
584 }
585
586 fn popcnt64_bitwise(x: u64) -> u64 {
589 const M1: u64 = 0x5555555555555555;
590 const M2: u64 = 0x3333333333333333;
591 const M4: u64 = 0x0F0F0F0F0F0F0F0F;
592 const H01: u64 = 0x0101010101010101;
593 let x = x - ((x >> 1) & M1);
594 let x = (x & M2) + ((x >> 2) & M2);
595 let x = (x + (x >> 4)) & M4;
596 x.wrapping_mul(H01) >> 56
597 }
598
599 #[test]
600 fn empty() {
601 assert_eq!(popcnt(&[]), 0);
602 }
603
604 #[test]
605 fn all_ones() {
606 for &size in &[
607 0, 1, 7, 8, 31, 32, 39, 40, 63, 64, 255, 256, 511, 512, 4095, 4096, 65537,
608 ] {
609 let bytes = vec![0xFFu8; size];
610 assert_eq!(popcnt(&bytes), size as u64 * 8, "size={size}");
611 }
612 }
613
614 #[test]
615 fn all_zeros() {
616 let bytes = vec![0u8; 65536];
617 assert_eq!(popcnt(&bytes), 0);
618 }
619
620 #[test]
621 fn single_bits() {
622 for bit in 0u64..64 {
623 let val = 1u64 << bit;
624 assert_eq!(popcnt(&val.to_le_bytes()), 1, "bit={bit}");
625 }
626 }
627
628 #[test]
631 fn ext_trait_widths() {
632 let u8s: &[u8] = &[0xFF, 0x0F, 0x00, 0xAB, 0x01];
633 assert_eq!(
634 u8s.popcnt(),
635 u8s.iter().map(|x| x.count_ones() as u64).sum()
636 );
637
638 let u16s: &[u16] = &[0xFFFF, 0x0F0F, 0x1234, 0];
639 assert_eq!(
640 u16s.popcnt(),
641 u16s.iter().map(|x| x.count_ones() as u64).sum()
642 );
643
644 let u32s: &[u32] = &[u32::MAX, 0, 0x8000_0001];
645 assert_eq!(
646 u32s.popcnt(),
647 u32s.iter().map(|x| x.count_ones() as u64).sum()
648 );
649
650 let u64s: &[u64] = &[u64::MAX, 0x0F0F_0F0F_0F0F_0F0F, 0];
651 assert_eq!(
652 u64s.popcnt(),
653 u64s.iter().map(|x| x.count_ones() as u64).sum()
654 );
655
656 let i32s = [-1i32, 0, 1, i32::MIN];
659 assert_eq!(
660 i32s.popcnt(),
661 i32s.iter().map(|x| x.count_ones() as u64).sum()
662 );
663 assert_eq!([u128::MAX, 0].popcnt(), 128);
664 }
665
666 #[test]
670 fn pseudorandom_all_sizes() {
671 let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
672 let mut next = || {
673 state ^= state << 13;
674 state ^= state >> 7;
675 state ^= state << 17;
676 state
677 };
678
679 const MAX_SIZE: usize = 4695;
682 const MAX_OFF: usize = 7;
683 let mut bytes = vec![0u8; MAX_SIZE + MAX_OFF + 1];
684 for b in bytes.iter_mut() {
685 *b = (next() & 0xFF) as u8;
686 }
687
688 let sizes =
691 (0usize..=600).chain([1023, 1024, 1025, 2048, 4095, 4096, 4097, 4608, MAX_SIZE]);
692 for size in sizes {
693 for &off in &[0usize, 1, 3, MAX_OFF] {
694 let slice = &bytes[off..off + size];
695 assert_eq!(popcnt(slice), reference(slice), "size={size} off={off}");
696 }
697 }
698 }
699
700 #[test]
708 fn suffix_sweep() {
709 let size = std::env::var("SIMD_POPCNT_TEST_SIZE")
710 .ok()
711 .and_then(|s| s.parse::<usize>().ok())
712 .unwrap_or(20_000);
713
714 let ones = vec![0xFFu8; size];
716 check_all_suffixes(&ones);
717
718 let mut state: u64 = 0x2545_F491_4F6C_DD1D;
720 let mut bytes = vec![0u8; size];
721 for b in bytes.iter_mut() {
722 state ^= state << 13;
723 state ^= state >> 7;
724 state ^= state << 17;
725 *b = state as u8;
726 }
727 check_all_suffixes(&bytes);
728 }
729
730 fn check_all_suffixes(bytes: &[u8]) {
733 let total: u64 = bytes.iter().map(|&b| popcnt64_bitwise(b as u64)).sum();
734 let mut prefix = 0u64; for (i, &byte) in bytes.iter().enumerate() {
736 assert_eq!(popcnt(&bytes[i..]), total - prefix, "suffix at offset {i}");
737 prefix += popcnt64_bitwise(byte as u64);
738 }
739 assert_eq!(popcnt(&bytes[bytes.len()..]), 0);
741 }
742}