Skip to content

Commit 8af8371

Browse files
DiamonDinoiaDiamonDinoia
committed
feat: dynamic batch_bool masks + avx_128 / avx2_128 mask overloads
Adds runtime batch_bool mask overloads of xsimd::load_masked and xsimd::store_masked across AVX, AVX2, AVX-512, SSE, SVE, RVV, and NEON; generic common-path fallback collapsed to a whole-vector select. SVE compile-time masked load/store forwarded through the runtime path so the per-lane predicate is correct on SVE wider than 128 bits. Adds arch-specific runtime-mask overloads of load_masked / store_masked for the avx_128 and avx2_128 arches so they inherit the hardware predicated load/store path on x86. Squashed from: b57a766 feat: add runtime batch_bool mask overloads for load_masked/store_masked d5f21c7 feat: add runtime batch_bool mask overloads for avx_128 / avx2_128
1 parent f320683 commit 8af8371

13 files changed

Lines changed: 422 additions & 47 deletions

‎docs/source/api/data_transfer.rst‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ Data Transfers
1212
From memory:
1313

1414
+---------------------------------------+----------------------------------------------------+
15-
| :cpp:func:`load` | load values from memory (optionally masked) |
15+
| :cpp:func:`load` | load values from memory (optionally masked) [#m]_ |
1616
+---------------------------------------+----------------------------------------------------+
1717
| :cpp:func:`load_aligned` | load values from aligned memory |
1818
+---------------------------------------+----------------------------------------------------+
@@ -32,7 +32,7 @@ From a scalar:
3232
To memory:
3333

3434
+---------------------------------------+----------------------------------------------------+
35-
| :cpp:func:`store` | store values to memory (optionally masked) |
35+
| :cpp:func:`store` | store values to memory (optionally masked) [#m]_ |
3636
+---------------------------------------+----------------------------------------------------+
3737
| :cpp:func:`store_aligned` | store values to aligned memory |
3838
+---------------------------------------+----------------------------------------------------+
@@ -84,3 +84,11 @@ The following empty types are used for tag dispatching:
8484

8585
.. doxygenstruct:: xsimd::unaligned_mode
8686
:project: xsimd
87+
88+
.. rubric:: Footnotes
89+
90+
.. [#m] Masked ``load`` / ``store`` come in two flavours. The
91+
:cpp:class:`batch_bool_constant` overload encodes the mask in the type and
92+
is resolved at compile time. The runtime :cpp:class:`batch_bool` overload
93+
accepts a mask computed at runtime. Prefer the compile-time mask whenever
94+
the selection is known at compile time.

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -437,13 +437,48 @@ namespace xsimd
437437
return detail::load_masked_common(mem, mask, cvt, mode, detail::masked_memory_uses_fp_bitcast<A, T_in, T_out> {});
438438
}
439439

440+
template <class A, class T, class Mode>
441+
XSIMD_INLINE batch<T, A>
442+
load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<common>) noexcept
443+
{
444+
// Per-lane validity contract: only active lanes of ``mem`` are
445+
// required to be addressable. An unconditional whole-vector load
446+
// would touch inactive lanes and trip ASan/Valgrind on partial
447+
// buffers, so stay scalar. Arches with hardware predicated loads
448+
// (AVX2 32/64-bit, AVX-512, SVE, RVV) override this with a single
449+
// intrinsic that suppresses inactive-lane reads in hardware.
450+
constexpr std::size_t size = batch<T, A>::size;
451+
alignas(A::alignment()) std::array<T, size> buffer;
452+
for (std::size_t i = 0; i < size; ++i)
453+
buffer[i] = mask.get(i) ? mem[i] : T(0);
454+
return batch<T, A>::load_aligned(buffer.data());
455+
}
456+
440457
template <class A, class T_in, class T_out, bool... Values, class alignment>
441458
XSIMD_INLINE void
442459
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
443460
{
444461
detail::store_masked_common(mem, src, mask, mode, detail::masked_memory_uses_fp_bitcast<A, T_in, T_out> {});
445462
}
446463

464+
template <class A, class T, class Mode>
465+
XSIMD_INLINE void
466+
store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
467+
{
468+
// Per-lane validity contract (matches native masked-store APIs):
469+
// only active lanes of ``mem`` are touched. A load+select+store
470+
// RMW would both read and write inactive bytes, breaking that
471+
// contract — stay scalar. Arches with hardware predicated stores
472+
// override this with a single intrinsic that suppresses inactive
473+
// lanes in hardware.
474+
constexpr std::size_t size = batch<T, A>::size;
475+
alignas(A::alignment()) std::array<T, size> src_buf;
476+
src.store_aligned(src_buf.data());
477+
for (std::size_t i = 0; i < size; ++i)
478+
if (mask.get(i))
479+
mem[i] = src_buf[i];
480+
}
481+
447482
template <class A, class T_in, class T_out>
448483
XSIMD_INLINE batch<T_out, A> load_stream(T_in const* mem, convert<T_out> cvt, requires_arch<common>) noexcept
449484
{

‎include/xsimd/arch/xsimd_avx.hpp‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -987,6 +987,23 @@ namespace xsimd
987987
}
988988
}
989989

