Improve generic version of complex masked store

Use the usual kernel mechanism which allows for specialization.
Implement specialization for avx and avx512.

Follow-up to #1391

Author: serge-sans-paille

include/xsimd/arch/common/xsimd_common_memory.hpp

@@ -865,6 +865,18 @@ namespace xsimd
             store_complex_aligned<A>(dst, src, A {});
         }
 
+        template <class A, class T, class Mode>
+        XSIMD_INLINE void
+        store_complex_masked(std::complex<T>* mem, batch<std::complex<T>, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
+        {
+            constexpr std::size_t size = batch<T, A>::size;
+            alignas(A::alignment()) std::array<std::complex<T>, size> buffer;
+            src.store_aligned(buffer.data());
+            for (std::size_t i = 0; i < size; ++i)
+                if (mask.get(i))
+                    mem[i] = buffer[i];
+        }
+
         // transpose
         template <class A, class T>
         XSIMD_INLINE void transpose(batch<T, A>* matrix_begin, batch<T, A>* matrix_end, requires_arch<common>) noexcept

include/xsimd/arch/xsimd_avx.hpp

@@ -1201,6 +1201,20 @@ namespace xsimd
             }
         }
 
+        template <class A, class T, class Mode>
+        XSIMD_INLINE void
+        store_complex_masked(std::complex<T>* mem, batch<std::complex<T>, A> const& src, batch_bool<T, A> mask, Mode mode, requires_arch<avx>) noexcept
+        {
+            using mask_register_type = typename batch_bool<T, A>::register_type;
+            mask_register_type nmask = mask.to_native();
+            batch_bool<T, A> lo_mask = zip_lo(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
+            batch_bool<T, A> hi_mask = zip_hi(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
+            batch<T, A> src_lo = zip_lo(src.real(), src.imag());
+            batch<T, A> src_hi = zip_hi(src.real(), src.imag());
+            store_masked(reinterpret_cast<T*>(mem), src_lo, lo_mask, mode, A {});
+            store_masked(reinterpret_cast<T*>(mem) + src.size, src_hi, hi_mask, mode, A {});
+        }
+
         namespace detail
         {
             // Reinterpret a constant-mask 4/8-byte load/store as same-width float

include/xsimd/arch/xsimd_avx512f.hpp

@@ -372,6 +372,28 @@ namespace xsimd
             detail::store_masked(mem, src, mask.mask(), Mode {});
         }
 
+        template <class A, class T, class Mode>
+        XSIMD_INLINE void
+        store_complex_masked(std::complex<T>* mem, batch<std::complex<T>, A> const& src, batch_bool<T, A> mask, Mode mode, requires_arch<avx512f>) noexcept
+        {
+            using mask_register_type = typename batch_bool<T, A>::register_type;
+            mask_register_type nmask = mask.to_native();
+
+            // manually zip mask
+            constexpr mask_register_type lo_bitmask = xsimd::utils::make_low_mask<mask_register_type>(src.size / 2);
+            mask_register_type lo_mask = nmask & lo_bitmask;
+            lo_mask |= lo_mask << (src.size / 2);
+
+            constexpr mask_register_type hi_bitmask = lo_bitmask << (src.size / 2);
+            mask_register_type hi_mask = nmask & hi_bitmask;
+            hi_mask |= hi_mask >> (src.size / 2);
+
+            batch<T, A> src_lo = zip_lo(src.real(), src.imag());
+            batch<T, A> src_hi = zip_hi(src.real(), src.imag());
+            detail::store_masked(reinterpret_cast<T*>(mem), src_lo, batch_bool<T, A> { lo_mask }, mode);
+            detail::store_masked(reinterpret_cast<T*>(mem) + src.size, src_hi, batch_bool<T, A> { hi_mask }, mode);
+        }
+
         // abs
         template <class A>
         XSIMD_INLINE batch<float, A> abs(batch<float, A> const& self, requires_arch<avx512f>) noexcept

include/xsimd/types/xsimd_batch.hpp

@@ -1512,13 +1512,9 @@ namespace xsimd
 
     template <class T, class A>
     template <class Mode>
-    XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode) const noexcept
+    XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode mode) const noexcept
     {
-        alignas(A::alignment()) std::array<value_type, size> buffer;
-        store_aligned(buffer.data());
-        for (std::size_t i = 0; i < size; ++i)
-            if (mask.get(i))
-                mem[i] = buffer[i];
+        kernel::store_complex_masked<A>(mem, *this, mask, mode, A {});
     }
 
     template <class T, class A>