Skip to content

Commit 82ff60f

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 28f6c62 commit 82ff60f

4 files changed

Lines changed: 130 additions & 41 deletions

File tree

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,19 @@ 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<std::complex<T>, size> buffer;
461+
for (std::size_t i = 0; i < size; ++i)
462+
buffer[i] = mask.get(i) ? mem[i] : std::complex<T>(0);
463+
return batch<std::complex<T>, A>::load_aligned(buffer.data());
464+
}
465+
453466
template <class A, class T_in, class T_out, bool... Values, class alignment>
454467
XSIMD_INLINE void
455468
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 +878,18 @@ namespace xsimd
865878
store_complex_aligned<A>(dst, src, A {});
866879
}
867880

881+
template <class A, class T, class Mode>
882+
XSIMD_INLINE void
883+
store_complex_masked(std::complex<T>* mem, batch<std::complex<T>, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
884+
{
885+
constexpr std::size_t size = batch<T, A>::size;
886+
alignas(A::alignment()) std::array<std::complex<T>, size> buffer;
887+
src.store_aligned(buffer.data());
888+
for (std::size_t i = 0; i < size; ++i)
889+
if (mask.get(i))
890+
mem[i] = buffer[i];
891+
}
892+
868893
// transpose
869894
template <class A, class T>
870895
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
@@ -373,6 +373,69 @@ namespace xsimd
373373
detail::store_masked(mem, src, mask.mask(), Mode {});
374374
}
375375

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

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

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

14981498
template <class T, class A>
14991499
template <class Mode>
1500-
XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode) const noexcept
1500+
XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode mode) const noexcept
15011501
{
1502-
alignas(A::alignment()) std::array<value_type, size> buffer;
1503-
store_aligned(buffer.data());
1504-
for (std::size_t i = 0; i < size; ++i)
1505-
if (mask.get(i))
1506-
mem[i] = buffer[i];
1502+
kernel::store_complex_masked<A>(mem, *this, mask, mode, A {});
15071503
}
15081504

15091505
template <class T, class A>
@@ -1546,12 +1542,9 @@ namespace xsimd
15461542

15471543
template <class T, class A>
15481544
template <class Mode>
1549-
XSIMD_INLINE batch<std::complex<T>, A> batch<std::complex<T>, A>::load(value_type const* mem, batch_bool<T, A> mask, Mode) noexcept
1545+
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
15501546
{
1551-
alignas(A::alignment()) std::array<value_type, size> buffer {};
1552-
for (std::size_t i = 0; i < size; ++i)
1553-
buffer[i] = mask.get(i) ? mem[i] : value_type(0);
1554-
return load_aligned(buffer.data());
1547+
return kernel::load_complex_masked<A>(mem, mask, mode, A {});
15551548
}
15561549

15571550
template <class T, class A>

0 commit comments

Comments
 (0)