Feat: Load/Store masked API

1. Adds new masked API compile time masks (store_masked and load_masked)
2. General use case optimization
3. New tests
4. x86 kernels
5. Adds new APIs to batch_bool_constant for convenience resembling #include<bit>
6. Tests the new APIs

Author: DiamonDinoia

docs/source/api/data_transfer.rst

@@ -10,7 +10,7 @@ Data transfer
 From memory:
 
 +---------------------------------------+----------------------------------------------------+
-| :cpp:func:`load`                      | load values from memory                            |
+| :cpp:func:`load`                      | load values from memory (optionally masked)        |
 +---------------------------------------+----------------------------------------------------+
 | :cpp:func:`load_aligned`              | load values from aligned memory                    |
 +---------------------------------------+----------------------------------------------------+
@@ -30,7 +30,7 @@ From a scalar:
 To memory:
 
 +---------------------------------------+----------------------------------------------------+
-| :cpp:func:`store`                     | store values to memory                             |
+| :cpp:func:`store`                     | store values to memory (optionally masked)         |
 +---------------------------------------+----------------------------------------------------+
 | :cpp:func:`store_aligned`             | store values to aligned memory                     |
 +---------------------------------------+----------------------------------------------------+

include/xsimd/arch/common/xsimd_common_arithmetic.hpp

@@ -16,6 +16,7 @@
 #include <limits>
 #include <type_traits>
 
+#include "../../types/xsimd_batch_constant.hpp"
 #include "./xsimd_common_details.hpp"
 
 namespace xsimd

include/xsimd/arch/common/xsimd_common_memory.hpp

@@ -13,6 +13,7 @@
 #define XSIMD_COMMON_MEMORY_HPP
 
 #include <algorithm>
+#include <array>
 #include <complex>
 #include <stdexcept>
 
@@ -341,6 +342,102 @@ namespace xsimd
             return detail::load_unaligned<A>(mem, cvt, common {}, detail::conversion_type<A, T_in, T_out> {});
         }
 
