Skip to content

Commit c17f5af

Browse files
serge-sans-pailleserge-sans-paille
committed
Improve generic version of complex masked store & load
Use the usual kernel mechanism which allows for specialization. Implement specialization for avx and avx512. Follow-up to #1391
1 parent 67e96b0 commit c17f5af

4 files changed

Lines changed: 103 additions & 11 deletions

File tree

‎include/xsimd/arch/common/xsimd_common_memory.hpp‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,29 @@ namespace xsimd
450450
return batch<T, A>::load_aligned(buffer.data());
451451
}
452452

453+
template <class A, class T, class Mode>
454+
XSIMD_INLINE batch<std::complex<T>, A>
455+
load_complex_masked(std::complex<T> const* mem, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
456+
{
457+
// Scalar fallback: only active lanes are touched. Arches with
458+
// hardware predicated loads should override this.
459+
constexpr std::size_t size = batch<T, A>::size;
460+
alignas(A::alignment()) std::array<T, size> buffer_real;
461+
alignas(A::alignment()) std::array<T, size> buffer_imag;
462+
for (std::size_t i = 0; i < size; ++i)
463+
if (mask.get(i))
464+
{
465+
buffer_real[i] = mem[i].real();
466+
buffer_imag[i] = mem[i].imag();
467+
}
468+
else
469+
{
470+
buffer_real[i] = T(0);
471+
buffer_imag[i] = T(0);
472+
}
473+
return batch<std::complex<T>, A>::load_aligned(buffer_real.data(), buffer_imag.data());
474+
}
475+
453476
template <class A, class T_in, class T_out, bool... Values, class alignment>
454477
XSIMD_INLINE void
455478
store_masked(T_out* mem, batch<T_in, A> const& src, batch_bool_constant<T_in, A, Values...> mask, alignment mode, requires_arch<common>) noexcept
@@ -865,6 +888,18 @@ namespace xsimd
865888
store_complex_aligned<A>(dst, src, A {});
866889
}
867890

891+
template <class A, class T, class Mode>
892+
XSIMD_INLINE void
893+
store_complex_masked(std::complex<T>* mem, batch<std::complex<T>, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
894+
{
895+
constexpr std::size_t size = batch<T, A>::size;
896+
alignas(A::alignment()) std::array<std::complex<T>, size> buffer;
897+
src.store_aligned(buffer.data());
898+
for (std::size_t i = 0; i < size; ++i)
899+
if (mask.get(i))
900+
mem[i] = buffer[i];
901+
}
902+
868903
// transpose
869904
template <class A, class T>
870905
XSIMD_INLINE void transpose(batch<T, A>* matrix_begin, batch<T, A>* matrix_end, requires_arch<common>) noexcept

‎include/xsimd/arch/xsimd_avx.hpp‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1025,6 +1025,19 @@ namespace xsimd
10251025
return _mm256_maskload_pd(mem, _mm256_castpd_si256(mask));
10261026
}
10271027

1028+
template <class A, class T, class Mode>
1029+
XSIMD_INLINE batch<std::complex<T>, A>
1030+
load_complex_masked(std::complex<T> const* mem, batch_bool<T, A> mask, Mode mode, requires_arch<avx>) noexcept
1031+
{
1032+
using mask_register_type = typename batch_bool<T, A>::register_type;
1033+
mask_register_type nmask = mask.to_native();
1034+
batch_bool<T, A> lo_mask = zip_lo(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
1035+
batch_bool<T, A> hi_mask = zip_hi(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
1036+
batch<T, A> res_lo = batch<T, A>::load(reinterpret_cast<T const*>(mem), lo_mask, mode);
1037+
batch<T, A> res_hi = batch<T, A>::load(reinterpret_cast<T const*>(mem) + mask.size, hi_mask, mode);
1038+
return detail::load_complex(res_lo, res_hi, A{});
1039+
}
1040+
10281041
// 4/8-byte ints: bitcast to same-width float, reuse the vmaskmov path.
10291042
template <class A, class T, class Mode>
10301043
XSIMD_INLINE std::enable_if_t<std::is_integral_v<T> && (sizeof(T) == 4 || sizeof(T) == 8), batch<T, A>>
@@ -1201,6 +1214,16 @@ namespace xsimd
12011214
}
12021215
}
12031216

