Simplify common_memory load_masked

Author: AntoinePrv

include/xsimd/arch/common/xsimd_common_memory.hpp

@@ -13,6 +13,7 @@
 #define XSIMD_COMMON_MEMORY_HPP
 
 #include "../../types/xsimd_batch_constant.hpp"
+#include "../../utils/xsimd_type_traits.hpp"
 #include "./xsimd_common_details.hpp"
 
 #include <algorithm>
@@ -360,88 +361,81 @@ namespace xsimd
             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
+        namespace detail
         {
-            constexpr std::size_t size = batch<T_in, A>::size;
-            constexpr bool mask[size] = { Values... };
+            // Compile-time dispatch tag for the common `load_masked`/ `store_masked`
+            // implementations: true iff we can use the int->float bitcast path (matching-size
+            // integer T_in/T_out with a SIMD register available for the matching
+            // floating-point type), false otherwise (use the scalar buffer fallback).
+            template <class A, class T_in, class T_out>
+            using common_masked_via_fp = std::integral_constant<bool,
+                                                                std::is_same<T_in, T_out>::value
+                                                                    && std::is_integral<T_out>::value
+                                                                    && !std::is_void<sized_fp_t<sizeof(T_out)>>::value
+                                                                    && types::has_simd_register<sized_fp_t<sizeof(T_out)>, A>::value>;
 
-            for (std::size_t i = 0; i < size; ++i)
-                if (mask[i])
-                {
-                    mem[i] = static_cast<T_out>(src.get(i));
-                }
-        }
+            // Scalar-buffer fallback: works for any T_in/T_out.
+            template <class A, class T_in, class T_out, bool... Values, class alignment>
+            XSIMD_INLINE batch<T_out, A>
+            load_masked_common(T_in const* mem, batch_bool_constant<T_out, A, Values...>, convert<T_out>, alignment, std::false_type /* via_fp */) 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... };
 
-        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);
-        }
+                for (std::size_t i = 0; i < size; ++i)
+                    buffer[i] = mask[i] ? static_cast<T_out>(mem[i]) : T_out(0);
 
-        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);
-        }
+                return batch<T_out, A>::load(buffer.data(), aligned_mode {});
+            }
 
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE std::enable_if_t<types::has_simd_register<double, A>::value, batch<int64_t, A>>
-        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);
-        }
+            // Integer-via-float bitcast: T_in == T_out == integral T with a matching
+            // `sized_fp_t<sizeof(T)>` for which the arch has a SIMD register.
+            // Dispatches to the floating `load_masked` (which is arch-specialized) and bitcasts back.
+            template <class A, class T, bool... Values, class Mode>
+            XSIMD_INLINE batch<T, A>
+            load_masked_common(T const* mem, batch_bool_constant<T, A, Values...>, convert<T>, Mode, std::true_type /* via_fp */) noexcept
+            {
+                using fp_t = sized_fp_t<sizeof(T)>;
+                const auto f = ::xsimd::kernel::load_masked<A>(reinterpret_cast<const fp_t*>(mem), batch_bool_constant<fp_t, A, Values...> {}, convert<fp_t> {}, Mode {}, A {});
+                return bitwise_cast<T>(f);
+            }
 
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE std::enable_if_t<types::has_simd_register<double, A>::value, batch<uint64_t, A>>
-        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, class T_in, class T_out, bool... Values, class alignment>
+            XSIMD_INLINE void
+            store_masked_common(T_out* mem, batch<T_in, A> const& src, batch_bool_constant<T_in, A, Values...>, alignment, std::false_type /* via_fp */) noexcept
+            {
+                constexpr std::size_t size = batch<T_in, A>::size;
+                constexpr bool mask[size] = { Values... };
 
-        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 {});
-        }
+                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 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, class T, bool... Values, class Mode>
+            XSIMD_INLINE void
+            store_masked_common(T* mem, batch<T, A> const& src, batch_bool_constant<T, A, Values...>, Mode, std::true_type /* via_fp */) noexcept
+            {
+                using fp_t = sized_fp_t<sizeof(T)>;
+                ::xsimd::kernel::store_masked<A>(reinterpret_cast<fp_t*>(mem), bitwise_cast<fp_t>(src), batch_bool_constant<fp_t, A, Values...> {}, Mode {}, A {});
+            }
+        } // namespace detail
 
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE std::enable_if_t<types::has_simd_register<double, A>::value>
-        store_masked(int64_t* mem, batch<int64_t, A> const& src, batch_bool_constant<int64_t, A, Values...>, Mode, requires_arch<A>) noexcept
+        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...> mask, convert<T_out> cvt, alignment mode, requires_arch<common>) noexcept
         {
-            store_masked<A>(reinterpret_cast<double*>(mem), bitwise_cast<double>(src), batch_bool_constant<double, A, Values...> {}, Mode {}, A {});
+            return detail::load_masked_common(mem, mask, cvt, mode, detail::common_masked_via_fp<A, T_in, T_out> {});
         }
 
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE std::enable_if_t<types::has_simd_register<double, A>::value>
-        store_masked(uint64_t* mem, batch<uint64_t, A> const& src, batch_bool_constant<uint64_t, A, Values...>, Mode, requires_arch<A>) noexcept
+        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
         {
-            store_masked<A>(reinterpret_cast<double*>(mem), bitwise_cast<double>(src), batch_bool_constant<double, A, Values...> {}, Mode {}, A {});
+            detail::store_masked_common(mem, src, mask, mode, detail::common_masked_via_fp<A, T_in, T_out> {});
         }
 
         template <class A, class T_in, class T_out>

