feat: add per-lane popcount on batch

.github/workflows/linux.yml

@@ -116,7 +116,8 @@ jobs:
           CMAKE_EXTRA_ARGS="$CMAKE_EXTRA_ARGS -DTARGET_ARCH=icelake-server"
         fi
         if [[ '${{ matrix.sys.flags }}' == 'avx512vnni' ]]; then
-          CMAKE_EXTRA_ARGS="$CMAKE_EXTRA_ARGS -DTARGET_ARCH=knm"
+          # knm enables avx5124vnniw, not avx512vnni: it selects the avx512pf arch
+          CMAKE_EXTRA_ARGS="$CMAKE_EXTRA_ARGS -DTARGET_ARCH=cascadelake"
         fi
         if [[ '${{ matrix.sys.flags }}' == 'i386' ]]; then
           CXX_FLAGS="$CXX_FLAGS -m32"

docs/source/api/bitwise_operators_index.rst

@@ -48,6 +48,8 @@ Bitwise Operators
 +---------------------------------------+----------------------------------------------------+
 | :cpp:func:`rotl`                      | per slot rotate left                               |
 +---------------------------------------+----------------------------------------------------+
+| :cpp:func:`popcount`                  | per slot population count                          |
++---------------------------------------+----------------------------------------------------+
 
 ----
 

include/xsimd/arch/common/xsimd_common_bit.hpp

@@ -7,7 +7,7 @@
 
 #include "../../config/xsimd_config.hpp"
 
-#if XSIMD_CPP_VERSION > 202002L
+#if XSIMD_CPP_VERSION >= 202002L
 
 #include <version>
 

include/xsimd/arch/common/xsimd_common_bitwise.hpp

@@ -0,0 +1,87 @@
+/***************************************************************************
+ * Copyright (c) Johan Mabille, Sylvain Corlay, Wolf Vollprecht and         *
+ * Martin Renou                                                             *
+ * Copyright (c) QuantStack                                                 *
+ * Copyright (c) Serge Guelton                                              *
+ * Copyright (c) Marco Barbone                                              *
+ *                                                                          *
+ * Distributed under the terms of the BSD 3-Clause License.                 *
+ *                                                                          *
+ * The full license is in the file LICENSE, distributed with this software. *
+ ****************************************************************************/
+
+#ifndef XSIMD_COMMON_BITWISE_HPP
+#define XSIMD_COMMON_BITWISE_HPP
+
+#include "./xsimd_common_details.hpp"
+
+#include <climits>
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+namespace xsimd
+{
+
+    namespace kernel
+    {
+
+        using namespace types;
+
+        namespace detail
+        {
+            // Pattern P repeated across U from bit 0 up. The type of P is the
+            // repeat unit, so repeat_pattern<uint64_t, uint8_t(0x55)>() is
+            // 0x5555555555555555 and repeat_pattern<uint32_t, uint16_t(0x0f0f)>()
+            // is 0x0f0f0f0f.
+            template <class U, auto P>
+            constexpr U repeat_pattern() noexcept
+            {
+                using unit = decltype(P);
+                static_assert(std::is_unsigned<unit>::value, "the repeat unit must be unsigned");
+                static_assert(P != 0, "the pattern must be non-zero");
+                static_assert(sizeof(U) % sizeof(unit) == 0, "the repeat unit must divide U");
+                U result = 0;
+                for (std::size_t i = 0; i < sizeof(U) / sizeof(unit); ++i)
+                    result = U(result | U(U(P) << (i * sizeof(unit) * CHAR_BIT)));
+                return result;
+            }
+        }
+
+        // popcount
+        // SWAR fold, popcount64b from Hacker's Delight, listed in
+        // https://en.wikipedia.org/wiki/Hamming_weight#Efficient_implementation
+        template <class A, class T, class /*=std::enable_if_t<std::is_integral_v<T>>*/>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<common>) noexcept
+        {
+            using U = as_unsigned_integer_t<T>;
+            using w_type = batch<uint64_t, A>;
+            constexpr std::size_t bits = sizeof(T) * CHAR_BIT;
+
+            // Every step runs on 64-bit lanes whatever T is. Each step confines
+            // its own carries, so a bit that crosses a T boundary is always
+            // masked off again, and targets with no narrow shift do not pay for
+            // one to be emulated.
+            w_type x = bitwise_cast<uint64_t>(self);
+            x = x - ((x >> 1) & w_type(detail::repeat_pattern<uint64_t, uint8_t(0x55)>()));
+            w_type const m2(detail::repeat_pattern<uint64_t, uint8_t(0x33)>());
+            x = (x & m2) + ((x >> 2) & m2);
+            x = (x + (x >> 4)) & w_type(detail::repeat_pattern<uint64_t, uint8_t(0x0f)>());
+            if constexpr (bits == 8)
+                return bitwise_cast<T>(x);
+
+            // Byte counts are at most 8, so a per-lane sum is at most 64 and no
+            // byte carries into the next. The final mask keeps the one byte per
+            // lane that holds the whole count and drops the bytes that summed
+            // across a lane boundary.
+            x = x + (x >> 8);
+            if constexpr (bits >= 32)
+                x = x + (x >> 16);
+            if constexpr (bits >= 64)
+                x = x + (x >> 32);
+            return bitwise_cast<T>(x) & batch<T, A>(T(U(0xff)));
+        }
+    }
+}
+
+#endif