990+
// Runtime-mask load for float/double on AVX. Both aligned_mode and
991+
// unaligned_mode map to _mm256_maskload_* — the intrinsic does not fault
992+
// on masked-off lanes, so partial loads across page boundaries are safe.
993+
template <class A, class Mode>
994+
XSIMD_INLINE batch<float, A>
995+
load_masked(float const* mem, batch_bool<float, A> mask, convert<float>, Mode, requires_arch<avx>) noexcept
996+
{
997+
return _mm256_maskload_ps(mem, _mm256_castps_si256(mask));
998+
}
999+
1000+
template <class A, class Mode>
1001+
XSIMD_INLINE batch<double, A>
1002+
load_masked(double const* mem, batch_bool<double, A> mask, convert<double>, Mode, requires_arch<avx>) noexcept
1003+
{
1004+
return _mm256_maskload_pd(mem, _mm256_castpd_si256(mask));
1005+
}
1006+
9901007
// load_masked (single overload for float/double)
9911008
template <class A, class T, bool... Values, class Mode, class = std::enable_if_t<std::is_floating_point<T>::value>>
9921009
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode, requires_arch<avx>) noexcept
@@ -1070,6 +1087,22 @@ namespace xsimd
10701087
}
10711088
}
10721089

1090+
// Runtime-mask store for float/double on AVX. Same fault-suppression
1091+
// semantics as the masked loads above; alignment mode is irrelevant.
1092+
template <class A, class Mode>
1093+
XSIMD_INLINE void
1094+
store_masked(float* mem, batch<float, A> const& src, batch_bool<float, A> mask, Mode, requires_arch<avx>) noexcept
1095+
{
1096+
detail::maskstore(mem, mask, src);
1097+
}
1098+
1099+
template <class A, class Mode>
1100+
XSIMD_INLINE void
1101+
store_masked(double* mem, batch<double, A> const& src, batch_bool<double, A> mask, Mode, requires_arch<avx>) noexcept
1102+
{
1103+
detail::maskstore(mem, mask, src);
1104+
}
1105+
10731106
// lt
10741107
template <class A>
10751108
XSIMD_INLINE batch_bool<float, A> lt(batch<float, A> const& self, batch<float, A> const& other, requires_arch<avx>) noexcept

‎include/xsimd/arch/xsimd_avx2.hpp‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,17 @@ namespace xsimd
174174
return bitwise_cast<uint64_t>(r);
175175
}
176176

177+
// Runtime-mask load for 32/64-bit integers on AVX2; narrower widths fall
178+
// back to the scalar common path. Aligned and unaligned share the same
179+
// intrinsic — masked-off lanes do not fault regardless of alignment.
180+
template <class A, class T, class Mode>
181+
XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), batch<T, A>>
182+
load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<avx2>) noexcept
183+
{
184+
using int_t = std::conditional_t<sizeof(T) == 4, int32_t, long long>;
185+
return detail::maskload(reinterpret_cast<const int_t*>(mem), __m256i(mask));
186+
}
187+
177188
// store_masked
178189
namespace detail
179190
{
@@ -232,6 +243,17 @@ namespace xsimd
232243
store_masked<A>(reinterpret_cast<int64_t*>(mem), s64, batch_bool_constant<int64_t, A, Values...> {}, Mode {}, avx2 {});
233244
}
234245

