Provide batch_bool::count() to count the number of positive value in …

…a batch Fix #1042
xtensor-stack · Sep 9, 2024 · c7ca125 · c7ca125
1 parent caeb777
commit c7ca125
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 0 deletions.
diff --git a/include/xsimd/arch/generic/xsimd_generic_logical.hpp b/include/xsimd/arch/generic/xsimd_generic_logical.hpp
@@ -14,6 +14,8 @@
 
 #include "./xsimd_generic_details.hpp"
 
+#include <climits>
+
 namespace xsimd
 {
 
@@ -22,6 +24,45 @@ namespace xsimd
 
         using namespace types;
 
+        // count
+        template <class A, class T>
+        XSIMD_INLINE size_t count(batch_bool<T, A> const& x, requires_arch<generic>) noexcept
+        {
+            uint64_t m = x.mask();
+            XSIMD_IF_CONSTEXPR(batch_bool<T, A>::size < 14)
+            {
+                // https://graphics.stanford.edu/~seander/bithacks.html#CountBitsSet64
+                return (m * 0x200040008001ULL & 0x111111111111111ULL) % 0xf;
+            }
+            else
+            {
+#if defined __has_builtin
+#if __has_builtin(__builtin_popcount)
+#define builtin_popcount(v) __builtin_popcount(v)
+#endif
+#endif
+#ifdef builtin_popcount
+                return builtin_popcount(m);
+#else
+                // FIXME: we could do better by dispatching to the appropriate
+                // popcount instruction depending on the arch...
+                XSIMD_IF_CONSTEXPR(batch_bool<T, A>::size < 32)
+                {
+                    m = m - ((m >> 1) & 0x55555555); // reuse input as temporary
+                    m = (m & 0x33333333) + ((m >> 2) & 0x33333333); // temp
+                    return ((m + (m >> 4) & 0xF0F0F0F) * 0x1010101) >> 24; // count
+                }
+                else
+                {
+                    m = m - ((m >> 1) & (uint64_t) ~(uint64_t)0 / 3); // temp
+                    m = (m & (uint64_t) ~(uint64_t)0 / 15 * 3) + ((m >> 2) & (uint64_t) ~(uint64_t)0 / 15 * 3); // temp
+                    m = (m + (m >> 4)) & (uint64_t) ~(uint64_t)0 / 255 * 15; // temp
+                    return (m * ((uint64_t) ~(uint64_t)0 / 255)) >> (sizeof(uint64_t) - 1) * CHAR_BIT; // count
+                }
+#endif
+            }
+        }
+
         // from  mask
         template <class A, class T>
         XSIMD_INLINE batch_bool<T, A> from_mask(batch_bool<T, A> const&, uint64_t mask, requires_arch<generic>) noexcept

diff --git a/include/xsimd/types/xsimd_api.hpp b/include/xsimd/types/xsimd_api.hpp
@@ -631,6 +631,20 @@ namespace xsimd
         return kernel::cosh<A>(x, A {});
     }
 
+    /**
+     * @ingroup batch_reducers
+     *
+     * Count the number of values set to true in the batch \c x
+     * @param x boolean or batch of boolean
+     * @return the result of the counting.
+     */
+    template <class T, class A>
+    XSIMD_INLINE size_t count(batch_bool<T, A> const& x) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::count<A>(x, A {});
+    }
+
     /**
      * @ingroup batch_arithmetic
      *

diff --git a/test/test_batch_bool.cpp b/test/test_batch_bool.cpp
@@ -441,6 +441,14 @@ struct batch_bool_test
         CHECK_EQ(batch_bool_type::from_mask(bool_g.interspersed.mask()).mask(), bool_g.interspersed.mask());
     }
 
+    void test_count() const
+    {
+        auto bool_g = xsimd::get_bool<batch_bool_type> {};
+        CHECK_EQ(count(bool_g.all_false), 0);
+        CHECK_EQ(count(bool_g.all_true), batch_bool_type::size);
+        CHECK_EQ(count(bool_g.half), batch_bool_type::size / 2);
+    }
+
     void test_comparison() const
     {
         auto bool_g = xsimd::get_bool<batch_bool_type> {};
@@ -485,6 +493,8 @@ TEST_CASE_TEMPLATE("[xsimd batch bool]", B, BATCH_TYPES)
 
     SUBCASE("mask") { Test.test_mask(); }
 
+    SUBCASE("count") { Test.test_count(); }
+
     SUBCASE("eq neq") { Test.test_comparison(); }
 }
 #endif