include/xsimd/arch/xsimd_avx.hpp

@@ -1411,6 +1411,15 @@ namespace xsimd
             return _mm256_castps_si256(_mm256_xor_ps(_mm256_castsi256_ps(self.data), _mm256_castsi256_ps(other.data)));
         }
 
+        // popcount
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<avx>) noexcept
+        {
+            return detail::fwd_to_sse([](__m128i s) noexcept
+                                      { return popcount(batch<T, ssse3>(s), ssse3 {}); },
+                                      self);
+        }
+
         // reciprocal
         template <class A>
         XSIMD_INLINE batch<float, A> reciprocal(batch<float, A> const& self,

include/xsimd/arch/xsimd_avx2.hpp

@@ -962,6 +962,38 @@ namespace xsimd
                                                { return batch<uint64_t, A>(_mm256_mul_epu32(a, b)); });
         }
 
+        // popcount
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<avx2>) noexcept
+        {
+            // per-byte counts from a nibble lookup indexed by VPSHUFB, after
+            // Wojciech Muła, http://0x80.pl/notesen/2008-05-24-sse-popcount.html
+            __m256i const low_mask = _mm256_set1_epi8(0x0f);
+            __m256i const lo = _mm256_and_si256(self, low_mask);
+            __m256i const hi = _mm256_and_si256(_mm256_srli_epi16(self, 4), low_mask);
+            if constexpr (sizeof(T) == 8)
+            {
+                // tables biased by +4 and -4 turn the VPSADBW difference into the
+                // per-byte count, so one instruction adds the nibble counts and
+                // sums the eight bytes, after https://github.com/kimwalisch/libpopcnt
+                __m256i const lookup_lo = _mm256_broadcastsi128_si256(_mm_setr_epi8(4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8));
+                __m256i const lookup_hi = _mm256_broadcastsi128_si256(_mm_setr_epi8(4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0));
+                return _mm256_sad_epu8(_mm256_shuffle_epi8(lookup_lo, lo), _mm256_shuffle_epi8(lookup_hi, hi));
+            }
+            else
+            {
+                __m256i const lookup = _mm256_broadcastsi128_si256(_mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4));
+                __m256i const counts = _mm256_add_epi8(_mm256_shuffle_epi8(lookup, lo), _mm256_shuffle_epi8(lookup, hi));
+                // wider elements sum their byte counts
+                if constexpr (sizeof(T) == 1)
+                    return counts;
+                else if constexpr (sizeof(T) == 2)
+                    return _mm256_maddubs_epi16(counts, _mm256_set1_epi8(1));
+                else
+                    return _mm256_madd_epi16(_mm256_maddubs_epi16(counts, _mm256_set1_epi8(1)), _mm256_set1_epi16(1));
+            }
+        }
+
         // reduce_add
         template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
         XSIMD_INLINE T reduce_add(batch<T, A> const& self, requires_arch<avx2>) noexcept

include/xsimd/arch/xsimd_avx512bw.hpp

@@ -561,6 +561,49 @@ namespace xsimd
             return detail::compare_int_avx512bw<A, T, _MM_CMPINT_NE>(self, other);
         }
 
