feat: add runtime batch_bool mask overloads for load_masked/store_masked

.github/workflows/doxygen.yml

@@ -9,6 +9,8 @@ jobs:
     steps:
     - uses: actions/checkout@v6
     - name: Install dependencies
-      run: sudo apt install doxygen python3-breathe python3-sphinx-rtd-theme
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y doxygen python3-breathe python3-sphinx-rtd-theme
     - name: Render
       run: make -C docs

docs/source/api/data_transfer.rst

@@ -12,7 +12,7 @@ Data Transfers
 From memory:
 
 +---------------------------------------+----------------------------------------------------+
-| :cpp:func:`load`                      | load values from memory (optionally masked)        |
+| :cpp:func:`load`                      | load values from memory (optionally masked) [#m]_  |
 +---------------------------------------+----------------------------------------------------+
 | :cpp:func:`load_aligned`              | load values from aligned memory                    |
 +---------------------------------------+----------------------------------------------------+
@@ -32,7 +32,7 @@ From a scalar:
 To memory:
 
 +---------------------------------------+----------------------------------------------------+
-| :cpp:func:`store`                     | store values to memory (optionally masked)         |
+| :cpp:func:`store`                     | store values to memory (optionally masked) [#m]_   |
 +---------------------------------------+----------------------------------------------------+
 | :cpp:func:`store_aligned`             | store values to aligned memory                     |
 +---------------------------------------+----------------------------------------------------+
@@ -84,3 +84,11 @@ The following empty types are used for tag dispatching:
 
 .. doxygenstruct:: xsimd::unaligned_mode
    :project: xsimd
+
+.. rubric:: Footnotes
+
+.. [#m] Masked ``load`` / ``store`` come in two flavours. The
+   :cpp:class:`batch_bool_constant` overload encodes the mask in the type and
+   is resolved at compile time. The runtime :cpp:class:`batch_bool` overload
+   accepts a mask computed at runtime. For performance reasons, prefer the
+   compile-time mask whenever possible.

include/xsimd/arch/common/xsimd_common_memory.hpp

@@ -437,13 +437,40 @@ namespace xsimd
             return detail::load_masked_common(mem, mask, cvt, mode, detail::masked_memory_uses_fp_bitcast<A, T_in, T_out> {});
         }
 
+        template <class A, class T, class Mode>
+        XSIMD_INLINE batch<T, A>
+        load_masked(T const* mem, batch_bool<T, A> mask, convert<T>, Mode, requires_arch<common>) noexcept
+        {
+            // Scalar fallback: only active lanes are touched. Arches with
+            // hardware predicated loads should override this.
+            constexpr std::size_t size = batch<T, A>::size;
+            alignas(A::alignment()) std::array<T, size> buffer;
+            for (std::size_t i = 0; i < size; ++i)
+                buffer[i] = mask.get(i) ? mem[i] : T(0);
+            return batch<T, A>::load_aligned(buffer.data());
+        }
+
         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...> mask, alignment mode, requires_arch<common>) noexcept
         {
             detail::store_masked_common(mem, src, mask, mode, detail::masked_memory_uses_fp_bitcast<A, T_in, T_out> {});
         }
 
+        template <class A, class T, class Mode>
+        XSIMD_INLINE void
+        store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<common>) noexcept
+        {
+            // Scalar fallback: only active lanes are touched. Arches with
+            // hardware predicated stores should override this.
+            constexpr std::size_t size = batch<T, A>::size;
+            alignas(A::alignment()) std::array<T, size> src_buf;
+            src.store_aligned(src_buf.data());
+            for (std::size_t i = 0; i < size; ++i)
+                if (mask.get(i))
+                    mem[i] = src_buf[i];
+        }
+
         template <class A, class T_in, class T_out>
         XSIMD_INLINE batch<T_out, A> load_stream(T_in const* mem, convert<T_out> cvt, requires_arch<common>) noexcept
         {

include/xsimd/arch/xsimd_avx.hpp

@@ -987,6 +987,21 @@ namespace xsimd
             }
         }
 