include/xsimd/arch/xsimd_avx512f.hpp

@@ -298,25 +298,11 @@ namespace xsimd
 
         } // namespace detail
 
-        // The AVX512F masked-load logic lives in this plain `*_avx512f` helper
-        // (no `requires_arch` tag) and is exposed through the concrete
-        // element-type overloads below.
-        //
-        // Why not a single generic `load_masked(T const*, ..., requires_arch<avx512f>)`?
-        // It is ambiguous against the concrete-type / generic-arch overloads in
-        // xsimd_common_memory.hpp (e.g. `load_masked(int32_t const*, ...,
-        // requires_arch<A>)`): the avx512f overload is more specialized on the
-        // architecture while the common one is more specialized on the pointer
-        // type, so partial ordering cannot pick a winner. When AVX512DQ/BW is
-        // available a fully concrete `requires_arch<avx512dq>` overload is the
-        // unique best match and hides this, but a pure-AVX512F target (the
-        // `avx512f` preset) has no such tie-breaker and the call fails to
-        // compile. Concrete element-type `requires_arch<avx512f>` overloads make
-        // the avx512f candidate the unique best match for every integer type.
-        template <class A, class T, bool... Values, class Mode>
-        XSIMD_INLINE batch<T, A> load_masked_avx512f(T const* mem,
-                                                     batch_bool_constant<T, A, Values...> mask,
-                                                     Mode) noexcept
+        template <class A, class T, bool... Values, class Mode,
+                  typename = std::enable_if_t<(sizeof(T) >= 4)>>
+        XSIMD_INLINE batch<T, A> load_masked(T const* mem,
+                                             batch_bool_constant<T, A, Values...> mask,
+                                             convert<T>, Mode, requires_arch<avx512f>) noexcept
         {
             constexpr auto half = batch<T, A>::size / 2;
             XSIMD_IF_CONSTEXPR(mask.countl_zero() >= half) // lower-half AVX2 forwarding
@@ -338,61 +324,12 @@ namespace xsimd
             }
         }
 
