Skip to content

Commit fe28329

Browse files
serge-sans-pailleserge-sans-paille
committed
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 fe28329

12 files changed

Lines changed: 45 additions & 6 deletions

‎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 (A::has_scalar_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_avx512f_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ namespace xsimd
2929
static constexpr bool available() noexcept { return true; }
3030
static constexpr std::size_t alignment() noexcept { return 64; }
3131
static constexpr bool requires_alignment() noexcept { return true; }
32+
static constexpr bool has_scalar_mask() noexcept { return true; }
3233
static constexpr char const* name() noexcept { return "avx512f"; }
3334
};
3435

‎include/xsimd/types/xsimd_avx_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ namespace xsimd
3030
static constexpr std::size_t alignment() noexcept { return 32; }
3131
static constexpr bool requires_alignment() noexcept { return true; }
3232
static constexpr char const* name() noexcept { return "avx"; }
33+
static constexpr bool has_scalar_mask() noexcept { return false; }
3334
};
3435

3536
/**

‎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>

‎include/xsimd/types/xsimd_emulated_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ namespace xsimd
3838
static constexpr bool requires_alignment() noexcept { return false; }
3939
static constexpr std::size_t alignment() noexcept { return 8; }
4040
static constexpr char const* name() noexcept { return "emulated"; }
41+
static constexpr bool has_scalar_mask() noexcept { return false; }
4142
};
4243

4344
namespace types

‎include/xsimd/types/xsimd_neon_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ namespace xsimd
3939
static constexpr bool requires_alignment() noexcept { return true; }
4040
static constexpr std::size_t alignment() noexcept { return 16; }
4141
static constexpr char const* name() noexcept { return "arm32+neon"; }
42+
static constexpr bool has_scalar_mask() noexcept { return true; }
4243
};
4344

4445
#if XSIMD_WITH_NEON

‎include/xsimd/types/xsimd_rvv_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ namespace xsimd
3939
static constexpr bool requires_alignment() noexcept { return true; }
4040
static constexpr std::size_t alignment() noexcept { return 16; }
4141
static constexpr char const* name() noexcept { return "riscv+rvv"; }
42+
static constexpr bool has_scalar_mask() noexcept { return true; }
4243
};
4344
}
4445

‎include/xsimd/types/xsimd_sse2_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ namespace xsimd
3434
static constexpr bool requires_alignment() noexcept { return true; }
3535
static constexpr std::size_t alignment() noexcept { return 16; }
3636
static constexpr char const* name() noexcept { return "sse2"; }
37+
static constexpr bool has_scalar_mask() noexcept { return false; }
3738
};
3839

3940
#if XSIMD_WITH_SSE2

‎include/xsimd/types/xsimd_sve_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ namespace xsimd
3737
static constexpr bool requires_alignment() noexcept { return true; }
3838
static constexpr std::size_t alignment() noexcept { return 16; }
3939
static constexpr char const* name() noexcept { return "arm64+sve"; }
40+
static constexpr bool has_scalar_mask() noexcept { return true; }
4041
};
4142
}
4243

‎include/xsimd/types/xsimd_vsx_register.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ namespace xsimd
3333
static constexpr bool requires_alignment() noexcept { return true; }
3434
static constexpr std::size_t alignment() noexcept { return 16; }
3535
static constexpr char const* name() noexcept { return "vmx+vsx"; }
36+
static constexpr bool has_scalar_mask() noexcept { return true; }
3637
};
3738

3839
#if XSIMD_WITH_VSX

0 commit comments

Comments
 (0)