+        // Runtime-mask load (float/double).
+        template <class A, class Mode>
+        XSIMD_INLINE batch<float, A>
+        load_masked(float const* mem, batch_bool<float, A> mask, convert<float>, Mode, requires_arch<avx>) noexcept
+        {
+            return _mm256_maskload_ps(mem, _mm256_castps_si256(mask));
+        }
+
+        template <class A, class Mode>
+        XSIMD_INLINE batch<double, A>
+        load_masked(double const* mem, batch_bool<double, A> mask, convert<double>, Mode, requires_arch<avx>) noexcept
+        {
+            return _mm256_maskload_pd(mem, _mm256_castpd_si256(mask));
+        }
+
         // load_masked (single overload for float/double)
         template <class A, class T, bool... Values, class Mode, class = std::enable_if_t<std::is_floating_point<T>::value>>
         XSIMD_INLINE batch<T, A> load_masked(T const* mem, batch_bool_constant<T, A, Values...> mask, convert<T>, Mode, requires_arch<avx>) noexcept
@@ -1019,12 +1034,8 @@ namespace xsimd
         // store_masked
         namespace detail
         {
-            // True when batch_bool<T, A> is the legacy VEX vector mask, i.e. it is stored
-            // in the same register as the data (__m256 / __m256d) rather than in an EVEX
-            // k-register (__mmask8) as on the avx512vl architectures. The _mm256_cast*_si256
-            // path below is only well-formed for the vector-mask representation. This names
-            // no architecture — it tests the mask's representation, in the spirit of
-            // detail::masked_memory_uses_fp_bitcast.
+            // True when batch_bool<T, A> shares the data register (__m256/__m256d) rather
+            // than an EVEX k-register; the _mm256_cast*_si256 path below needs the former.
             template <class T, class A>
             using uses_vector_mask = std::is_same<typename batch_bool<T, A>::register_type,
                                                   typename batch<T, A>::register_type>;
@@ -1070,6 +1081,21 @@ namespace xsimd
             }
         }
 
+        // Runtime-mask store (float/double).
+        template <class A, class Mode>
+        XSIMD_INLINE void
+        store_masked(float* mem, batch<float, A> const& src, batch_bool<float, A> mask, Mode, requires_arch<avx>) noexcept
+        {
+            detail::maskstore(mem, mask, src);
+        }
+
+        template <class A, class Mode>
+        XSIMD_INLINE void
+        store_masked(double* mem, batch<double, A> const& src, batch_bool<double, A> mask, Mode, requires_arch<avx>) noexcept
+        {
+            detail::maskstore(mem, mask, 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

include/xsimd/arch/xsimd_avx2.hpp

@@ -117,18 +117,38 @@ namespace xsimd
             }
         }
 
-        // load_masked
-        // AVX2 low-level helpers (operate on raw SIMD registers)
+        // load_masked / store_masked: AVX2 has _mm256_maskload/maskstore_epi{32,64};
+        // 8/16-bit integers fall back to the common scalar path.
         namespace detail
         {
-            XSIMD_INLINE __m256i maskload(const int32_t* mem, __m256i mask) noexcept
+            template <class T>
+            XSIMD_INLINE __m256i maskload(T const* mem, __m256i mask) noexcept
             {
-                return _mm256_maskload_epi32(mem, mask);
+                XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
+                {
+                    static_assert(sizeof(int) == 4, "_mm256_maskload_epi32 requires a 4-byte int");
+                    return _mm256_maskload_epi32(reinterpret_cast<int const*>(mem), mask);
+                }
+                else
+                {
+                    static_assert(sizeof(long long) == 8, "_mm256_maskload_epi64 requires an 8-byte long long");
+                    return _mm256_maskload_epi64(reinterpret_cast<long long const*>(mem), mask);
+                }
             }
 
-            XSIMD_INLINE __m256i maskload(const long long* mem, __m256i mask) noexcept
+            template <class T>
+            XSIMD_INLINE void maskstore(T* mem, __m256i mask, __m256i src) noexcept
             {
-                return _mm256_maskload_epi64(reinterpret_cast<long long const*>(mem), mask);
+                XSIMD_IF_CONSTEXPR(sizeof(T) == 4)
+                {
+                    static_assert(sizeof(int) == 4, "_mm256_maskstore_epi32 requires a 4-byte int");
+                    _mm256_maskstore_epi32(reinterpret_cast<int*>(mem), mask, src);
+                }
+                else
+                {
+                    static_assert(sizeof(long long) == 8, "_mm256_maskstore_epi64 requires an 8-byte long long");
+                    _mm256_maskstore_epi64(reinterpret_cast<long long*>(mem), mask, src);
+                }
             }
 
             XSIMD_INLINE __m256i zero_extend(__m128i hi) noexcept
@@ -137,61 +157,22 @@ namespace xsimd
             }
         }
 
