Skip to content

Commit 6e28269

Browse files
DiamonDinoiaDiamonDinoia
committed
perf: native AVX512BW masked load/store for 8/16-bit integers
8/16-bit int masked load/store on AVX512BW previously fell through to the branchy common scalar fallback because xsimd_avx512bw.hpp had no load_masked/store_masked overloads. Add four requires_arch<avx512bw> overloads (runtime batch_bool + compile-time batch_bool_constant, load + store) constrained to sizeof(T)==1||2, emitting the native vmovdqu8 / vmovdqu16 predicated moves (2 instructions, no branch). The size branch lives only in the runtime overloads; the constant overloads delegate via mask.as_batch_bool(), which also avoids batch_bool_constant::mask() (return type int) truncating a 64-lane int8 compile-time mask. 32/64-bit stays on the avx512f path; SSE/AVX2 8/16-bit scalar fallback is hardware-forced and unchanged.
1 parent 262f5a7 commit 6e28269

8 files changed

Lines changed: 124 additions & 9 deletions

File tree

‎docs/source/api/data_transfer.rst‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,5 +90,5 @@ The following empty types are used for tag dispatching:
9090
.. [#m] Masked ``load`` / ``store`` come in two flavours. The
9191
:cpp:class:`batch_bool_constant` overload encodes the mask in the type and
9292
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.
93+
accepts a mask computed at runtime. For performance reasons, prefer the
94+
compile-time mask whenever possible.

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,7 @@ namespace xsimd
442442
load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<common>) noexcept
443443
{
444444
// Scalar fallback: only active lanes are touched. Arches with
445-
// hardware predicated loads override this.
445+
// hardware predicated loads should override this.
446446
constexpr std::size_t size = batch<T, A>::size;
447447
alignas(A::alignment()) std::array<T, size> buffer;
448448
for (std::size_t i = 0; i < size; ++i)
@@ -462,7 +462,7 @@ namespace xsimd
462462
store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
463463
{
464464
// Scalar fallback: only active lanes are touched. Arches with
465-
// hardware predicated stores override this.
465+
// hardware predicated stores should override this.
466466
constexpr std::size_t size = batch<T, A>::size;
467467
alignas(A::alignment()) std::array<T, size> src_buf;
468468
src.store_aligned(src_buf.data());

‎include/xsimd/arch/xsimd_avx2.hpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,10 +126,12 @@ namespace xsimd
126126
{
127127
XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
128128
{
129+
static_assert(sizeof(int) == 4, "_mm256_maskload_epi32 requires a 4-byte int");
129130
return _mm256_maskload_epi32(reinterpret_cast<int const*>(mem), mask);
130131
}
131132
else
132133
{
134+
static_assert(sizeof(long long) == 8, "_mm256_maskload_epi64 requires an 8-byte long long");
133135
return _mm256_maskload_epi64(reinterpret_cast<long long const*>(mem), mask);
134136
}
135137
}
@@ -139,10 +141,12 @@ namespace xsimd
139141
{
140142
XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
141143
{
144+
static_assert(sizeof(int) == 4, "_mm256_maskstore_epi32 requires a 4-byte int");
142145
_mm256_maskstore_epi32(reinterpret_cast<int*>(mem), mask, src);
143146
}
144147
else
145148
{
149+
static_assert(sizeof(long long) == 8, "_mm256_maskstore_epi64 requires an 8-byte long long");
146150
_mm256_maskstore_epi64(reinterpret_cast<long long*>(mem), mask, src);
147151
}
148152
}

‎include/xsimd/arch/xsimd_avx512bw.hpp‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,53 @@ namespace xsimd
378378
}
379379
}
380380

381+
// load_masked / store_masked: native vmovdqu8 / vmovdqu16 predication for
382+
// 8/16-bit, replacing the common scalar fallback. No aligned masked 8/16
383+
// intrinsic exists and masked moves never fault, so loadu fits both modes.
384+
template <class A, class T, class Mode,
385+
class = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 1 || sizeof(T) == 2)>>
386+
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<avx512bw>) noexcept
387+
{
388+
XSIMD_IF_CONSTEXPR(sizeof(T) == 1)
389+
{
390+
return _mm512_maskz_loadu_epi8((__mmask64)mask.mask(), mem);
391+
}
392+
else
393+
{
394+
return _mm512_maskz_loadu_epi16((__mmask32)mask.mask(), mem);
395+
}
396+
}
397+
398+
template <class A, class T, class Mode,
399+
class = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 1 || sizeof(T) == 2)>>
400+
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<avx512bw>) noexcept
401+
{
402+
XSIMD_IF_CONSTEXPR(sizeof(T) == 1)
403+
{
404+
_mm512_mask_storeu_epi8((void*)mem, (__mmask64)mask.mask(), src);
405+
}
406+
else
407+
{
408+
_mm512_mask_storeu_epi16((void*)mem, (__mmask32)mask.mask(), src);
409+
}
410+
}
411+
412+
// Constant masks reuse the runtime overloads; as_batch_bool() also avoids
413+
// batch_bool_constant::mask() truncating a 64-lane int8 mask to int.
414+
template <class A, class T, bool... Values, class Mode,
415+
class = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 1 || sizeof(T) == 2)>>
416+
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode, requires_arch<avx512bw>) noexcept
417+
{
418+
return load_masked(mem, mask.as_batch_bool(), convert<T> {}, Mode {}, avx512bw {});
419+
}
420+
421+
template <class A, class T, bool... Values, class Mode,
422+
class = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 1 || sizeof(T) == 2)>>
423+
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, Values...> mask, Mode, requires_arch<avx512bw>) noexcept
424+
{
425+
store_masked(mem, src, mask.as_batch_bool(), Mode {}, avx512bw {});
426+
}
427+
381428
// max
382429
template <class A, class T, class = std::enable_if_t<std::is_integral<T>::value>>
383430
XSIMD_INLINE batch<T, A> max(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx512bw>) noexcept

