Skip to content

Commit f9732b7

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 f9732b7

4 files changed

Lines changed: 140 additions & 41 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: 28 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,21 @@ 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+
using mask_register_type = typename batch_bool<T, A>::register_type;
1222+
mask_register_type nmask = mask.to_native();
1223+
batch_bool<T, A> lo_mask = zip_lo(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
1224+
batch_bool<T, A> hi_mask = zip_hi(batch<T, A>(nmask), batch<T, A>(nmask)).to_native();
1225+
1226+
auto src_lo = detail::complex_low(src, A {});
1227+
auto src_hi = detail::complex_high(src, A {});
1228+
store_masked(reinterpret_cast<T*>(mem), src_lo, lo_mask, mode, A {});
1229+
store_masked(reinterpret_cast<T*>(mem) + src.size, src_hi, hi_mask, mode, A {});
1230+
}
1231+
12041232
namespace detail
12051233
{
12061234
// Reinterpret a constant-mask 4/8-byte load/store as same-width float

‎include/xsimd/arch/xsimd_avx512f.hpp‎

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

375+
namespace detail
376+
{
377+
// complex_low
378+
template <class A>
379+
XSIMD_INLINE batch<float, A> complex_low(batch<std::complex<float>, A> const& self, requires_arch<avx512f>) noexcept
380+
{
381+
__m512i idx = _mm512_setr_epi32(0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21, 6, 22, 7, 23);
382+
return _mm512_permutex2var_ps(self.real(), idx, self.imag());
383+
}
384+
template <class A>
385+
XSIMD_INLINE batch<double, A> complex_low(batch<std::complex<double>, A> const& self, requires_arch<avx512f>) noexcept
386+
{
387+
__m512i idx = _mm512_setr_epi64(0, 8, 1, 9, 2, 10, 3, 11);
388+
return _mm512_permutex2var_pd(self.real(), idx, self.imag());
389+
}
390+
391+
// complex_high
392+
template <class A>
393+
XSIMD_INLINE batch<float, A> complex_high(batch<std::complex<float>, A> const& self, requires_arch<avx512f>) noexcept
394+
{
395+
__m512i idx = _mm512_setr_epi32(8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13, 29, 14, 30, 15, 31);
396+
return _mm512_permutex2var_ps(self.real(), idx, self.imag());
397+
}
398+
template <class A>
399+
XSIMD_INLINE batch<double, A> complex_high(batch<std::complex<double>, A> const& self, requires_arch<avx512f>) noexcept
400+
{
401+
__m512i idx = _mm512_setr_epi64(4, 12, 5, 13, 6, 14, 7, 15);
402+
return _mm512_permutex2var_pd(self.real(), idx, self.imag());
403+
}
404+
}
405+
406+
namespace detail
407+
{
408+
template <class A, class T>
409+
std::array<batch_bool<T, A>, 2> zip_complex_mask(batch_bool<T, A> mask)
410+
{
411+
using mask_register_type = typename batch_bool<T, A>::register_type;
412+
mask_register_type nmask = mask.to_native();
413+
414+
constexpr mask_register_type lo_bitmask = xsimd::utils::make_low_mask<mask_register_type>(mask.size / 2);
415+
mask_register_type lo_mask = nmask & lo_bitmask;
416+
lo_mask = detail::interleave(lo_mask);
417+
418+
constexpr mask_register_type hi_bitmask = lo_bitmask << (mask.size / 2);
419+
mask_register_type hi_mask = (nmask & hi_bitmask) >> (mask.size / 2);
420+
hi_mask = detail::interleave(hi_mask);
421+
422+
return { batch_bool<T, A> { lo_mask }, batch_bool<T, A> { hi_mask } };
423+
}
424+
}
425+
426+
template <class A, class T, class Mode>
427+
XSIMD_INLINE void
428+
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
429+
{
430+
auto [lo_mask, hi_mask] = detail::zip_complex_mask(mask);
431+
432+
auto src_lo = detail::complex_low(src, A {});
433+
auto src_hi = detail::complex_high(src, A {});
434+
store_masked(reinterpret_cast<T*>(mem), src_lo, lo_mask, mode, A {});
435+
store_masked(reinterpret_cast<T*>(mem) + src.size, src_hi, hi_mask, mode, A {});
436+
}
437+
375438
// abs
376439
template <class A>
377440
XSIMD_INLINE batch<float, A> abs(batch<float, A> const& self, requires_arch<avx512f>) noexcept
@@ -972,36 +1035,6 @@ namespace xsimd
9721035
}
9731036
}
9741037

975-
namespace detail
976-
{
977-
// complex_low
978-
template <class A>
979-
XSIMD_INLINE batch<float, A> complex_low(batch<std::complex<float>, A> const& self, requires_arch<avx512f>) noexcept
980-
{
981-
__m512i idx = _mm512_setr_epi32(0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21, 6, 22, 7, 23);
982-
return _mm512_permutex2var_ps(self.real(), idx, self.imag());
983-
}
984-
template <class A>
985-
XSIMD_INLINE batch<double, A> complex_low(batch<std::complex<double>, A> const& self, requires_arch<avx512f>) noexcept
986-
{
987-
__m512i idx = _mm512_setr_epi64(0, 8, 1, 9, 2, 10, 3, 11);
988-
return _mm512_permutex2var_pd(self.real(), idx, self.imag());
989-
}
990-
991-
// complex_high
992-
template <class A>
993-
XSIMD_INLINE batch<float, A> complex_high(batch<std::complex<float>, A> const& self, requires_arch<avx512f>) noexcept
994-
{
995-
__m512i idx = _mm512_setr_epi32(8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13, 29, 14, 30, 15, 31);
996-
return _mm512_permutex2var_ps(self.real(), idx, self.imag());
997-
}
998-
template <class A>
999-
XSIMD_INLINE batch<double, A> complex_high(batch<std::complex<double>, A> const& self, requires_arch<avx512f>) noexcept
1000-
{
1001-
__m512i idx = _mm512_setr_epi64(4, 12, 5, 13, 6, 14, 7, 15);
1002-
return _mm512_permutex2var_pd(self.real(), idx, self.imag());
1003-
}
1004-
}
10051038
// incr_if
10061039
template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
10071040
XSIMD_INLINE batch<T, A> decr_if(batch<T, A> const& self, batch_bool<T, A> const& mask, requires_arch<avx512f>) noexcept
@@ -1633,6 +1666,16 @@ namespace xsimd
16331666
}
16341667
}
16351668

1669+
template <class A, class T, class Mode>
1670+
XSIMD_INLINE batch<std::complex<T>, A>
1671+
load_complex_masked(std::complex<T> const* mem, batch_bool<T, A> mask, Mode mode, requires_arch<avx512f>) noexcept
1672+
{
1673+
auto [lo_mask, hi_mask] = detail::zip_complex_mask(mask);
1674+
batch<T, A> res_lo = batch<T, A>::load(reinterpret_cast<T const*>(mem), lo_mask, mode);
1675+
batch<T, A> res_hi = batch<T, A>::load(reinterpret_cast<T const*>(mem) + mask.size, hi_mask, mode);
1676+
return detail::load_complex(res_lo, res_hi, A {});
1677+
}
1678+
16361679
// load_unaligned
16371680
template <class A, class T, class = std::enable_if_t<std::is_integral_v<T>>>
16381681
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)