246+
// Runtime-mask store for 32/64-bit integers on AVX2; narrower widths fall
247+
// back to the scalar common path. Same fault-suppression semantics as the
248+
// masked loads above; alignment mode is irrelevant.
249+
template <class A, class T, class Mode>
250+
XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), void>
251+
store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<avx2>) noexcept
252+
{
253+
using int_t = std::conditional_t<sizeof(T) == 4, int32_t, int64_t>;
254+
detail::maskstore<int_t, A>(reinterpret_cast<int_t*>(mem), __m256i(mask), __m256i(src));
255+
}
256+
235257
// load_stream
236258
template <class A, class T, class = std::enable_if_t<std::is_integral<T>::value, void>>
237259
XSIMD_INLINE batch<T, A> load_stream(T const* mem, convert<T>, requires_arch<avx2>) noexcept

‎include/xsimd/arch/xsimd_avx2_128.hpp‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,45 @@ namespace xsimd
137137
return _mm_maskstore_epi64(reinterpret_cast<long long*>(mem), mask.as_batch(), src);
138138
}
139139

140+
// Runtime-mask path for 32/64-bit integers; narrower widths fall back to
141+
// the common scalar path. Aligned and unaligned share the same intrinsic
142+
// — masked-off lanes do not fault regardless of alignment.
143+
namespace detail
144+
{
145+
XSIMD_INLINE __m128i maskload_128(int32_t const* mem, __m128i mask) noexcept
146+
{
147+
return _mm_maskload_epi32(mem, mask);
148+
}
149+
XSIMD_INLINE __m128i maskload_128(long long const* mem, __m128i mask) noexcept
150+
{
151+
return _mm_maskload_epi64(mem, mask);
152+
}
153+
XSIMD_INLINE void maskstore_128(int32_t* mem, __m128i mask, __m128i src) noexcept
154+
{
155+
_mm_maskstore_epi32(mem, mask, src);
156+
}
157+
XSIMD_INLINE void maskstore_128(long long* mem, __m128i mask, __m128i src) noexcept
158+
{
159+
_mm_maskstore_epi64(mem, mask, src);
160+
}
161+
}
162+
163+
template <class A, class T, class Mode>
164+
XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), batch<T, A>>
165+
load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<avx2_128>) noexcept
166+
{
167+
using int_t = std::conditional_t<sizeof(T) == 4, int32_t, long long>;
168+
return detail::maskload_128(reinterpret_cast<int_t const*>(mem), __m128i(mask));
169+
}
170+
171+
template <class A, class T, class Mode>
172+
XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), void>
173+
store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<avx2_128>) noexcept
174+
{
175+
using int_t = std::conditional_t<sizeof(T) == 4, int32_t, long long>;
176+
detail::maskstore_128(reinterpret_cast<int_t*>(mem), __m128i(mask), __m128i(src));
177+
}
178+
140179
// gather
141180
template <class T, class A, class U, detail::enable_sized_integral_t<T, 4> = 0, detail::enable_sized_integral_t<U, 4> = 0>
142181
XSIMD_INLINE batch<T, A> gather(batch<T, A> const&, T const* src, batch<U, A> const& index,

‎include/xsimd/arch/xsimd_avx_128.hpp‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,22 @@ namespace xsimd
115115
return _mm_maskload_pd(mem, mask.as_batch());
116116
}
117117

118+
// Runtime-mask load for float/double on AVX-128. Both aligned_mode and
119+
// unaligned_mode map to _mm_maskload_* — the intrinsic does not fault
120+
// on masked-off lanes, so partial loads across page boundaries are safe.
121+
template <class A, class Mode>
122+
XSIMD_INLINE batch<float, A>
123+
load_masked(float const* mem, batch_bool<float, A> mask, convert<float>, Mode, requires_arch<avx_128>) noexcept
124+
{
125+
return _mm_maskload_ps(mem, _mm_castps_si128(mask));
126+
}
127+
template <class A, class Mode>
128+
XSIMD_INLINE batch<double, A>
129+
load_masked(double const* mem, batch_bool<double, A> mask, convert<double>, Mode, requires_arch<avx_128>) noexcept
130+
{
131+
return _mm_maskload_pd(mem, _mm_castpd_si128(mask));
132+
}
133+
118134
// store_masked
119135
template <class A, bool... Values, class Mode>
120136
XSIMD_INLINE void store_masked(float* mem, batch<float, A> const& src, batch_bool_constant<float, A, Values...> mask, Mode, requires_arch<avx_128>) noexcept
@@ -128,6 +144,21 @@ namespace xsimd
128144
return _mm_maskstore_pd(mem, mask.as_batch(), src);
129145
}
130146