‎include/xsimd/arch/xsimd_avx512f.hpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,8 @@ namespace xsimd
356356

357357
// Runtime-mask load/store: same native k-register path as the constant
358358
// overloads above, minus the compile-time half-forwarding. 8/16-bit
359-
// elements fall back to the common scalar path.
359+
// elements are handled natively by avx512bw (vmovdqu8 / vmovdqu16);
360+
// without AVX512BW they fall back to the common scalar path.
360361
template <class A, class T, class Mode,
361362
typename = std::enable_if_t<(sizeof(T) >= 4)>>
362363
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<avx512f>) noexcept

‎include/xsimd/arch/xsimd_avx512vl_128.hpp‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,66 +212,98 @@ namespace xsimd
212212
XSIMD_INLINE __m128i maskload128(T const* mem, uint64_t m, Mode) noexcept
213213
{
214214
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
215+
{
215216
return _mm_maskz_load_epi32((__mmask8)m, mem);
217+
}
216218
else
219+
{
217220
return _mm_maskz_loadu_epi32((__mmask8)m, mem);
221+
}
218222
}
219223
template <class T, class Mode, enable_sized_integral_t<T, 8> = 0>
220224
XSIMD_INLINE __m128i maskload128(T const* mem, uint64_t m, Mode) noexcept
221225
{
222226
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
227+
{
223228
return _mm_maskz_load_epi64((__mmask8)m, mem);
229+
}
224230
else
231+
{
225232
return _mm_maskz_loadu_epi64((__mmask8)m, mem);
233+
}
226234
}
227235
template <class Mode>
228236
XSIMD_INLINE __m128 maskload128(float const* mem, uint64_t m, Mode) noexcept
229237
{
230238
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
239+
{
231240
return _mm_maskz_load_ps((__mmask8)m, mem);
241+
}
232242
else
243+
{
233244
return _mm_maskz_loadu_ps((__mmask8)m, mem);
245+
}
234246
}
235247
template <class Mode>
236248
XSIMD_INLINE __m128d maskload128(double const* mem, uint64_t m, Mode) noexcept
237249
{
238250
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
251+
{
239252
return _mm_maskz_load_pd((__mmask8)m, mem);
253+
}
240254
else
255+
{
241256
return _mm_maskz_loadu_pd((__mmask8)m, mem);
257+
}
242258
}
243259

244260
template <class T, class Mode, enable_sized_integral_t<T, 4> = 0>
245261
XSIMD_INLINE void maskstore128(T* mem, __m128i src, uint64_t m, Mode) noexcept
246262
{
247263
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
264+
{
248265
_mm_mask_store_epi32(mem, (__mmask8)m, src);
266+
}
249267
else
268+
{
250269
_mm_mask_storeu_epi32(mem, (__mmask8)m, src);
270+
}
251271
}
252272
template <class T, class Mode, enable_sized_integral_t<T, 8> = 0>
253273
XSIMD_INLINE void maskstore128(T* mem, __m128i src, uint64_t m, Mode) noexcept
254274
{
255275
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
276+
{
256277
_mm_mask_store_epi64(mem, (__mmask8)m, src);
278+
}
257279
else
280+
{
258281
_mm_mask_storeu_epi64(mem, (__mmask8)m, src);
282+
}
259283
}
260284
template <class Mode>
261285
XSIMD_INLINE void maskstore128(float* mem, __m128 src, uint64_t m, Mode) noexcept
262286
{
263287
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
288+
{
264289
_mm_mask_store_ps(mem, (__mmask8)m, src);
290+
}
265291
else
292+
{
266293
_mm_mask_storeu_ps(mem, (__mmask8)m, src);
294+
}
267295
}
268296
template <class Mode>
269297
XSIMD_INLINE void maskstore128(double* mem, __m128d src, uint64_t m, Mode) noexcept
270298
{
271299
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
300+
{
272301
_mm_mask_store_pd(mem, (__mmask8)m, src);
302+
}
273303
else
304+
{
274305
_mm_mask_storeu_pd(mem, (__mmask8)m, src);
306+
}
275307
}
276308
}
277309

