Skip to content

perf(avx2/avx512): plain-move lowering for int & EVEX constant-mask load/store #1379

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DiamonDinoia wants to merge 1 commit into
base: master
Choose a base branch
from
Open

perf(avx2/avx512): plain-move lowering for int & EVEX constant-mask load/store #1379

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 59 additions & 0 deletions include/xsimd/arch/xsimd_avx.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "../types/xsimd_avx_register.hpp"
#include "../types/xsimd_batch_constant.hpp"

#include <cassert>
#include <complex>
#include <limits>
#include <type_traits>
Expand Down Expand Up @@ -1041,6 +1042,8 @@ namespace xsimd
// exactly the lower 128-bit half: one plain load, upper lanes zero
XSIMD_IF_CONSTEXPR(mask.prefix() == half_size)
{
// cross-check the plain move via countr_one/countl_zero (independent of prefix())
assert(mask.countr_one() >= half_size && mask.countl_zero() >= half_size && "lower half fully active, upper empty");
return detail::zero_extend_lo<A>(half_batch::load(mem, Mode {}));
}
// lower 128-bit half: stay in the value domain so the half kernel can
Expand All @@ -1051,9 +1054,21 @@ namespace xsimd
const auto lo = load_masked(mem, mlo, convert<T> {}, Mode {}, half_arch {});
return detail::zero_extend_lo<A>(lo);
}
// prefix crossing the 128-bit boundary: plain lower half +
// prefix-masked upper half (mirrors the store side)
else XSIMD_IF_CONSTEXPR(mask.prefix() > half_size && mask.prefix() < batch<T, A>::size)
{
// the plain lower-half load reads every lower lane, so they must all be active
assert(mask.countr_one() >= half_size && "plain lower-half load needs the lower half fully active");
const half_batch lo = half_batch::load(mem, Mode {});
constexpr auto mhi = ::xsimd::detail::upper_half<half_arch>(mask);
const half_batch hi = load_masked(mem + half_size, mhi, convert<T> {}, Mode {}, half_arch {});
return detail::merge_sse(lo.data, hi.data);
}
// exactly the upper 128-bit half: one plain load into the upper lanes
else XSIMD_IF_CONSTEXPR(mask.suffix() == half_size)
{
assert(mask.countl_one() >= half_size && mask.countr_zero() >= half_size && "upper half fully active, lower empty");
return detail::zero_extend<A>(half_batch::load(mem + half_size, Mode {}));
}
// upper 128-bit half
Expand Down Expand Up @@ -1103,13 +1118,17 @@ namespace xsimd
// exactly the lower 128-bit half: one plain store
XSIMD_IF_CONSTEXPR(mask.prefix() == half_size)
{
// a plain store writes every lower lane and no upper lane, so the mask
// must have the lower half fully active and the upper half empty
assert(mask.countr_one() >= half_size && mask.countl_zero() >= half_size && "lower half fully active, upper empty");
const half_batch lo = detail::lower_half(src);
lo.store(mem, Mode {});
}
// prefix crossing the 128-bit boundary: plain lower half + prefix-masked
// upper half. Never emits vmaskmov, which does not store-forward.
else XSIMD_IF_CONSTEXPR(mask.prefix() > half_size && mask.prefix() < batch<T, A>::size)
{
assert(mask.countr_one() >= half_size && "plain lower-half store needs the lower half fully active");
const half_batch lo = detail::lower_half(src);
lo.store(mem, Mode {});
constexpr auto mhi = ::xsimd::detail::upper_half<half_arch>(mask);
Expand All @@ -1119,6 +1138,7 @@ namespace xsimd
// exactly the upper 128-bit half: one plain store
else XSIMD_IF_CONSTEXPR(mask.suffix() == half_size)
{
assert(mask.countl_one() >= half_size && mask.countr_zero() >= half_size && "upper half fully active, lower empty");
const half_batch hi = detail::upper_half(src);
hi.store(mem + half_size, Mode {});
}
Expand Down Expand Up @@ -1172,6 +1192,45 @@ namespace xsimd
}
}

