feat: add per-lane popcount on batch

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_bitwise.hpp

@@ -0,0 +1,70 @@
+/***************************************************************************
+ * 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>
+
+namespace xsimd
+{
+
+    namespace kernel
+    {
+
+        using namespace types;
+
+        namespace detail
+        {
+            // bit i is set when i / s is even: 0x55.. for s == 1, 0x33.. for
+            // s == 2, 0x0f0f.. for s == 4, and so on up to s == bits / 2
+            template <class U>
+            constexpr U alternating_mask(unsigned s) noexcept
+            {
+                U m = 0;
+                for (unsigned i = 0; i < sizeof(U) * CHAR_BIT; ++i)
+                    if (((i / s) & 1u) == 0u)
+                        m = U(m | U(U(1) << i));
+                return m;
+            }
+        }
+
+        // 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 b_type = batch<U, A>;
+            constexpr std::size_t bits = sizeof(T) * CHAR_BIT;
+
+            b_type x = bitwise_cast<U>(self);
+            x = x - ((x >> 1) & b_type(detail::alternating_mask<U>(1)));
+            b_type const m2(detail::alternating_mask<U>(2));
+            x = (x & m2) + ((x >> 2) & m2);
+            x = (x + (x >> 4)) & b_type(detail::alternating_mask<U>(4));
+            if constexpr (bits >= 16)
+                x = (x + (x >> 8)) & b_type(detail::alternating_mask<U>(8));
+            if constexpr (bits >= 32)
+                x = (x + (x >> 16)) & b_type(detail::alternating_mask<U>(16));
+            if constexpr (bits >= 64)
+                x = (x + (x >> 32)) & b_type(detail::alternating_mask<U>(32));
+            return bitwise_cast<T>(x);
+        }
+    }
+}
+
+#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,28 @@ 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/articles/sse-popcount.html
+            __m256i const low_mask = _mm256_set1_epi8(0x0f);
+            __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 lo = _mm256_and_si256(self, low_mask);
+            __m256i const hi = _mm256_and_si256(_mm256_srli_epi16(self, 4), low_mask);
+            __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 if constexpr (sizeof(T) == 4)
+                return _mm256_madd_epi16(_mm256_maddubs_epi16(counts, _mm256_set1_epi8(1)), _mm256_set1_epi16(1));
+            else
+                return _mm256_sad_epu8(counts, _mm256_setzero_si256());
+        }
+
         // 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,28 @@ namespace xsimd
             return detail::compare_int_avx512bw<A, T, _MM_CMPINT_NE>(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<avx512bw>) noexcept
+        {
+            // per-byte counts from a nibble lookup indexed by VPSHUFB, after
+            // Wojciech Muła, http://0x80.pl/articles/sse-popcount.html
+            __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);
+            __m512i const counts = _mm512_add_epi8(_mm512_shuffle_epi8(lookup, lo), _mm512_shuffle_epi8(lookup, hi));
+            // 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 if constexpr (sizeof(T) == 4)
+                return _mm512_madd_epi16(_mm512_maddubs_epi16(counts, _mm512_set1_epi8(1)), _mm512_set1_epi16(1));
+            else
+                return _mm512_sad_epu8(counts, _mm512_setzero_si512());
+        }
+
         // 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_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,26 @@ 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, http://0x80.pl/articles/sse-popcount.html
+            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 if constexpr (sizeof(T) == 4)
+                return bitwise_cast<T>(batch<uint32_t, A>(vpaddlq_u16(vpaddlq_u8(counts))));
+            else
+                return bitwise_cast<T>(batch<uint64_t, A>(vpaddlq_u32(vpaddlq_u16(vpaddlq_u8(counts)))));
+        }
     }
 
 }

include/xsimd/arch/xsimd_ssse3.hpp

@@ -82,6 +82,28 @@ namespace xsimd
             return detail::extract_pair(self, other, i, std::make_index_sequence<size>());
         }
 
+        // 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<ssse3>) noexcept
+        {
+            // per-byte counts from a nibble lookup indexed by PSHUFB, after
+            // Wojciech Muła, http://0x80.pl/articles/sse-popcount.html
+            __m128i const low_mask = _mm_set1_epi8(0x0f);
+            __m128i const lookup = _mm_setr_epi8(0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4);
+            __m128i const lo = _mm_and_si128(self, low_mask);
+            __m128i const hi = _mm_and_si128(_mm_srli_epi16(self, 4), low_mask);
+            __m128i const counts = _mm_add_epi8(_mm_shuffle_epi8(lookup, lo), _mm_shuffle_epi8(lookup, hi));
+            // wider elements sum their byte counts
+            if constexpr (sizeof(T) == 1)
+                return counts;
+            else if constexpr (sizeof(T) == 2)
+                return _mm_maddubs_epi16(counts, _mm_set1_epi8(1));
+            else if constexpr (sizeof(T) == 4)
+                return _mm_madd_epi16(_mm_maddubs_epi16(counts, _mm_set1_epi8(1)), _mm_set1_epi16(1));
+            else
+                return _mm_sad_epu8(counts, _mm_setzero_si128());
+        }
+
         // 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<ssse3>) noexcept

include/xsimd/arch/xsimd_sve.hpp

@@ -1206,6 +1206,14 @@ namespace xsimd
             return svscale_x(detail_sve::ptrue<T>(), x, exp);
         }
 
+        // popcount
+        template <class A, class T, detail::enable_integral_t<T> = 0>
+        XSIMD_INLINE batch<T, A> popcount(const batch<T, A>& self, requires_arch<sve>) noexcept
+        {
+            using U = as_unsigned_integer_t<T>;
+            return bitwise_cast<T>(batch<U, A>(svcnt_x(detail_sve::ptrue<T>(), self)));
+        }
+
     } // namespace kernel
 } // namespace xsimd
 

include/xsimd/arch/xsimd_vsx.hpp

@@ -536,6 +536,15 @@ namespace xsimd
             return ~vec_cmpeq(self.data, 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<vsx>) noexcept
+        {
+            // VPOPCNTB/H/W/D count the bits of each element in one instruction
+            using U = as_unsigned_integer_t<T>;
+            return bitwise_cast<T>(batch<U, A>(vec_popcnt(bitwise_cast<U>(self).data)));
+        }
+
         // reciprocal
         template <class A>
         XSIMD_INLINE batch<float, A> reciprocal(batch<float, A> const& self,

include/xsimd/arch/xsimd_vxe.hpp

@@ -408,6 +408,15 @@ namespace xsimd
             return vec_mergeh(row[0].data, row[1].data) + vec_mergel(row[0].data, row[1].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<vxe>) noexcept
+        {
+            // VPOPCT counts the bits of each element in one instruction
+            using U = as_unsigned_integer_t<T>;
+            return bitwise_cast<T>(batch<U, A>(vec_popcnt(bitwise_cast<U>(self).data)));
+        }
+
         // reduce_add
         template <class A>
         XSIMD_INLINE float reduce_add(batch<float, A> const& self, requires_arch<vxe>) noexcept

include/xsimd/arch/xsimd_wasm.hpp

@@ -1813,6 +1813,27 @@ namespace xsimd
         {
             return wasm_i64x2_shuffle(self, other, 0, 2);
         }
+
+        // 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<wasm>) noexcept
+        {
+            // i8x16.popcnt only counts bytes, and pairwise widening addition
+            // stops at 32-bit elements
+            if constexpr (sizeof(T) == 8)
+            {
+                return popcount(self, common {});
+            }
+            else
+            {
+                v128_t counts = wasm_i8x16_popcnt(self);
+                if constexpr (sizeof(T) == 2)
+                    counts = wasm_u16x8_extadd_pairwise_u8x16(counts);
+                else if constexpr (sizeof(T) == 4)
+                    counts = wasm_u32x4_extadd_pairwise_u16x8(wasm_u16x8_extadd_pairwise_u8x16(counts));
+                return counts;
+            }
+        }
     }
 }
 

include/xsimd/types/xsimd_api.hpp

@@ -1946,6 +1946,20 @@ namespace xsimd
         return kernel::polar<A>(r, theta, A {});
     }
 
+    /**
+     * @ingroup batch_bitwise
+     *
+     * Counts the bits set in each element of \c x.
+     * @param x batch of integer values.
+     * @return per slot population count.
+     */
+    template <class T, class A>
+    XSIMD_INLINE batch<T, A> popcount(batch<T, A> const& x) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::popcount<A>(x, A {});
+    }
+
     /**
      * @ingroup batch_arithmetic
      *

test/test_batch_int.cpp

@@ -303,6 +303,57 @@ struct batch_int_test
         t.run();
     }
 
+    // 0, ~0, single bits, prefixes, suffixes and pseudo-random words: the
+    // patterns the bit ops actually branch on
+    static array_type bit_patterns(size_t seed)
+    {
+        constexpr size_t bits = sizeof(value_type) * CHAR_BIT;
+        using U = std::make_unsigned_t<value_type>;
+        array_type a;
+        for (size_t i = 0; i < size; ++i)
+        {
+            size_t k = seed * size + i;
+            size_t sh = k % bits;
+            U u;
+            switch (k % 6)
+            {
+            case 0:
+                u = U(0);
+                break;
+            case 1:
+                u = U(~U(0));
+                break;
+            case 2:
+                u = U(U(1) << sh);
+                break;
+            case 3:
+                u = U(~U(0)) << sh;
+                break;
+            case 4:
+                u = U(U(U(1) << sh) - U(1));
+                break;
+            default:
+                u = U(k * 0x9E3779B9u + 0x7F4A7C15u);
+                break;
+            }
+            a[i] = value_type(u);
+        }
+        return a;
+    }
+
+    void test_popcount() const
+    {
+        using U = std::make_unsigned_t<value_type>;
+        for (size_t s = 0; s < 6; ++s)
+        {
+            array_type in = bit_patterns(s), expected;
+            std::transform(in.cbegin(), in.cend(), expected.begin(), [](value_type v)
+                           { return value_type(xsimd::detail::popcount(U(v))); });
+            INFO("popcount, pattern " << s);
+            CHECK_BATCH_EQ(xsimd::popcount(batch_type::load_unaligned(in.data())), expected);
+        }
+    }
+
     void test_less_than_underflow() const
     {
         batch_type test_negative_compare = batch_type(5) - 6;
@@ -361,5 +412,10 @@ TEST_CASE_TEMPLATE("[batch int tests]", B, BATCH_INT_TYPES)
     {
         Test.test_less_than_underflow();
     }
+
+    SUBCASE("popcount")
+    {
+        Test.test_popcount();
+    }
 }
 #endif