+        namespace detail
+        {
+            // per-byte counts from a nibble lookup indexed by VPSHUFB, after
+            // Wojciech Muła, http://0x80.pl/notesen/2008-05-24-sse-popcount.html
+            XSIMD_INLINE __m512i popcount_bytes(__m512i self) noexcept
+            {
+                __m512i const low_mask = _mm512_set1_epi8(0x0f);
+                __m512i const lookup = _mm512_broadcast_i32x4(_mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4));
+                __m512i const lo = _mm512_and_si512(self, low_mask);
+                __m512i const hi = _mm512_and_si512(_mm512_srli_epi16(self, 4), low_mask);
+                return _mm512_add_epi8(_mm512_shuffle_epi8(lookup, lo), _mm512_shuffle_epi8(lookup, hi));
+            }
+        }
+
+        // popcount
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<avx512bw>) noexcept
+        {
+            if constexpr (sizeof(T) == 8)
+            {
+                __m512i const low_mask = _mm512_set1_epi8(0x0f);
+                __m512i const lo = _mm512_and_si512(self, low_mask);
+                __m512i const hi = _mm512_and_si512(_mm512_srli_epi16(self, 4), low_mask);
+                // tables biased by +4 and -4 turn the VPSADBW difference into the
+                // per-byte count, so one instruction adds the nibble counts and
+                // sums the eight bytes, after https://github.com/kimwalisch/libpopcnt
+                __m512i const lookup_lo = _mm512_broadcast_i32x4(_mm_setr_epi8(4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8));
+                __m512i const lookup_hi = _mm512_broadcast_i32x4(_mm_setr_epi8(4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0));
+                return _mm512_sad_epu8(_mm512_shuffle_epi8(lookup_lo, lo), _mm512_shuffle_epi8(lookup_hi, hi));
+            }
+            else
+            {
+                __m512i const counts = detail::popcount_bytes(self);
+                // wider elements sum their byte counts
+                if constexpr (sizeof(T) == 1)
+                    return counts;
+                else if constexpr (sizeof(T) == 2)
+                    return _mm512_maddubs_epi16(counts, _mm512_set1_epi8(1));
+                else
+                    return _mm512_madd_epi16(_mm512_maddubs_epi16(counts, _mm512_set1_epi8(1)), _mm512_set1_epi16(1));
+            }
+        }
+
         // sadd
         template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
         XSIMD_INLINE batch<T, A> sadd(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx512bw>) noexcept

include/xsimd/arch/xsimd_avx512vnni_avx512bw.hpp

@@ -13,5 +13,26 @@
 #define XSIMD_AVX512VNNI_AVX512_BW_HPP
 
 #include "../types/xsimd_avx512vnni_avx512bw_register.hpp"
+#include "./xsimd_avx512bw.hpp"
+
+namespace xsimd
+{
+
+    namespace kernel
+    {
+
+        using namespace types;
+
+        // popcount
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T> && sizeof(T) == 4>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<avx512vnni<avx512bw>>) noexcept
+        {
+            // VPDPBUSD does in one uop what the VPMADDUBSW + VPMADDWD pair of
+            // the avx512bw kernel does in two; the zero accumulator it needs
+            // costs nothing, since the register copy is eliminated at rename
+            return _mm512_dpbusd_epi32(_mm512_setzero_si512(), detail::popcount_bytes(self), _mm512_set1_epi8(1));
+        }
+    }
+}
 
 #endif

include/xsimd/arch/xsimd_avx512vnni_avx512vbmi2.hpp

@@ -13,5 +13,26 @@
 #define XSIMD_AVX512VNNI_AVX512VBMI2_HPP
 
 #include "../types/xsimd_avx512vnni_avx512vbmi2_register.hpp"
+#include "./xsimd_avx512bw.hpp"
+
+namespace xsimd
+{
+
+    namespace kernel
+    {
+
+        using namespace types;
+
+        // popcount
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T> && sizeof(T) == 4>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<avx512vnni<avx512vbmi2>>) noexcept
+        {
+            // VPDPBUSD does in one uop what the VPMADDUBSW + VPMADDWD pair of
+            // the avx512bw kernel does in two; the zero accumulator it needs
+            // costs nothing, since the register copy is eliminated at rename
+            return _mm512_dpbusd_epi32(_mm512_setzero_si512(), detail::popcount_bytes(self), _mm512_set1_epi8(1));
+        }
+    }
+}
 
 #endif

include/xsimd/arch/xsimd_common.hpp

@@ -14,6 +14,7 @@
 
 #include "./common/xsimd_common_arithmetic.hpp"
 #include "./common/xsimd_common_bit.hpp"
+#include "./common/xsimd_common_bitwise.hpp"
 #include "./common/xsimd_common_cast.hpp"
 #include "./common/xsimd_common_complex.hpp"
 #include "./common/xsimd_common_logical.hpp"

include/xsimd/arch/xsimd_common_fwd.hpp

