Generic, simple implementation fox xsimd::compress (#981)

serge-sans-paille · web-flow · commit d94d7dd21b72 · 2023-11-28T22:08:42.000Z
* Generic, simple implementation fox xsimd::compress Related to #975 * fixup! Generic, simple implementation fox xsimd::compress * fixup! Generic, simple implementation fox xsimd::compress
diff --git a/docs/source/api/data_transfer.rst b/docs/source/api/data_transfer.rst
@@ -56,6 +56,8 @@ In place:
 +---------------------------------------+----------------------------------------------------+
 | :cpp:func:`insert`                    | modify a single batch slot                         |
 +---------------------------------------+----------------------------------------------------+
+| :cpp:func:`compress`                  | pack elements according to a mask                  |
++---------------------------------------+----------------------------------------------------+
 
 Between batches:
 
diff --git a/include/xsimd/arch/generic/xsimd_generic_memory.hpp b/include/xsimd/arch/generic/xsimd_generic_memory.hpp
@@ -32,6 +32,36 @@ namespace xsimd
 
         using namespace types;
 
+        // compress
+        namespace detail
+        {
+            template <class IT, class A, class I, size_t... Is>
+            inline batch<IT, A> create_compress_swizzle_mask(I bitmask, ::xsimd::detail::index_sequence<Is...>)
+            {
+                batch<IT, A> swizzle_mask(IT(0));
+                size_t i = 0;
+                alignas(A::alignment()) IT mask_buffer[batch<IT, A>::size] = { Is... };
+                size_t inserted = 0;
+                for (size_t i = 0; i < sizeof...(Is); ++i)
+                    if ((bitmask >> i) & 1u)
+                        std::swap(mask_buffer[inserted++], mask_buffer[i]);
+                return batch<IT, A>::load_aligned(&mask_buffer[0]);
+            }
+        }
+
+        template <typename A, typename T>
+        inline batch<T, A>
+        compress(batch<T, A> const& x, batch_bool<T, A> const& mask,
+                 kernel::requires_arch<generic>) noexcept
+        {
+            using IT = as_unsigned_integer_t<T>;
+            constexpr std::size_t size = batch_bool<T, A>::size;
+            auto bitmask = mask.mask();
+            auto z = select(mask, x, batch<T, A>((T)0));
+            auto compress_mask = detail::create_compress_swizzle_mask<IT, A>(bitmask, ::xsimd::detail::make_index_sequence<size>());
+            return swizzle(z, compress_mask);
+        }
+
         // extract_pair
         template <class A, class T>
         inline batch<T, A> extract_pair(batch<T, A> const& self, batch<T, A> const& other, std::size_t i, requires_arch<generic>) noexcept
diff --git a/include/xsimd/arch/xsimd_avx512f.hpp b/include/xsimd/arch/xsimd_avx512f.hpp
@@ -661,6 +661,38 @@ namespace xsimd
             return _mm512_roundscale_pd(self, _MM_FROUND_TO_POS_INF);
         }
 
+        // compress
+        template <class A>
+        inline batch<float, A> compress(batch<float, A> const& self, batch_bool<float, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_compress_ps(mask.mask(), self);
+        }
+        template <class A>
+        inline batch<double, A> compress(batch<double, A> const& self, batch_bool<double, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_compress_pd(mask.mask(), self);
+        }
+        template <class A>
+        inline batch<int32_t, A> compress(batch<int32_t, A> const& self, batch_bool<int32_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_compress_epi32(mask.mask(), self);
+        }
+        template <class A>
+        inline batch<uint32_t, A> compress(batch<uint32_t, A> const& self, batch_bool<uint32_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_compress_epi32(mask.mask(), self);
+        }
+        template <class A>
+        inline batch<int64_t, A> compress(batch<int64_t, A> const& self, batch_bool<int64_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_compress_epi64(mask.mask(), self);
+        }
+        template <class A>
+        inline batch<uint64_t, A> compress(batch<uint64_t, A> const& self, batch_bool<uint64_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_compress_epi64(mask.mask(), self);
+        }
+
         // convert
         namespace detail
         {
diff --git a/include/xsimd/types/xsimd_api.hpp b/include/xsimd/types/xsimd_api.hpp
@@ -530,6 +530,19 @@ namespace xsimd
         return kernel::clip(x, lo, hi, A {});
     }
 
