From 96866bc5828a2244bf5a3a57465b33c6a8605643 Mon Sep 17 00:00:00 2001 From: Marco Barbone Date: Thu, 9 Jul 2026 12:18:33 -0400 Subject: [PATCH] fix: correct constant-mask load/store kernels and add shape predicates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sse2 and neon constant-mask load_masked kernels were unreachable: they lacked the convert overload parameter the dispatcher passes, so calls fell through to the runtime-mask path. Add the parameter and fix the mask conditions, which used ambiguous count-based tests (e.g. countr_one()==2 also matches 0b1011, silently dropping a lane). Add batch_bool_constant::is_prefix/is_suffix shape predicates plus prefix()/ suffix() returning the set-run length (or size+1 when the mask is not that shape, keeping == exact). The sse2 kernels now dispatch on these instead of raw bit patterns, so each condition reads as the lane shape it handles. avx_128 delegates prefix/upper-half masks to sse2: those lower to plain moves (vmovss/movlps/…) which store-forward, whereas vmaskmov never does on Intel and its stores are microcoded on AMD. Tests cover the shape predicates including off-by-one boundaries (size-1 vs the size+1 sentinel) and near-full mask patterns through the load/store kernels. --- include/xsimd/arch/xsimd_avx_128.hpp | 21 +++++- include/xsimd/arch/xsimd_neon.hpp | 22 +++--- include/xsimd/arch/xsimd_sse2.hpp | 78 +++++++++++--------- include/xsimd/types/xsimd_batch_constant.hpp | 28 +++++++ test/test_batch_constant.cpp | 74 +++++++++++++++++++ test/test_load_store.cpp | 21 ++++++ 6 files changed, 195 insertions(+), 49 deletions(-) diff --git a/include/xsimd/arch/xsimd_avx_128.hpp b/include/xsimd/arch/xsimd_avx_128.hpp index e3d8b326b..498db1066 100644 --- a/include/xsimd/arch/xsimd_avx_128.hpp +++ b/include/xsimd/arch/xsimd_avx_128.hpp @@ -103,11 +103,19 @@ namespace xsimd return _mm_cmp_pd(self, other, _CMP_NEQ_UQ); } - // constant masks gain nothing on a single register; forward to runtime + // Masks that lower to plain moves go to sse2; the rest gain nothing on a + // single register, so take the runtime path. template ::value>> XSIMD_INLINE batch load_masked(T const* mem, batch_bool_constant mask, convert, Mode, requires_arch) noexcept { - return load_masked(mem, mask.as_batch_bool(), convert {}, Mode {}, avx_128 {}); + XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask)) + { + return load_masked(mem, mask, convert {}, Mode {}, sse2 {}); + } + else + { + return load_masked(mem, mask.as_batch_bool(), convert {}, Mode {}, avx_128 {}); + } } // Runtime-mask load (float/double). @@ -142,7 +150,14 @@ namespace xsimd template ::value>> XSIMD_INLINE void store_masked(T* mem, batch const& src, batch_bool_constant mask, Mode, requires_arch) noexcept { - store_masked(mem, src, mask.as_batch_bool(), Mode {}, avx_128 {}); + XSIMD_IF_CONSTEXPR(detail::lowers_to_plain_moves(mask)) + { + store_masked(mem, src, mask, Mode {}, sse2 {}); + } + else + { + store_masked(mem, src, mask.as_batch_bool(), Mode {}, avx_128 {}); + } } // Runtime-mask store (float/double). diff --git a/include/xsimd/arch/xsimd_neon.hpp b/include/xsimd/arch/xsimd_neon.hpp index 48be9b88e..20bb67746 100644 --- a/include/xsimd/arch/xsimd_neon.hpp +++ b/include/xsimd/arch/xsimd_neon.hpp @@ -539,8 +539,8 @@ namespace xsimd template <> struct load_masked<> { - template - static XSIMD_INLINE batch apply(T const* /* mem */, batch acc, std::integral_constant) noexcept + template + static XSIMD_INLINE batch apply(T const* /* mem */, batch acc) noexcept { return acc; } @@ -549,23 +549,21 @@ namespace xsimd struct load_masked { template - static XSIMD_INLINE batch apply(T const* mem, batch acc, std::true_type) noexcept - { - return load_masked::template apply(mem, insert(acc, mem[I], index {}), std::integral_constant {}); - } - template - static XSIMD_INLINE batch apply(T const* mem, batch acc, std::false_type) noexcept + static XSIMD_INLINE batch apply(T const* mem, batch acc) noexcept { - return load_masked::template apply(mem, acc, std::integral_constant {}); + XSIMD_IF_CONSTEXPR(Value) + { + acc = insert(acc, mem[I], index {}); + } + return load_masked::template apply(mem, acc); } }; } template - XSIMD_INLINE batch load_masked(T const* mem, batch_bool_constant /* mask */, Mode, requires_arch) noexcept + XSIMD_INLINE batch load_masked(T const* mem, batch_bool_constant /* mask */, convert, Mode, requires_arch) noexcept { - // Call insert whenever Values... are true - return detail::load_masked::template apply<0>(mem, broadcast(T(0), A {}), std::integral_constant {}); + return detail::load_masked::template apply<0>(mem, batch(T(0))); } /********* diff --git a/include/xsimd/arch/xsimd_sse2.hpp b/include/xsimd/arch/xsimd_sse2.hpp index 5748ed831..920814c74 100644 --- a/include/xsimd/arch/xsimd_sse2.hpp +++ b/include/xsimd/arch/xsimd_sse2.hpp @@ -1100,32 +1100,42 @@ namespace xsimd return { load_unaligned(mem, batch_bool {}, r).data }; } + namespace detail + { + // Plain moves store-forward; vmaskmov never does on Intel and its + // stores are microcoded on AMD. So wider archs delegate the masks + // that lower to plain moves here rather than taking their runtime path. + template + constexpr bool lowers_to_plain_moves(batch_bool_constant mask) noexcept + { + return (mask.is_prefix() && mask.any() && !mask.all()) + || mask.suffix() == mask.size / 2; + } + } + // load_masked template ::value>> - XSIMD_INLINE batch load_masked(T const* mem, batch_bool_constant mask, Mode, requires_arch) noexcept + XSIMD_INLINE batch load_masked(T const* mem, batch_bool_constant mask, convert, Mode, requires_arch) noexcept { - XSIMD_IF_CONSTEXPR(mask.mask() == 0x1) + XSIMD_IF_CONSTEXPR(sizeof(T) == 2 && mask.prefix() == 1) { - XSIMD_IF_CONSTEXPR(sizeof(T) == 2) - { - return mm_loadu_si16(mem); - } - XSIMD_IF_CONSTEXPR(sizeof(T) == 4) - { - return mm_loadu_si32(mem); - } - XSIMD_IF_CONSTEXPR(sizeof(T) == 8) - { - return mm_loadu_si64(mem); - } + return _mm_loadu_si16(mem); + } + else XSIMD_IF_CONSTEXPR(sizeof(T) == 4 && mask.prefix() == 1) + { + return _mm_loadu_si32(mem); + } + else XSIMD_IF_CONSTEXPR(sizeof(T) == 8 && mask.prefix() == 1) + { + return _mm_loadu_si64(mem); } - else XSIMD_IF_CONSTEXPR(sizeof(T) == 2 && mask.mask() == 0x3) + else XSIMD_IF_CONSTEXPR(sizeof(T) == 2 && mask.prefix() == 2) { - return mm_loadu_si32(mem); + return _mm_loadu_si32(mem); } - else XSIMD_IF_CONSTEXPR(sizeof(T) == 4 && mask.mask() == 0x3) + else XSIMD_IF_CONSTEXPR(sizeof(T) == 4 && mask.prefix() == 2) { - return mm_loadu_si64(mem); + return _mm_loadu_si64(mem); } else { @@ -1133,21 +1143,21 @@ namespace xsimd } } template - XSIMD_INLINE batch load_masked(float const* mem, batch_bool_constant mask, Mode, requires_arch) noexcept + XSIMD_INLINE batch load_masked(float const* mem, batch_bool_constant mask, convert, Mode, requires_arch) noexcept { - XSIMD_IF_CONSTEXPR(mask.mask() == 0x1) + XSIMD_IF_CONSTEXPR(mask.prefix() == 1) { return _mm_load_ss(mem); } - else XSIMD_IF_CONSTEXPR(mask.countr_one() == 2) + else XSIMD_IF_CONSTEXPR(mask.prefix() == 2) { return _mm_loadl_pi(_mm_setzero_ps(), reinterpret_cast<__m64 const*>(mem)); } - else XSIMD_IF_CONSTEXPR(mask.countl_one() == 2) + else XSIMD_IF_CONSTEXPR(mask.suffix() == 2) { return _mm_loadh_pi(_mm_setzero_ps(), reinterpret_cast<__m64 const*>(mem + 2)); } - else XSIMD_IF_CONSTEXPR(mask.countr_one() == 3) + else XSIMD_IF_CONSTEXPR(mask.prefix() == 3) { __m128 const lo2 = _mm_castsi128_ps(_mm_loadl_epi64(reinterpret_cast<__m128i const*>(mem))); return _mm_shuffle_ps(lo2, _mm_load_ss(mem + 2), _MM_SHUFFLE(3, 0, 1, 0)); @@ -1158,13 +1168,13 @@ namespace xsimd } } template - XSIMD_INLINE batch load_masked(double const* mem, batch_bool_constant mask, Mode, requires_arch) noexcept + XSIMD_INLINE batch load_masked(double const* mem, batch_bool_constant mask, convert, Mode, requires_arch) noexcept { - XSIMD_IF_CONSTEXPR(mask.countr_one() == 1) + XSIMD_IF_CONSTEXPR(mask.prefix() == 1) { return _mm_load_sd(mem); } - else XSIMD_IF_CONSTEXPR(mask.countl_one() == 1) + else XSIMD_IF_CONSTEXPR(mask.suffix() == 1) { return _mm_loadh_pd(_mm_setzero_pd(), mem + 1); } @@ -1178,19 +1188,19 @@ namespace xsimd template XSIMD_INLINE void store_masked(float* mem, batch const& src, batch_bool_constant mask, Mode, requires_arch) noexcept { - XSIMD_IF_CONSTEXPR(mask.mask() == 0x1) + XSIMD_IF_CONSTEXPR(mask.prefix() == 1) { _mm_store_ss(mem, src); } - else XSIMD_IF_CONSTEXPR(mask.countr_one() == 2) + else XSIMD_IF_CONSTEXPR(mask.prefix() == 2) { _mm_storel_pi(reinterpret_cast<__m64*>(mem), src); } - else XSIMD_IF_CONSTEXPR(mask.countl_one() == 2) + else XSIMD_IF_CONSTEXPR(mask.suffix() == 2) { _mm_storeh_pi(reinterpret_cast<__m64*>(mem + 2), src); } - else XSIMD_IF_CONSTEXPR(mask.countr_one() == 3) + else XSIMD_IF_CONSTEXPR(mask.prefix() == 3) { _mm_storel_pi(reinterpret_cast<__m64*>(mem), src); _mm_store_ss(mem + 2, _mm_movehl_ps(src, src)); @@ -1204,11 +1214,11 @@ namespace xsimd template XSIMD_INLINE void store_masked(double* mem, batch const& src, batch_bool_constant mask, Mode, requires_arch) noexcept { - XSIMD_IF_CONSTEXPR(mask.countr_one() == 1) + XSIMD_IF_CONSTEXPR(mask.prefix() == 1) { _mm_store_sd(mem, src); } - else XSIMD_IF_CONSTEXPR(mask.countl_one() == 1) + else XSIMD_IF_CONSTEXPR(mask.suffix() == 1) { _mm_storeh_pd(mem + 1, src); } @@ -2331,11 +2341,11 @@ namespace xsimd aligned_mode, requires_arch) noexcept { - XSIMD_IF_CONSTEXPR(mask.countr_one() == 2) + XSIMD_IF_CONSTEXPR(mask.prefix() == 2) { _mm_storel_pi(reinterpret_cast<__m64*>(mem), src); } - else XSIMD_IF_CONSTEXPR(mask.countl_one() == 2) + else XSIMD_IF_CONSTEXPR(mask.suffix() == 2) { _mm_storeh_pi(reinterpret_cast<__m64*>(mem + 2), src); } diff --git a/include/xsimd/types/xsimd_batch_constant.hpp b/include/xsimd/types/xsimd_batch_constant.hpp index edecee3c3..54cdade1a 100644 --- a/include/xsimd/types/xsimd_batch_constant.hpp +++ b/include/xsimd/types/xsimd_batch_constant.hpp @@ -99,6 +99,34 @@ namespace xsimd return countl_one_impl(truncated_mask(), size); } + // true when the set lanes form one contiguous run starting at lane 0 + // (the empty and full masks are prefixes) + static constexpr bool is_prefix() noexcept + { + return (truncated_mask() & (truncated_mask() + 1u)) == 0u; + } + + // true when the set lanes form one contiguous run ending at the last + // lane (the empty and full masks are suffixes) + static constexpr bool is_suffix() noexcept + { + return ((truncated_mask() ^ low_mask(size)) & ((truncated_mask() ^ low_mask(size)) + 1u)) == 0u; + } + + // length of the set run when the mask is a pure prefix (first k lanes + // set, rest clear), else size + 1 so `prefix() == k` stays exact + static constexpr std::size_t prefix() noexcept + { + return is_prefix() ? countr_one() : size + 1; + } + + // length of the set run when the mask is a pure suffix (last k lanes + // set, rest clear), else size + 1 so `suffix() == k` stays exact + static constexpr std::size_t suffix() noexcept + { + return is_suffix() ? countl_one() : size + 1; + } + private: static constexpr int mask_helper(int acc) noexcept { return acc; } diff --git a/test/test_batch_constant.cpp b/test/test_batch_constant.cpp index 9877cd17d..e92cf18e8 100644 --- a/test/test_batch_constant.cpp +++ b/test/test_batch_constant.cpp @@ -389,6 +389,76 @@ struct constant_bool_batch_test constexpr auto inv_x = ~x; static_assert(std::is_same::value, "~x == y"); } + + struct first_half + { + static constexpr bool get(size_t index, size_t size) { return index < size / 2; } + }; + struct second_half + { + static constexpr bool get(size_t index, size_t size) { return index >= size / 2; } + }; + struct ends + { + static constexpr bool get(size_t index, size_t size) { return index == 0 || index + 1 == size; } + }; + struct all_but_last + { + static constexpr bool get(size_t index, size_t size) { return index + 1 < size; } + }; + struct all_but_first + { + static constexpr bool get(size_t index, size_t) { return index != 0; } + }; + struct first_only + { + static constexpr bool get(size_t index, size_t) { return index == 0; } + }; + struct last_only + { + static constexpr bool get(size_t index, size_t size) { return index + 1 == size; } + }; + + void test_shape() const + { + constexpr auto all_true = xsimd::make_batch_bool_constant, arch_type>(); + constexpr auto all_false = xsimd::make_batch_bool_constant, arch_type>(); + constexpr auto lo = xsimd::make_batch_bool_constant(); + constexpr auto hi = xsimd::make_batch_bool_constant(); + constexpr auto edges = xsimd::make_batch_bool_constant(); + constexpr auto size = decltype(all_true)::size; + + static_assert(all_true.is_prefix() && all_true.is_suffix(), "full mask is prefix and suffix"); + static_assert(all_false.is_prefix() && all_false.is_suffix(), "empty mask is prefix and suffix"); + static_assert(lo.is_prefix() && !lo.is_suffix(), "first half is a prefix only"); + static_assert(hi.is_suffix() && !hi.is_prefix(), "second half is a suffix only"); + // {first, last} lanes: contiguous only when that covers the whole batch + static_assert(edges.is_prefix() == (size <= 2), "non-contiguous mask is not a prefix"); + static_assert(edges.is_suffix() == (size <= 2), "non-contiguous mask is not a suffix"); + + // prefix()/suffix() return the set-run length, or size+1 when not that shape + static_assert(lo.prefix() == size / 2 && lo.suffix() == size + 1, "lo is the size/2 prefix, no suffix"); + static_assert(hi.suffix() == size / 2 && hi.prefix() == size + 1, "hi is the size/2 suffix, no prefix"); + static_assert(all_true.prefix() == size && all_true.suffix() == size, "full mask runs the whole width"); + static_assert(all_false.prefix() == 0 && all_false.suffix() == 0, "empty mask has zero-length runs"); + static_assert(edges.prefix() == (size <= 2 ? size : size + 1), "non-contiguous mask has no prefix length"); + + // off-by-one boundary: the length adjacent to the sentinel (size-1) must + // count exactly, and the sentinel must be exactly size+1 (never size, size+2) + constexpr auto pre1 = xsimd::make_batch_bool_constant(); + constexpr auto suf1 = xsimd::make_batch_bool_constant(); + constexpr auto f1 = xsimd::make_batch_bool_constant(); + constexpr auto l1 = xsimd::make_batch_bool_constant(); + static_assert(size < 2 || pre1.prefix() == size - 1, "size-1 prefix counts exactly"); + static_assert(size < 2 || pre1.suffix() == size + 1, "a size-1 prefix is not a suffix"); + static_assert(size < 2 || suf1.suffix() == size - 1, "size-1 suffix counts exactly"); + static_assert(size < 2 || suf1.prefix() == size + 1, "a size-1 suffix is not a prefix"); + static_assert(size < 2 || (pre1.suffix() != size && pre1.suffix() != size + 2), "sentinel is exactly size+1"); + static_assert(f1.prefix() == 1, "single first lane is the 1-prefix"); + static_assert(size < 2 || f1.suffix() == size + 1, "single first lane is not a suffix"); + static_assert(l1.suffix() == 1, "single last lane is the 1-suffix"); + static_assert(size < 2 || l1.prefix() == size + 1, "single last lane is not a prefix"); + } }; TEST_CASE_TEMPLATE("[constant bool batch]", B, BATCH_INT_TYPES) @@ -410,5 +480,9 @@ TEST_CASE_TEMPLATE("[constant bool batch]", B, BATCH_INT_TYPES) { Test.test_ops(); } + SUBCASE("shape") + { + Test.test_shape(); + } } #endif diff --git a/test/test_load_store.cpp b/test/test_load_store.cpp index 6a1b18c02..f89f89c0f 100644 --- a/test/test_load_store.cpp +++ b/test/test_load_store.cpp @@ -125,6 +125,21 @@ struct load_store_test static constexpr bool get(std::size_t, std::size_t) noexcept { return true; } }; + struct mask_all_but_last + { + static constexpr bool get(std::size_t index, std::size_t size) noexcept { return index + 1 < size; } + }; + + struct mask_all_but_first + { + static constexpr bool get(std::size_t index, std::size_t) noexcept { return index != 0; } + }; + + struct mask_last + { + static constexpr bool get(std::size_t index, std::size_t size) noexcept { return index + 1 == size; } + }; + template static batch_bool_type make_runtime_mask() noexcept { @@ -436,6 +451,9 @@ struct load_store_test run_load_mask_pattern(v, name, b, expected, p + " last half"); run_load_mask_pattern(v, name, b, expected, p + " even elements"); run_load_mask_pattern(v, name, b, expected, p + " odd elements"); + run_load_mask_pattern(v, name, b, expected, p + " last element"); + run_load_mask_pattern(v, name, b, expected, p + " all but last"); + run_load_mask_pattern(v, name, b, expected, p + " all but first"); run_load_mask_pattern(v, name, b, expected, p + " pseudo random"); run_load_mask_pattern(v, name, b, expected, p + " all elements"); } @@ -499,6 +517,9 @@ struct load_store_test run_store_mask_pattern(v, name, b, res, expected_masked, p + " last half"); run_store_mask_pattern(v, name, b, res, expected_masked, p + " even elements"); run_store_mask_pattern(v, name, b, res, expected_masked, p + " odd elements"); + run_store_mask_pattern(v, name, b, res, expected_masked, p + " last element"); + run_store_mask_pattern(v, name, b, res, expected_masked, p + " all but last"); + run_store_mask_pattern(v, name, b, res, expected_masked, p + " all but first"); run_store_mask_pattern(v, name, b, res, expected_masked, p + " pseudo random"); run_store_mask_pattern(v, name, b, res, expected_masked, p + " all elements"); }