namespace detail
{
// Reinterpret a constant-mask 4/8-byte load/store as same-width float
// and run DstArch's kernel, which lowers prefix/suffix shapes to plain
// moves. Shared by the int/EVEX 128- and 256-bit archs.
template <class DstArch, class A, class T, bool... V, class Mode>
XSIMD_INLINE batch<T, A> plain_move_load(T const* mem, batch_bool_constant<T, A, V...>, convert<T>, Mode) noexcept
{
static_assert(sizeof(T) == 4 || sizeof(T) == 8, "plain-move delegation only supports 4/8-byte lanes");
using F = std::conditional_t<sizeof(T) == 4, float, double>;
// same-width float has the same lane count, so the V... mask pack and the
// memory reinterpret line up one-to-one; wrong here would load the wrong lanes
static_assert(batch<F, A>::size == batch<T, A>::size, "same-width float must preserve lane count");
static_assert(sizeof...(V) == batch<T, A>::size, "mask pack width must match the batch");
// the plain-move path emits aligned moves in aligned/stream mode, which fault
// on a misaligned pointer (the old vmaskmov tolerated it)
assert((std::is_same<Mode, unaligned_mode>::value || ::xsimd::is_aligned<A>(mem)) && "aligned/stream masked load needs an aligned pointer");
// qualify: an unqualified call resolves to detail::load_masked (a different
// helper) under MSVC's two-phase lookup; we want the kernel-level overload
return bitwise_cast<T>(batch<F, A>(::xsimd::kernel::load_masked(reinterpret_cast<F const*>(mem), batch_bool_constant<F, A, V...> {}, convert<F> {}, Mode {}, DstArch {})));
}

// Re-tag to DstArch so the vector-mask store kernel is used (the AVX
// store is gated off EVEX k-register archs).
template <class DstArch, class A, class T, bool... V, class Mode>
XSIMD_INLINE void plain_move_store(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, V...>, Mode) noexcept
{
static_assert(sizeof(T) == 4 || sizeof(T) == 8, "plain-move delegation only supports 4/8-byte lanes");
using F = std::conditional_t<sizeof(T) == 4, float, double>;
static_assert(batch<F, A>::size == batch<T, A>::size, "same-width float must preserve lane count");
static_assert(sizeof...(V) == batch<T, A>::size, "mask pack width must match the batch");
assert((std::is_same<Mode, unaligned_mode>::value || ::xsimd::is_aligned<A>(mem)) && "aligned/stream masked store needs an aligned pointer");
const auto fsrc = bitwise_cast<F>(src);
// qualify: an unqualified call resolves to detail::store_masked (a different
// helper) under MSVC's two-phase lookup; we want the kernel-level overload
::xsimd::kernel::store_masked(reinterpret_cast<F*>(mem), batch<F, DstArch>(fsrc.data), batch_bool_constant<F, DstArch, V...> {}, Mode {}, DstArch {});
}
}

// lt
template <class A>
XSIMD_INLINE batch_bool<float, A> lt(batch<float, A> const& self, batch<float, A> const& other, requires_arch<avx>) noexcept
Expand Down
32 changes: 3 additions & 29 deletions include/xsimd/arch/xsimd_avx2.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,19 +150,14 @@ namespace xsimd
_mm256_maskstore_epi64(reinterpret_cast<long long*>(mem), mask, src);
}
}

XSIMD_INLINE __m256i zero_extend(__m128i hi) noexcept
{
return _mm256_insertf128_si256(_mm256_setzero_si256(), hi, 1);
}
}

// no half-split shortcut for load; forward to runtime
// Constant masks: prefix/suffix shapes lower to plain moves.
template <class A, class T, bool... Values, class Mode>
XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), batch<T, A>>
load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode, requires_arch<avx2>) noexcept
{
return load_masked(mem, mask.as_batch_bool(), convert<T> {}, Mode {}, avx2 {});
return detail::plain_move_load<avx>(mem, mask, convert<T> {}, Mode {});
}

template <class A, class T, class Mode>
Expand All @@ -176,28 +171,7 @@ namespace xsimd
typename = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, Values...> mask, Mode, requires_arch<avx2>) noexcept
{
constexpr size_t lanes_per_half = batch<T, A>::size / 2;
using half_batch = ::xsimd::make_sized_batch_t<T, lanes_per_half>;
using half_arch = typename half_batch::arch_type;

// lower 128-bit half
XSIMD_IF_CONSTEXPR(mask.countl_zero() >= lanes_per_half)
{
constexpr auto mlo = ::xsimd::detail::lower_half<half_arch>(mask);
const half_batch lo = detail::lower_half(src);
store_masked<half_arch>(mem, lo, mlo, Mode {}, half_arch {});
}
// upper 128-bit half
else XSIMD_IF_CONSTEXPR(mask.countr_zero() >= lanes_per_half)
{
constexpr auto mhi = ::xsimd::detail::upper_half<half_arch>(mask);
const half_batch hi = detail::upper_half(src);
store_masked<half_arch>(mem + lanes_per_half, hi, mhi, Mode {}, half_arch {});
}
else
{
detail::maskstore(mem, mask.as_batch(), src);
}
detail::plain_move_store<avx>(mem, src, mask, Mode {});
}