1217+
template <class A, class T, class Mode>
1218+
XSIMD_INLINE void
1219+
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
1220+
{
1221+
auto src_lo = detail::complex_low(src, A{});
1222+
auto src_hi = detail::complex_high(src, A{});
1223+
store_masked(reinterpret_cast<T*>(mem), src_lo, mask, mode, A{});
1224+
store_masked(reinterpret_cast<T*>(mem) + src.size, src_hi, mask, mode, A{});
1225+
}
1226+
12041227
namespace detail
12051228
{
12061229
// Reinterpret a constant-mask 4/8-byte load/store as same-width float

‎include/xsimd/arch/xsimd_avx512f.hpp‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,37 @@ namespace xsimd
372372
detail::store_masked(mem, src, mask.mask(), Mode {});
373373
}
374374

375+
namespace detail
376+
{
377+
template <class A, class T>
378+
std::array<batch_bool<T, A>, 2> zip_complex_mask(batch_bool<T, A> mask)
379+
{
380+
using mask_register_type = typename batch_bool<T, A>::register_type;
381+
mask_register_type nmask = mask.to_native();
382+
383+
constexpr mask_register_type lo_bitmask = xsimd::utils::make_low_mask<mask_register_type>(mask.size / 2);
384+
mask_register_type lo_mask = nmask & lo_bitmask;
385+
lo_mask |= lo_mask << (mask.size / 2);
386+
387+
constexpr mask_register_type hi_bitmask = lo_bitmask << (mask.size / 2);
388+
mask_register_type hi_mask = nmask & hi_bitmask;
389+
hi_mask |= hi_mask >> (mask.size / 2);
390+
391+
return { batch_bool<T, A>{ lo_mask }, batch_bool<T, A>{ hi_mask } };
392+
}
393+
}
394+
395+
template <class A, class T, class Mode>
396+
XSIMD_INLINE void
397+
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
398+
{
399+
auto [lo_mask, hi_mask] = detail::zip_complex_mask(mask);
400+
batch<T, A> src_lo = zip_lo(src.real(), src.imag());
401+
batch<T, A> src_hi = zip_hi(src.real(), src.imag());
402+
detail::store_masked(reinterpret_cast<T*>(mem), src_lo, lo_mask, mode);
403+
detail::store_masked(reinterpret_cast<T*>(mem) + src.size, src_hi, hi_mask, mode);
404+
}
405+
375406
// abs
376407
template <class A>
377408
XSIMD_INLINE batch<float, A> abs(batch<float, A> const& self, requires_arch<avx512f>) noexcept
@@ -1633,6 +1664,16 @@ namespace xsimd
16331664
}
16341665
}
16351666

1667+
template <class A, class T, class Mode>
1668+
XSIMD_INLINE batch<std::complex<T>, A>
1669+
load_complex_masked(std::complex<T> const* mem, batch_bool<T, A> mask, Mode mode, requires_arch<avx512f>) noexcept
1670+
{
1671+
auto [lo_mask, hi_mask] = detail::zip_complex_mask(mask);
1672+
batch<T, A> res_lo = batch<T, A>::load(reinterpret_cast<T const*>(mem), lo_mask, mode);
1673+
batch<T, A> res_hi = batch<T, A>::load(reinterpret_cast<T const*>(mem) + mask.size, hi_mask, mode);
1674+
return detail::load_complex(res_lo, res_hi, A{});
1675+
}
1676+
16361677
// load_unaligned
16371678
template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
16381679
XSIMD_INLINE batch<T, A> load_unaligned(T const* mem, convert<T>, requires_arch<avx512f>) noexcept

‎include/xsimd/types/xsimd_batch.hpp‎

Lines changed: 4 additions & 11 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>
@@ -1561,12 +1557,9 @@ namespace xsimd
15611557

15621558
template <class T, class A>
15631559
template <class Mode>
1564-
XSIMD_INLINE batch<std::complex<T>, A> batch<std::complex<T>, A>::load(value_type const* mem, batch_bool<T, A> mask, Mode) noexcept
1560+
XSIMD_INLINE batch<std::complex<T>, A> batch<std::complex<T>, A>::load(value_type const* mem, batch_bool<T, A> mask, Mode mode) noexcept
15651561
{
1566-
alignas(A::alignment()) std::array<value_type, size> buffer {};
1567-
for (std::size_t i = 0; i < size; ++i)
1568-
buffer[i] = mask.get(i) ? mem[i] : value_type(0);
1569-
return load_aligned(buffer.data());
1562+
return kernel::load_complex_masked<A>(mem, mask, mode, A {});
15701563
}
15711564

15721565
template <class T, class A>

0 commit comments

Comments
 (0)