@@ -85,6 +85,8 @@ namespace xsimd
         XSIMD_INLINE batch<T, A> rotr(batch<T, A> const& self, STy other, requires_arch<common>) noexcept;
         template <size_t count, class A, class T>
         XSIMD_INLINE batch<T, A> rotr(batch<T, A> const& self, requires_arch<common>) noexcept;
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<common>) noexcept;
         template <class A, class T>
         XSIMD_INLINE batch<T, A> load(T const* mem, aligned_mode, requires_arch<A>) noexcept;
         template <class A, class T>

include/xsimd/arch/xsimd_neon.hpp

@@ -3680,6 +3680,31 @@ namespace xsimd
         WRAP_MASK_OP(countr_one)
 
 #undef WRAP_MASK_OP
+
+        /************
+         * popcount *
+         ************/
+
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<neon>) noexcept
+        {
+            // CNT counts bytes; wider elements fold them with pairwise widening
+            // adds, after Wojciech Muła's popcnt_neon_vcnt,
+            // https://github.com/WojciechMula/sse-popcount
+            uint8x16_t counts = vcntq_u8(bitwise_cast<uint8_t>(self).data);
+            if constexpr (sizeof(T) == 1)
+                return bitwise_cast<T>(batch<uint8_t, A>(counts));
+            else if constexpr (sizeof(T) == 2)
+                return bitwise_cast<T>(batch<uint16_t, A>(vpaddlq_u8(counts)));
+            else
+            {
+                uint32x4_t const wide = vpaddlq_u16(vpaddlq_u8(counts));
+                if constexpr (sizeof(T) == 4)
+                    return bitwise_cast<T>(batch<uint32_t, A>(wide));
+                else
+                    return bitwise_cast<T>(batch<uint64_t, A>(vpaddlq_u32(wide)));
+            }
+        }
     }
 
 }

include/xsimd/arch/xsimd_scalar.hpp

@@ -13,6 +13,7 @@
 #define XSIMD_SCALAR_HPP
 
 #include "../config/xsimd_macros.hpp"
+#include "./common/xsimd_common_bit.hpp"
 
 #include <cassert>
 #include <cmath>
@@ -408,6 +409,14 @@ namespace xsimd
         return +x;
     }
 
+    // returns T rather than int, so that the scalar and the batch overload agree
+    template <class T>
+    XSIMD_INLINE std::enable_if_t<std::is_integral_v<T>, T>
+    popcount(T x) noexcept
+    {
+        return T(detail::popcount(std::make_unsigned_t<T>(x)));
+    }
+
     XSIMD_INLINE float reciprocal(float const& x) noexcept
     {
         return 1.f / x;

include/xsimd/arch/xsimd_sse2.hpp

@@ -1577,6 +1577,39 @@ namespace xsimd
             return _mm_xor_pd(self, other);
         }
 
+        // popcount
+        template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
+        XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& self, requires_arch<sse2>) noexcept
+        {
+            // SWAR fold to per-byte counts, after popcount64b from Hacker's
+            // Delight, vectorized as in Wojciech Muła's popcnt_SSE_bit_parallel,
+            // https://github.com/WojciechMula/sse-popcount. The byte shifts need
+            // no mask of their own, since every bit they leak across a byte
+            // boundary falls under one of the three masks the fold already
+            // applies.
+            __m128i const m1 = _mm_set1_epi8(0x55);
+            __m128i const m2 = _mm_set1_epi8(0x33);
+            __m128i const m4 = _mm_set1_epi8(0x0f);
+            __m128i x = _mm_sub_epi8(self, _mm_and_si128(_mm_srli_epi16(self, 1), m1));
+            x = _mm_add_epi8(_mm_and_si128(x, m2), _mm_and_si128(_mm_srli_epi16(x, 2), m2));
+            __m128i const counts = _mm_and_si128(_mm_add_epi8(x, _mm_srli_epi16(x, 4)), m4);
+            if constexpr (sizeof(T) == 1)
+                return counts;
+            // PSADBW sums the eight byte counts of a 64-bit element in one
+            // instruction, which the shift-and-mask fold needs nine to do
+            else if constexpr (sizeof(T) == 8)
+                return _mm_sad_epu8(counts, _mm_setzero_si128());
+            else
+            {
+                __m128i const pairs = _mm_and_si128(_mm_add_epi8(counts, _mm_srli_epi16(counts, 8)), _mm_set1_epi16(0xff));
+                if constexpr (sizeof(T) == 2)
+                    return pairs;
+                // PMADDWD adds the two halves of a 32-bit element in one instruction
+                else
+                    return _mm_madd_epi16(pairs, _mm_set1_epi16(1));
+            }
+        }
+
         // reciprocal
         template <class A>
         XSIMD_INLINE batch<float, A> reciprocal(batch<float, A> const& self,