From b6827c125ee00482bc81a5739363cdd1030c5d41 Mon Sep 17 00:00:00 2001
From: serge-sans-paille <serge.guelton@telecom-bretagne.eu>
Date: Tue, 21 Nov 2023 23:56:14 +0100
Subject: [PATCH] Generic, simple implementation fox xsimd::expand

Also provide a specialization for avx512f.

Related to #975
---
 .../arch/generic/xsimd_generic_memory.hpp     | 25 +++++
 include/xsimd/arch/xsimd_avx512f.hpp          | 32 +++++++
 include/xsimd/types/xsimd_api.hpp             | 13 +++
 test/test_shuffle.cpp                         | 96 +++++++++++++++++++
 4 files changed, 166 insertions(+)

diff --git a/include/xsimd/arch/generic/xsimd_generic_memory.hpp b/include/xsimd/arch/generic/xsimd_generic_memory.hpp
index ee65b0d35..33daffe6d 100644
--- a/include/xsimd/arch/generic/xsimd_generic_memory.hpp
+++ b/include/xsimd/arch/generic/xsimd_generic_memory.hpp
@@ -62,6 +62,31 @@ namespace xsimd
             return swizzle(z, compress_mask);
         }
 
+        // expand
+        namespace detail
+        {
+            template <class IT, class A, class I, size_t... Is>
+            batch<IT, A> create_expand_swizzle_mask(I bitmask, ::xsimd::detail::index_sequence<Is...>)
+            {
+                batch<IT, A> swizzle_mask(IT(0));
+                IT j = 0;
+                (void)std::initializer_list<bool> { ((swizzle_mask = insert(swizzle_mask, j, index<Is>())), (j += ((bitmask >> Is) & 1u)), true)... };
+                return swizzle_mask;
+            }
+        }
+
+        template <typename A, typename T>
+        inline batch<T, A>
+        expand(batch<T, A> const& x, batch_bool<T, A> const& mask,
+               kernel::requires_arch<generic>) noexcept
+        {
+            constexpr std::size_t size = batch_bool<T, A>::size;
+            auto bitmask = mask.mask();
+            auto swizzle_mask = detail::create_expand_swizzle_mask<as_unsigned_integer_t<T>, A>(bitmask, ::xsimd::detail::make_index_sequence<size>());
+            auto z = swizzle(x, swizzle_mask);
+            return select(mask, z, batch<T, A>(T(0)));
+        }
+
         // 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
index b25280cb3..1243f7516 100644
--- a/include/xsimd/arch/xsimd_avx512f.hpp
+++ b/include/xsimd/arch/xsimd_avx512f.hpp
@@ -788,6 +788,38 @@ namespace xsimd
             return register_type(~self.data ^ other.data);
         }
 
+        // expand
+        template <class A>
+        inline batch<float, A> expand(batch<float, A> const& self, batch_bool<float, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_expand_ps(self, mask.mask());
+        }
+        template <class A>
+        inline batch<double, A> expand(batch<double, A> const& self, batch_bool<double, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_expand_pd(self, mask.mask());
+        }
+        template <class A>
+        inline batch<int32_t, A> expand(batch<int32_t, A> const& self, batch_bool<int32_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_expand_epi32(self, mask.mask());
+        }
+        template <class A>
+        inline batch<uint32_t, A> expand(batch<uint32_t, A> const& self, batch_bool<uint32_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_expand_epi32(self, mask.mask());
+        }
+        template <class A>
+        inline batch<int64_t, A> expand(batch<int64_t, A> const& self, batch_bool<int64_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_expand_epi64(self, mask.mask());
+        }
+        template <class A>
+        inline batch<uint64_t, A> expand(batch<uint64_t, A> const& self, batch_bool<uint64_t, A> const& mask, requires_arch<avx512f>) noexcept
+        {
+            return _mm512_maskz_expand_epi64(self, mask.mask());
+        }
+
         // floor
         template <class A>
         inline batch<float, A> floor(batch<float, A> const& self, requires_arch<avx512f>) noexcept
diff --git a/include/xsimd/types/xsimd_api.hpp b/include/xsimd/types/xsimd_api.hpp
index b4a84a66d..0420f0a09 100644
--- a/include/xsimd/types/xsimd_api.hpp
+++ b/include/xsimd/types/xsimd_api.hpp
@@ -718,6 +718,19 @@ namespace xsimd
         return kernel::exp2<A>(x, A {});
     }
 
+    /**
+     * @ingroup batch_data_transfer
+     *
+     * Load contiguous elements from \c x and place them in slots selected by \c
+     * mask, zeroing the other slots
+     */
+    template <class T, class A>
+    inline batch<T, A> expand(batch<T, A> const& x, batch_bool<T, A> const& mask) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::expand<A>(x, mask, A {});
+    }
+
     /**
      * @ingroup batch_math
      *
diff --git a/test/test_shuffle.cpp b/test/test_shuffle.cpp
index beafeb1f9..d7e8458fa 100644
--- a/test/test_shuffle.cpp
+++ b/test/test_shuffle.cpp
@@ -368,6 +368,102 @@ TEST_CASE_TEMPLATE("[compress]", B, BATCH_FLOAT_TYPES, xsimd::batch<uint32_t>, x
     // }
 }
 
+template <class B>
+struct expand_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;
+
+    expand_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::expand(
+            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::expand(
+            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] = mask[i] ? input[j++] : 0;
+
+        auto b = xsimd::expand(
+            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] = mask[i] ? input[j++] : 0;
+
+        auto b = xsimd::expand(
+            batch_type::load_unaligned(input.data()),
+            mask_batch_type::load_unaligned(mask.data()));
+        CHECK_BATCH_EQ(b, expected);
+    }
+};
+
+TEST_CASE_TEMPLATE("[expand]", B, BATCH_FLOAT_TYPES, xsimd::batch<uint32_t>, xsimd::batch<int32_t>, xsimd::batch<uint64_t>, xsimd::batch<int64_t>)
+{
+    expand_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
 {