Skip to content

Commit 29f64dc

Browse files
Improve generic version of complex masked store
Use the usual kernel mechanism which allows for specialization. Leverage existing masked store mechanism instead of implementing a new one. Follow-up to #1391
1 parent 67e96b0 commit 29f64dc

2 files changed

Lines changed: 35 additions & 6 deletions

File tree

include/xsimd/arch/common/xsimd_common_memory.hpp

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -865,6 +865,39 @@ namespace xsimd
865865
store_complex_aligned<A>(dst, src, A {});
866866
}
867867

868+
template <class A, class T, class Mode>
869+
XSIMD_INLINE void
870+
store_complex_masked(std::complex<T>* mem, batch<std::complex<T>, A> const& src, batch_bool<T, A> mask, Mode mode, requires_arch<common>) noexcept
871+
{
872+
// Generic fallback: mask and real /imag part are zipped before
873+
// calling the generic masked store routine.
874+
using mask_register_type = typename batch_bool<T, A>::register_type;
875+
mask_register_type nmask = mask.to_native();
876+
batch_bool<T, A> lo_mask;
877+
batch_bool<T, A> hi_mask;
878+
879+
// Generic zip_lo/hi of batch_bool depending on native register type {
880+
if constexpr (sizeof(mask_register_type) < sizeof(mask))
881+
{
882+
constexpr mask_register_type lo_bitmask = xsimd::utils::make_low_mask<mask_register_type>(src.size / 2);
883+
constexpr mask_register_type hi_bitmask = lo_bitmask << (src.size / 2);
884+
lo_mask = nmask & lo_bitmask;
885+
lo_mask |= lo_mask << (src.size / 2);
886+
hi_mask = nmask & hi_bitmask;
887+
hi_mask |= hi_mask >> (src.size / 2);
888+
}
889+
else
890+
{
891+
lo_mask = zip_lo(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
892+
hi_mask = zip_hi(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
893+
}
894+
// }.
895+
batch<T, A> src_lo = zip_lo(src.real(), src.imag());
896+
batch<T, A> src_hi = zip_hi(src.real(), src.imag());
897+
src_lo.store(reinterpret_cast<T*>(mem), lo_mask, mode);
898+
src_hi.store(reinterpret_cast<T*>(mem) + src.size, hi_mask, mode);
899+
}
900+
868901
// transpose
869902
template <class A, class T>
870903
XSIMD_INLINE void transpose(batch<T, A>* matrix_begin, batch<T, A>* matrix_end, requires_arch<common>) noexcept

include/xsimd/types/xsimd_batch.hpp

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1512,13 +1512,9 @@ namespace xsimd
15121512

15131513
template <class T, class A>
15141514
template <class Mode>
1515-
XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode) const noexcept
1515+
XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode mode) const noexcept
15161516
{
1517-
alignas(A::alignment()) std::array<value_type, size> buffer;
1518-
store_aligned(buffer.data());
1519-
for (std::size_t i = 0; i < size; ++i)
1520-
if (mask.get(i))
1521-
mem[i] = buffer[i];
1517+
kernel::store_complex_masked<A>(mem, *this, mask, mode, A { });
15221518
}
15231519

15241520
template <class T, class A>

0 commit comments

Comments
 (0)