-        // single templated implementation for integer masked loads (32/64-bit)
         template <class A, class T, bool... Values, class Mode>
-        XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) >= 4), batch<T, A>>
+        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
         {
-            static_assert(sizeof(T) == 4 || sizeof(T) == 8, "load_masked supports only 32/64-bit integers on AVX2");
-            using int_t = std::conditional_t<sizeof(T) == 4, int32_t, long long>;
-            // Use the raw register-level maskload helpers for the remaining cases.
-            return detail::maskload(reinterpret_cast<const int_t*>(mem), mask.as_batch());
+            return detail::maskload(mem, mask.as_batch());
         }
 
-        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...> mask, convert<int32_t>, Mode, requires_arch<avx2>) noexcept
+        template <class A, class T, 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<T, A> mask, convert<T>, Mode, requires_arch<avx2>) noexcept
         {
-            return load_masked<A, int32_t>(mem, mask, convert<int32_t> {}, Mode {}, avx2 {});
-        }
-
-        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<avx2>) noexcept
-        {
-            const auto r = load_masked<A, int32_t>(reinterpret_cast<int32_t const*>(mem), batch_bool_constant<int32_t, A, Values...> {}, convert<int32_t> {}, Mode {}, avx2 {});
-            return bitwise_cast<uint32_t>(r);
-        }
-
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE batch<int64_t, A> load_masked(int64_t const* mem, batch_bool_constant<int64_t, A, Values...> mask, convert<int64_t>, Mode, requires_arch<avx2>) noexcept
-        {
-            return load_masked<A, int64_t>(mem, mask, convert<int64_t> {}, Mode {}, avx2 {});
-        }
-
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE batch<uint64_t, A> load_masked(uint64_t const* mem, batch_bool_constant<uint64_t, A, Values...>, convert<uint64_t>, Mode, requires_arch<avx2>) noexcept
-        {
-            const auto r = load_masked<A, int64_t>(reinterpret_cast<int64_t const*>(mem), batch_bool_constant<int64_t, A, Values...> {}, convert<int64_t> {}, Mode {}, avx2 {});
-            return bitwise_cast<uint64_t>(r);
-        }
-
-        // store_masked
-        namespace detail
-        {
-            template <class T, class A>
-            XSIMD_INLINE void maskstore(int32_t* mem, __m256i mask, __m256i src) noexcept
-            {
-                _mm256_maskstore_epi32(reinterpret_cast<int*>(mem), mask, src);
-            }
-
-            template <class T, class A>
-            XSIMD_INLINE void maskstore(int64_t* mem, __m256i mask, __m256i src) noexcept
-            {
-                _mm256_maskstore_epi64(reinterpret_cast<long long*>(mem), mask, src);
-            }
+            return detail::maskload(mem, __m256i(mask));
         }
 
         template <class A, class T, bool... Values, class Mode,
-                  typename = std::enable_if_t<std::is_integral<T>::value && (sizeof(T) >= 4)>>
+                  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;
@@ -214,22 +195,15 @@ namespace xsimd
             }
             else
             {
-                detail::maskstore<T, A>(mem, mask.as_batch(), src);
+                detail::maskstore(mem, mask.as_batch(), src);
             }
         }
 
-        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<avx2>) noexcept
-        {
-            const auto s32 = bitwise_cast<int32_t>(src);
-            store_masked<A>(reinterpret_cast<int32_t*>(mem), s32, batch_bool_constant<int32_t, A, Values...> {}, Mode {}, avx2 {});
-        }
-
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE void store_masked(uint64_t* mem, batch<uint64_t, A> const& src, batch_bool_constant<uint64_t, A, Values...>, Mode, requires_arch<avx2>) noexcept
+        template <class A, class T, class Mode>
+        XSIMD_INLINE std::enable_if_t<std::is_integral<T>::value && (sizeof(T) == 4 || sizeof(T) == 8), void>
+        store_masked(T* mem, batch<T, A> const& src, batch_bool<T, A> mask, Mode, requires_arch<avx2>) noexcept
         {
-            const auto s64 = bitwise_cast<int64_t>(src);
-            store_masked<A>(reinterpret_cast<int64_t*>(mem), s64, batch_bool_constant<int64_t, A, Values...> {}, Mode {}, avx2 {});
+            detail::maskstore(mem, __m256i(mask), __m256i(src));
         }
 
         // load_stream