+    /**
+     * @ingroup batch_data_transfer
+     *
+     * Pick elements from \c x selected by \c mask, and append them to the
+     * resulting vector, zeroing the remaining slots
+     */
+    template <class T, class A>
+    inline batch<T, A> compress(batch<T, A> const& x, batch_bool<T, A> const& mask) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::compress<A>(x, mask, A {});
+    }
+
     /**
      * @ingroup batch_complex
      *
diff --git a/test/test_shuffle.cpp b/test/test_shuffle.cpp
@@ -272,6 +272,102 @@ TEST_CASE_TEMPLATE("[slide]", B, BATCH_INT_TYPES)
 
 #endif
 
+template <class B>
+struct compress_test
+{
+    using batch_type = B;
+    using value_type = typename B::value_type;
+    using mask_batch_type = typename B::batch_bool_type;
+
+    static constexpr size_t size = B::size;
+    std::array<value_type, size> input;
+    std::array<bool, size> mask;
+    std::array<value_type, size> expected;
+
+    compress_test()
+    {
+        for (size_t i = 0; i < size; ++i)
+        {
+            input[i] = i;
+        }
+    }
+
+    void full()
+    {
+        std::fill(mask.begin(), mask.end(), true);
+
+        for (size_t i = 0; i < size; ++i)
+            expected[i] = input[i];
+
+        auto b = xsimd::compress(
+            batch_type::load_unaligned(input.data()),
+            mask_batch_type::load_unaligned(mask.data()));
+        CHECK_BATCH_EQ(b, expected);
+    }
+
+    void empty()
+    {
+        std::fill(mask.begin(), mask.end(), false);
+
+        for (size_t i = 0; i < size; ++i)
+            expected[i] = 0;
+
+        auto b = xsimd::compress(
+            batch_type::load_unaligned(input.data()),
+            mask_batch_type::load_unaligned(mask.data()));
+        CHECK_BATCH_EQ(b, expected);
+    }
+
+    void interleave()
+    {
+        for (size_t i = 0; i < size; ++i)
+            mask[i] = i % 2 == 0;
+
+        for (size_t i = 0, j = 0; i < size; ++i)
+            expected[i] = i < size / 2 ? input[2 * i] : 0;
+
+        auto b = xsimd::compress(
+            batch_type::load_unaligned(input.data()),
+            mask_batch_type::load_unaligned(mask.data()));
+        CHECK_BATCH_EQ(b, expected);
+    }
+
+    void generic()
+    {
+        for (size_t i = 0; i < size; ++i)
+            mask[i] = i % 3 == 0;
+
+        for (size_t i = 0, j = 0; i < size; ++i)
+            expected[i] = i < size / 3 ? input[3 * i] : 0;
+
+        auto b = xsimd::compress(
+            batch_type::load_unaligned(input.data()),
+            mask_batch_type::load_unaligned(mask.data()));
+        CHECK_BATCH_EQ(b, expected);
+    }
+};
+
+TEST_CASE_TEMPLATE("[compress]", B, BATCH_FLOAT_TYPES, xsimd::batch<uint32_t>, xsimd::batch<int32_t>, xsimd::batch<uint64_t>, xsimd::batch<int64_t>)
+{
+    compress_test<B> Test;
+    SUBCASE("empty")
+    {
+        Test.empty();
+    }
+    SUBCASE("full")
+    {
+        Test.full();
+    }
+    // SUBCASE("interleave")
+    //{
+    //     Test.interleave();
+    // }
+    // SUBCASE("generic")
+    //{
+    //     Test.generic();
+    // }
+}
+
 template <class B>
 struct shuffle_test
 {