‎include/xsimd/arch/xsimd_avx512vl_256.hpp‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,66 +211,98 @@ namespace xsimd
211211
XSIMD_INLINE __m256i maskload256(T const* mem, uint64_t m, Mode) noexcept
212212
{
213213
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
214+
{
214215
return _mm256_maskz_load_epi32((__mmask8)m, mem);
216+
}
215217
else
218+
{
216219
return _mm256_maskz_loadu_epi32((__mmask8)m, mem);
220+
}
217221
}
218222
template <class T, class Mode, enable_sized_integral_t<T, 8> = 0>
219223
XSIMD_INLINE __m256i maskload256(T const* mem, uint64_t m, Mode) noexcept
220224
{
221225
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
226+
{
222227
return _mm256_maskz_load_epi64((__mmask8)m, mem);
228+
}
223229
else
230+
{
224231
return _mm256_maskz_loadu_epi64((__mmask8)m, mem);
232+
}
225233
}
226234
template <class Mode>
227235
XSIMD_INLINE __m256 maskload256(float const* mem, uint64_t m, Mode) noexcept
228236
{
229237
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
238+
{
230239
return _mm256_maskz_load_ps((__mmask8)m, mem);
240+
}
231241
else
242+
{
232243
return _mm256_maskz_loadu_ps((__mmask8)m, mem);
244+
}
233245
}
234246
template <class Mode>
235247
XSIMD_INLINE __m256d maskload256(double const* mem, uint64_t m, Mode) noexcept
236248
{
237249
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
250+
{
238251
return _mm256_maskz_load_pd((__mmask8)m, mem);
252+
}
239253
else
254+
{
240255
return _mm256_maskz_loadu_pd((__mmask8)m, mem);
256+
}
241257
}
242258

243259
template <class T, class Mode, enable_sized_integral_t<T, 4> = 0>
244260
XSIMD_INLINE void maskstore256(T* mem, __m256i src, uint64_t m, Mode) noexcept
245261
{
246262
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
263+
{
247264
_mm256_mask_store_epi32(mem, (__mmask8)m, src);
265+
}
248266
else
267+
{
249268
_mm256_mask_storeu_epi32(mem, (__mmask8)m, src);
269+
}
250270
}
251271
template <class T, class Mode, enable_sized_integral_t<T, 8> = 0>
252272
XSIMD_INLINE void maskstore256(T* mem, __m256i src, uint64_t m, Mode) noexcept
253273
{
254274
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
275+
{
255276
_mm256_mask_store_epi64(mem, (__mmask8)m, src);
277+
}
256278
else
279+
{
257280
_mm256_mask_storeu_epi64(mem, (__mmask8)m, src);
281+
}
258282
}
259283
template <class Mode>
260284
XSIMD_INLINE void maskstore256(float* mem, __m256 src, uint64_t m, Mode) noexcept
261285
{
262286
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
287+
{
263288
_mm256_mask_store_ps(mem, (__mmask8)m, src);
289+
}
264290
else
291+
{
265292
_mm256_mask_storeu_ps(mem, (__mmask8)m, src);
293+
}
266294
}
267295
template <class Mode>
268296
XSIMD_INLINE void maskstore256(double* mem, __m256d src, uint64_t m, Mode) noexcept
269297
{
270298
XSIMD_IF_CONSTEXPR(std::is_same<Mode, aligned_mode>::value)
299+
{
271300
_mm256_mask_store_pd(mem, (__mmask8)m, src);
301+
}
272302
else
303+
{
273304
_mm256_mask_storeu_pd(mem, (__mmask8)m, src);
305+
}
274306
}
275307
}
276308

‎include/xsimd/arch/xsimd_isa.hpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -81,18 +81,17 @@
8181
// unqualified dependent call, which clang only resolves through ordinary lookup
8282
// at the point of definition (ADL cannot reach xsimd::kernel from xsimd-namespace
8383
// arguments). The sub-arch overloads must therefore be declared beforehand.
84+
// clang-format off
8485
#include "./xsimd_avx512vl_128.hpp"
8586
#include "./xsimd_avx512vl_256.hpp"
87+
#include "./xsimd_avx512vl.hpp"
88+
// clang-format on
8689
#endif
8790

8891
#if XSIMD_WITH_AVX512F
8992
#include "./xsimd_avx512f.hpp"
9093
#endif
9194

92-
#if XSIMD_WITH_AVX512VL
93-
#include "./xsimd_avx512vl.hpp"
94-
#endif
95-
9695
#if XSIMD_WITH_AVX512DQ
9796
#include "./xsimd_avx512dq.hpp"
9897
#endif

0 commit comments

Comments
 (0)