template <class A, class T, class Mode>
Expand Down
6 changes: 2 additions & 4 deletions include/xsimd/arch/xsimd_avx2_128.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,7 @@ namespace xsimd
{
XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
{
using F = std::conditional_t<sizeof(T) == 4, float, double>;
return bitwise_cast<T>(batch<F, A>(load_masked(reinterpret_cast<F const*>(mem), batch_bool_constant<F, A, Values...> {}, convert<F> {}, Mode {}, sse2 {})));
return detail::plain_move_load<sse2>(mem, mask, convert<T> {}, Mode {});
}
else
{
Expand All @@ -164,8 +163,7 @@ namespace xsimd
{
XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
{
using F = std::conditional_t<sizeof(T) == 4, float, double>;
store_masked(reinterpret_cast<F*>(mem), bitwise_cast<F>(src), batch_bool_constant<F, A, Values...> {}, Mode {}, sse2 {});
detail::plain_move_store<sse2>(mem, src, mask, Mode {});
}
else
{
Expand Down
20 changes: 18 additions & 2 deletions include/xsimd/arch/xsimd_avx512vl_128.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -307,11 +307,20 @@ namespace xsimd
}
}

// Constant masks: prefix/suffix shapes lower to plain moves; interior
// masks keep the EVEX path.
template <class A, class T, bool... V, class Mode,
typename>
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, V...> mask, convert<T>, Mode, requires_arch<avx512vl_128>) noexcept
{
return detail::maskload128(mem, mask.mask(), Mode {});
XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
{
return detail::plain_move_load<sse2>(mem, mask, convert<T> {}, Mode {});
}
else
{
return detail::maskload128(mem, mask.mask(), Mode {});
}
}

template <class A, class T, class Mode,
Expand All @@ -325,7 +334,14 @@ namespace xsimd
typename>
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, V...> mask, Mode, requires_arch<avx512vl_128>) noexcept
{
detail::maskstore128(mem, src, mask.mask(), Mode {});
XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
{
detail::plain_move_store<sse2>(mem, src, mask, Mode {});
}
else
{
detail::maskstore128(mem, src, mask.mask(), Mode {});
}
}

template <class A, class T, class Mode,
Expand Down
29 changes: 27 additions & 2 deletions include/xsimd/arch/xsimd_avx512vl_256.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,11 +306,25 @@ namespace xsimd
}
}

// Constant masks: prefix/suffix shapes lower to plain moves; interior
// masks keep the EVEX path.
template <class A, class T, bool... V, class Mode,
typename = std::enable_if_t<std::is_arithmetic<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, V...> mask, convert<T>, Mode, requires_arch<avx512vl_256>) noexcept
{
return detail::maskload256(mem, mask.mask(), Mode {});
// all() reaches here only via the avx512f half-split cascade.
XSIMD_IF_CONSTEXPR(mask.all())
{
return batch<T, A>::load(mem, Mode {});
}
else XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
{
return detail::plain_move_load<avx>(mem, mask, convert<T> {}, Mode {});
}
else
{
return detail::maskload256(mem, mask.mask(), Mode {});
}
}

template <class A, class T, class Mode,
Expand All @@ -324,7 +338,18 @@ namespace xsimd
typename = std::enable_if_t<std::is_arithmetic<T>::value && (sizeof(T) == 4 || sizeof(T) == 8)>>
XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, V...> mask, Mode, requires_arch<avx512vl_256>) noexcept
{
detail::maskstore256(mem, src, mask.mask(), Mode {});
XSIMD_IF_CONSTEXPR(mask.all())
{
src.store(mem, Mode {});
}
else XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask))
{
detail::plain_move_store<avx2>(mem, src, mask, Mode {});
}
else
{
detail::maskstore256(mem, src, mask.mask(), Mode {});
}
}

template <class A, class T, class Mode,
Expand Down
Loading