+        template <class A, class T>
+        XSIMD_INLINE batch<T, A> load(T const* mem, aligned_mode, requires_arch<A>) noexcept
+        {
+            return load_aligned<A>(mem, convert<T> {}, A {});
+        }
+
+        template <class A, class T>
+        XSIMD_INLINE batch<T, A> load(T const* mem, unaligned_mode, requires_arch<A>) noexcept
+        {
+            return load_unaligned<A>(mem, convert<T> {}, A {});
+        }
+
+        template <class A, class T_in, class T_out, bool... Values, class alignment>
+        XSIMD_INLINE batch<T_out, A>
+        load_masked(T_in const* mem, batch_bool_constant<T_out, A, Values...>, convert<T_out>, alignment, requires_arch<common>) noexcept
+        {
+            constexpr std::size_t size = batch<T_out, A>::size;
+            alignas(A::alignment()) std::array<T_out, size> buffer {};
+            constexpr bool mask[size] = { Values... };
+
+            for (std::size_t i = 0; i < size; ++i)
+                buffer[i] = mask[i] ? static_cast<T_out>(mem[i]) : T_out(0);
+
+            return batch<T_out, A>::load(buffer.data(), aligned_mode {});
+        }
+
+        template <class A, class T_in, class T_out, bool... Values, class alignment>
+        XSIMD_INLINE void
+        store_masked(T_out* mem, batch<T_in, A> const& src, batch_bool_constant<T_in, A, Values...>, alignment, requires_arch<common>) noexcept
+        {
+            constexpr std::size_t size = batch<T_in, A>::size;
+            constexpr bool mask[size] = { Values... };
+
+            for (std::size_t i = 0; i < size; ++i)
+                if (mask[i])
+                {
+                    mem[i] = static_cast<T_out>(src.get(i));
+                }
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE batch<int32_t, A> load_masked(int32_t const* mem, batch_bool_constant<int32_t, A, Values...>, convert<int32_t>, Mode, requires_arch<A>) noexcept
+        {
+            const auto f = load_masked<A>(reinterpret_cast<const float*>(mem), batch_bool_constant<float, A, Values...> {}, convert<float> {}, Mode {}, A {});
+            return bitwise_cast<int32_t>(f);
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE batch<uint32_t, A> load_masked(uint32_t const* mem, batch_bool_constant<uint32_t, A, Values...>, convert<uint32_t>, Mode, requires_arch<A>) noexcept
+        {
+            const auto f = load_masked<A>(reinterpret_cast<const float*>(mem), batch_bool_constant<float, A, Values...> {}, convert<float> {}, Mode {}, A {});
+            return bitwise_cast<uint32_t>(f);
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE typename std::enable_if<types::has_simd_register<double, A>::value, batch<int64_t, A>>::type
+        load_masked(int64_t const* mem, batch_bool_constant<int64_t, A, Values...>, convert<int64_t>, Mode, requires_arch<A>) noexcept
+        {
+            const auto d = load_masked<A>(reinterpret_cast<const double*>(mem), batch_bool_constant<double, A, Values...> {}, convert<double> {}, Mode {}, A {});
+            return bitwise_cast<int64_t>(d);
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE typename std::enable_if<types::has_simd_register<double, A>::value, batch<uint64_t, A>>::type
+        load_masked(uint64_t const* mem, batch_bool_constant<uint64_t, A, Values...>, convert<uint64_t>, Mode, requires_arch<A>) noexcept
+        {
+            const auto d = load_masked<A>(reinterpret_cast<const double*>(mem), batch_bool_constant<double, A, Values...> {}, convert<double> {}, Mode {}, A {});
+            return bitwise_cast<uint64_t>(d);
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE void store_masked(int32_t* mem, batch<int32_t, A> const& src, batch_bool_constant<int32_t, A, Values...>, Mode, requires_arch<A>) noexcept
+        {
+            store_masked<A>(reinterpret_cast<float*>(mem), bitwise_cast<float>(src), batch_bool_constant<float, A, Values...> {}, Mode {}, A {});
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE void store_masked(uint32_t* mem, batch<uint32_t, A> const& src, batch_bool_constant<uint32_t, A, Values...>, Mode, requires_arch<A>) noexcept
+        {
+            store_masked<A>(reinterpret_cast<float*>(mem), bitwise_cast<float>(src), batch_bool_constant<float, A, Values...> {}, Mode {}, A {});
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE typename std::enable_if<types::has_simd_register<double, A>::value, void>::type
+        store_masked(int64_t* mem, batch<int64_t, A> const& src, batch_bool_constant<int64_t, A, Values...>, Mode, requires_arch<A>) noexcept
+        {
+            store_masked<A>(reinterpret_cast<double*>(mem), bitwise_cast<double>(src), batch_bool_constant<double, A, Values...> {}, Mode {}, A {});
+        }
+
+        template <class A, bool... Values, class Mode>
+        XSIMD_INLINE typename std::enable_if<types::has_simd_register<double, A>::value, void>::type
+        store_masked(uint64_t* mem, batch<uint64_t, A> const& src, batch_bool_constant<uint64_t, A, Values...>, Mode, requires_arch<A>) noexcept
+        {
+            store_masked<A>(reinterpret_cast<double*>(mem), bitwise_cast<double>(src), batch_bool_constant<double, A, Values...> {}, Mode {}, A {});
+        }
+
         // rotate_right
         template <size_t N, class A, class T>
         XSIMD_INLINE batch<T, A> rotate_right(batch<T, A> const& self, requires_arch<common>) noexcept

include/xsimd/arch/xsimd_avx.hpp

@@ -18,6 +18,7 @@
 #include <type_traits>
 
 #include "../types/xsimd_avx_register.hpp"
+#include "../types/xsimd_batch_constant.hpp"
 
 namespace xsimd
 {
@@ -871,6 +872,132 @@ namespace xsimd
             return _mm256_loadu_pd(mem);
         }
 
+        // AVX helpers to avoid type-based branching in the generic load_masked
+        namespace detail
+        {
+            template <class A>
+            XSIMD_INLINE batch<float, A> maskload(float const* mem, batch<as_integer_t<float>, A> const& mask) noexcept
+            {
+                return _mm256_maskload_ps(mem, mask);
+            }
+
+            template <class A>
+            XSIMD_INLINE batch<double, A> maskload(double const* mem, batch<as_integer_t<double>, A> const& mask) noexcept
+            {
+                return _mm256_maskload_pd(mem, mask);
+            }
+
+            template <class A>
+            XSIMD_INLINE batch<float, A> zero_extend(batch<float, A> const& hi) noexcept
+            {
+                return _mm256_insertf128_ps(_mm256_setzero_ps(), hi, 1);
+            }
+
+            template <class A>
+            XSIMD_INLINE batch<double, A> zero_extend(batch<double, A> const& hi) noexcept
+            {
+                return _mm256_insertf128_pd(_mm256_setzero_pd(), hi, 1);
+            }
+
+            // allow inserting a 128-bit SSE batch into the upper half of an AVX batch
+            template <class A, class SrcA>
+            XSIMD_INLINE batch<float, A> zero_extend(batch<float, SrcA> const& hi) noexcept
+            {
+                return _mm256_insertf128_ps(_mm256_setzero_ps(), hi, 1);
+            }
+
+            template <class A, class SrcA>
+            XSIMD_INLINE batch<double, A> zero_extend(batch<double, SrcA> const& hi) noexcept
+            {
+                return _mm256_insertf128_pd(_mm256_setzero_pd(), hi, 1);
+            }
+        }
+
+        // load_masked (single overload for float/double)
+        template <class A, class T, bool... Values, class Mode, class = typename std::enable_if<std::is_floating_point<T>::value>::type>
+        XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode, requires_arch<avx>) noexcept
+        {
+            using int_t = as_integer_t<T>;
+            constexpr size_t half_size = batch<T, A>::size / 2;
+
+            XSIMD_IF_CONSTEXPR(mask.none())
+            {
+                return batch<T, A>(T { 0 });
+            }
+            else XSIMD_IF_CONSTEXPR(mask.all())
+            {
+                return load<A>(mem, Mode {});
+            }
+            // confined to lower 128-bit half → forward to SSE2
+            else XSIMD_IF_CONSTEXPR(mask.countl_zero() >= half_size)
+            {
+                constexpr auto mlo = ::xsimd::detail::lower_half<sse4_2>(batch_bool_constant<int_t, A, Values...> {});
+                const auto lo = load_masked(reinterpret_cast<int_t const*>(mem), mlo, convert<int_t> {}, Mode {}, sse4_2 {});
+                return bitwise_cast<T>(batch<int_t, A>(_mm256_zextsi128_si256(lo)));
+            }
+            // confined to upper 128-bit half → forward to SSE2
+            else XSIMD_IF_CONSTEXPR(mask.countr_zero() >= half_size)
+            {
+                constexpr auto mhi = ::xsimd::detail::upper_half<sse4_2>(mask);
+                const auto hi = load_masked(mem + half_size, mhi, convert<T> {}, Mode {}, sse4_2 {});
+                return detail::zero_extend<A>(hi);
+            }
+            else
+            {
+                // crossing 128-bit boundary → use 256-bit masked load
+                return detail::maskload<A>(mem, mask.as_batch());
+            }
+        }
+
+        // store_masked
+        namespace detail
+        {
+            template <class A>
+            XSIMD_INLINE void maskstore(float* mem, batch_bool<float, A> const& mask, batch<float, A> const& src) noexcept
+            {
+                _mm256_maskstore_ps(mem, mask, src);
+            }
+
+            template <class A>
+            XSIMD_INLINE void maskstore(double* mem, batch_bool<double, A> const& mask, batch<double, A> const& src) noexcept
+            {
+                _mm256_maskstore_pd(mem, mask, src);
+            }
+        }
+
+        template <class A, class T, bool... Values, class Mode>
+        XSIMD_INLINE void store_masked(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, Values...> mask, Mode, requires_arch<avx>) noexcept
+        {
+            constexpr size_t half_size = batch<T, A>::size / 2;
+
+            XSIMD_IF_CONSTEXPR(mask.none())
+            {
+                return;
+            }
+            else XSIMD_IF_CONSTEXPR(mask.all())
+            {
+                src.store(mem, Mode {});
+            }
+            // confined to lower 128-bit half → forward to SSE2
+            else XSIMD_IF_CONSTEXPR(mask.countl_zero() >= half_size)
+            {
+                constexpr auto mlo = ::xsimd::detail::lower_half<sse4_2>(mask);
+                const auto lo = detail::lower_half(src);
+                store_masked<sse4_2>(mem, lo, mlo, Mode {}, sse4_2 {});
+            }
+            // confined to upper 128-bit half → forward to SSE2
+            else XSIMD_IF_CONSTEXPR(mask.countr_zero() >= half_size)
+            {
+                constexpr auto mhi = ::xsimd::detail::upper_half<sse4_2>(mask);
+                const auto hi = detail::upper_half(src);
+                store_masked<sse4_2>(mem + half_size, hi, mhi, Mode {}, sse4_2 {});
+            }
+            else
+            {
+                detail::maskstore(mem, mask.as_batch(), src);
+            }
+        }
+
         // 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