147+
// Runtime-mask store for float/double on AVX-128. Same fault-suppression
148+
// semantics as the masked loads above; alignment mode is irrelevant.
149+
template <class A, class Mode>
150+
XSIMD_INLINE void
151+
store_masked(float* mem, batch<float, A> const& src, batch_bool<float, A> mask, Mode, requires_arch<avx_128>) noexcept
152+
{
153+
_mm_maskstore_ps(mem, _mm_castps_si128(mask), src);
154+
}
155+
template <class A, class Mode>
156+
XSIMD_INLINE void
157+
store_masked(double* mem, batch<double, A> const& src, batch_bool<double, A> mask, Mode, requires_arch<avx_128>) noexcept
158+
{
159+
_mm_maskstore_pd(mem, _mm_castpd_si128(mask), src);
160+
}
161+
131162
// swizzle (dynamic mask)
132163
template <class A, class ITy>
133164
XSIMD_INLINE batch<float, A> swizzle(batch<float, A> const& self, batch<ITy, A> mask, requires_arch<avx_128>) noexcept

‎include/xsimd/arch/xsimd_common_fwd.hpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,12 @@ namespace xsimd
8888
XSIMD_INLINE batch<T, A> load(T const* mem, unaligned_mode, requires_arch<A>) noexcept;
8989
template <class A, class T_in, class T_out, bool... Values, class alignment>
9090
XSIMD_INLINE batch<T_out, A> load_masked(T_in const* mem, batch_bool_constant<T_out, A, Values...> mask, convert<T_out>, alignment, requires_arch<common>) noexcept;
91+
template <class A, class T, class Mode>
92+
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<common>) noexcept;
9193
template <class A, class T_in, class T_out, bool... Values, class alignment>
9294
XSIMD_INLINE void store_masked(T_out* mem, batch<T_in, A> const& src, batch_bool_constant<T_in, A, Values...> mask, alignment, requires_arch<common>) noexcept;
95+
template <class A, class T, class Mode>
96+
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept;
9397

9498
// Forward declarations for pack-level helpers
9599
namespace detail

‎include/xsimd/arch/xsimd_rvv.hpp‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,11 @@ namespace xsimd
409409
{
410410
XSIMD_RVV_OVERLOAD(rvvle, (__riscv_vle XSIMD_RVV_S _v_ XSIMD_RVV_TSM), , vec(T const*))
411411
XSIMD_RVV_OVERLOAD(rvvse, (__riscv_vse XSIMD_RVV_S _v_ XSIMD_RVV_TSM), , void(T*, vec))
412+
// Masked load (mask-undisturbed with zero passthrough): inactive lanes read as 0,
413+
// no memory access is performed for inactive lanes (page-fault safe).
414+
XSIMD_RVV_OVERLOAD(rvvle_mu, (__riscv_vle XSIMD_RVV_S _v_ XSIMD_RVV_TSM _mu), , vec(bvec, vec, T const*))
415+
// Masked store: inactive lanes are not written.
416+
XSIMD_RVV_OVERLOAD(rvvse_m, (__riscv_vse XSIMD_RVV_S _v_ XSIMD_RVV_TSM _m), , void(bvec, T*, vec))
412417
}
413418

414419
template <class A, class T, detail::enable_arithmetic_t<T> = 0>
@@ -423,6 +428,16 @@ namespace xsimd
423428
return load_aligned<A>(src, convert<T>(), rvv {});
424429
}
425430