-        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<avx512f>) noexcept
-        {
-            return load_masked_avx512f(mem, mask, Mode {});
-        }
-
-        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...> mask,
-                                                    convert<uint32_t>, Mode, requires_arch<avx512f>) noexcept
-        {
-            return load_masked_avx512f(mem, mask, Mode {});
-        }
-
-        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<avx512f>) noexcept
-        {
-            return load_masked_avx512f(mem, mask, Mode {});
-        }
-
-        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...> mask,
-                                                    convert<uint64_t>, Mode, requires_arch<avx512f>) noexcept
-        {
-            return load_masked_avx512f(mem, mask, Mode {});
-        }
-
-        // Non-integer element types only (float, double, ...): integer types
-        // are handled by the concrete overloads above. gcc 10's partial
-        // ordering cannot break the tie between a concrete-type avx512f
-        // overload and a generic-T avx512f overload, so this catch-all must
-        // exclude the integer types we already specialized.
         template <class A, class T, bool... Values, class Mode,
-                  typename = std::enable_if_t<(sizeof(T) >= 4) && !std::is_integral<T>::value>>
-        XSIMD_INLINE batch<T, A> load_masked(T const* mem,
-                                             batch_bool_constant<T, A, Values...> mask,
-                                             convert<T>, Mode, requires_arch<avx512f>) noexcept
-        {
-            return load_masked_avx512f(mem, mask, Mode {});
-        }
-
-        // Same ambiguity as load_masked above (see comment there): factor the
-        // AVX512F store logic into a plain helper and expose it via concrete
-        // element-type `requires_arch<avx512f>` overloads so a pure-AVX512F
-        // target has a unique best match.
-        template <class A, class T, bool... Values, class Mode>
-        XSIMD_INLINE void store_masked_avx512f(T* mem,
-                                               batch<T, A> const& src,
-                                               batch_bool_constant<T, A, Values...> mask,
-                                               Mode) noexcept
+                  typename = std::enable_if_t<(sizeof(T) >= 4)>>
+        XSIMD_INLINE void store_masked(T* mem,
+                                       batch<T, A> const& src,
+                                       batch_bool_constant<T, A, Values...> mask,
+                                       Mode, requires_arch<avx512f>) noexcept
         {
             constexpr auto half = batch<T, A>::size / 2;
             XSIMD_IF_CONSTEXPR(mask.countl_zero() >= half) // lower-half AVX2 forwarding
@@ -414,50 +351,6 @@ namespace xsimd
             }
         }
 
-        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...> mask,
-                                       Mode, requires_arch<avx512f>) noexcept
-        {
-            store_masked_avx512f(mem, src, mask, Mode {});
-        }
-
-        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...> mask,
-                                       Mode, requires_arch<avx512f>) noexcept
-        {
-            store_masked_avx512f(mem, src, mask, Mode {});
-        }
-
-        template <class A, bool... Values, class Mode>
-        XSIMD_INLINE void store_masked(int64_t* mem, batch<int64_t, A> const& src,
-                                       batch_bool_constant<int64_t, A, Values...> mask,
-                                       Mode, requires_arch<avx512f>) noexcept
-        {
-            store_masked_avx512f(mem, src, mask, Mode {});
-        }
-
-        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...> mask,
-                                       Mode, requires_arch<avx512f>) noexcept
-        {
-            store_masked_avx512f(mem, src, mask, Mode {});
-        }
-
-        // Non-integer element types only: see load_masked above for the gcc 10
-        // partial ordering rationale.
-        template <class A, class T, bool... Values, class Mode,
-                  typename = std::enable_if_t<(sizeof(T) >= 4) && !std::is_integral<T>::value>>
-        XSIMD_INLINE void store_masked(T* mem,
-                                       batch<T, A> const& src,
-                                       batch_bool_constant<T, A, Values...> mask,
-                                       Mode, requires_arch<avx512f>) noexcept
-        {
-            store_masked_avx512f(mem, src, mask, Mode {});
-        }
-
         // abs
         template <class A>
         XSIMD_INLINE batch<float, A> abs(batch<float, A> const& self, requires_arch<avx512f>) noexcept