431+
// load_masked (runtime mask): native vle*.v vd, (rs1), v0.t with zero-init
432+
// passthrough so inactive lanes read as 0, matching xsimd's contract.
433+
template <class A, class T, class Mode, detail::enable_arithmetic_t<T> = 0>
434+
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<rvv>) noexcept
435+
{
436+
using proj_t = map_to_sized_type_t<T>;
437+
const auto zero = detail_rvv::rvvmv_splat(proj_t {});
438+
return detail_rvv::rvvle_mu(mask, zero, reinterpret_cast<proj_t const*>(mem));
439+
}
440+
426441
// load_complex
427442
namespace detail_rvv
428443
{
@@ -500,6 +515,15 @@ namespace xsimd
500515
store_aligned<A>(dst, src, rvv {});
501516
}
502517

518+
// store_masked (runtime mask): native vse*.v vd, (rs1), v0.t — inactive lanes
519+
// are not written (page-fault safe).
520+
template <class A, class T, class Mode, detail::enable_arithmetic_t<T> = 0>
521+
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<rvv>) noexcept
522+
{
523+
using proj_t = map_to_sized_type_t<T>;
524+
detail_rvv::rvvse_m(mask, reinterpret_cast<proj_t*>(mem), src);
525+
}
526+
503527
/******************
504528
* scatter/gather *
505529
******************/

‎include/xsimd/arch/xsimd_sve.hpp‎

Lines changed: 39 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,11 +101,28 @@ namespace xsimd
101101
return load_aligned<A>(src, convert<T>(), sve {});
102102
}
103103

104-
// load_masked
105-
template <class A, class T, bool... Values, class Mode, detail::enable_arithmetic_t<T> = 0>
106-
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<float, A, Values...>, Mode, requires_arch<sve>) noexcept
104+
// load_masked (compile-time mask): build a runtime predicate from
105+
// the constant mask and reuse the runtime-mask path. ``pmask`` only
106+
// constructs a 128-bit chunk predicate (svdupq_b{8,16,32,64}), which
107+
// is replication-based and does not correctly express a per-lane
108+
// mask on SVE wider than 128 bits — going through ``as_batch_bool``
109+
// gives the right predicate for every vector width. ``int32``/
110+
// ``int64``/``uint32``/``uint64`` are excluded so the common-arch
111+
// dispatchers that reinterpret to ``float``/``double`` win partial
112+
// ordering (otherwise we'd be ambiguous with ``requires_arch<A>``).
113+
template <class A, class T, bool... Values, class Mode,
114+
detail::enable_arithmetic_t<T> = 0,
115+
std::enable_if_t<!(std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)), int> = 0>
116+
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode m, requires_arch<sve>) noexcept
107117
{
108-
return svld1(detail_sve::pmask<Values...>(), reinterpret_cast<map_to_sized_type_t<T> const*>(mem));
118+
return load_masked<A>(mem, mask.as_batch_bool(), convert<T> {}, m, sve {});
119+
}
120+
121+
// load_masked (runtime mask)
122+
template <class A, class T, class Mode, detail::enable_arithmetic_t<T> = 0>
123+
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<sve>) noexcept
124+
{
125+
return svld1(mask, reinterpret_cast<map_to_sized_type_t<T> const*>(mem));
109126
}
110127

111128
// load_complex
@@ -141,6 +158,24 @@ namespace xsimd
141158
store_aligned<A>(dst, src, sve {});
142159
}
143160

161+
// store_masked (compile-time mask): forward to the runtime-mask
162+
// path for the same reason as load_masked above; same exclusion of
163+
// 32/64-bit integers to defer to the common dispatchers.
164+
template <class A, class T, bool... Values, class Mode,
165+
detail::enable_arithmetic_t<T> = 0,
166+
std::enable_if_t<!(std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)), int> = 0>
167+
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, Values...> mask, Mode m, requires_arch<sve>) noexcept
168+
{
169+
store_masked<A>(mem, src, mask.as_batch_bool(), m, sve {});
170+
}
171+
172+
// store_masked (runtime mask)
173+
template <class A, class T, class Mode, detail::enable_arithmetic_t<T> = 0>
174+
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<sve>) noexcept
175+
{
176+
svst1(mask, reinterpret_cast<map_to_sized_type_t<T>*>(mem), src);
177+
}
178+
144179
// store_complex
145180
template <class A, class T, detail::enable_floating_point_t<T> = 0>
146181
XSIMD_INLINE void store_complex_aligned(std::complex<T>* dst, batch<std::complex<T>, A> const& src, requires_arch<sve>) noexcept

0 commit comments

Comments
 (0)