Merge pull request #439 from serge-sans-paille/feature/improve-select

Provide batch_constant and batch_bool_constant, and use them to speci…

Author: JohanMabille

docs/source/index.rst

@@ -77,6 +77,7 @@ This software is licensed under the BSD-3-Clause license. See the LICENSE file f
    api/instr_macros
    api/batch_index
    api/data_transfer
+   api/batch_manip
    api/math_index
    api/aligned_allocator
 

include/xsimd/types/xsimd_avx512_int_base.hpp

@@ -114,45 +114,64 @@ namespace xsimd
 
     namespace avx512_detail
     {
-        template<class Tup, std::size_t... Is>
-        __m512i revert_args_set_epi8(Tup&& t, detail::index_sequence<Is...>)
+        inline __m512i int_init(std::integral_constant<std::size_t, 1>,
+                         int8_t t0, int8_t t1, int8_t t2, int8_t t3,
+                         int8_t t4, int8_t t5, int8_t t6, int8_t t7,
+                         int8_t t8, int8_t t9, int8_t t10, int8_t t11,
+                         int8_t t12, int8_t t13, int8_t t14, int8_t t15,
+                         int8_t t16, int8_t t17, int8_t t18, int8_t t19,
+                         int8_t t20, int8_t t21, int8_t t22, int8_t t23,
+                         int8_t t24, int8_t t25, int8_t t26, int8_t t27,
+                         int8_t t28, int8_t t29, int8_t t30, int8_t t31,
+                         int8_t t32, int8_t t33, int8_t t34, int8_t t35,
+                         int8_t t36, int8_t t37, int8_t t38, int8_t t39,
+                         int8_t t40, int8_t t41, int8_t t42, int8_t t43,
+                         int8_t t44, int8_t t45, int8_t t46, int8_t t47,
+                         int8_t t48, int8_t t49, int8_t t50, int8_t t51,
+                         int8_t t52, int8_t t53, int8_t t54, int8_t t55,
+                         int8_t t56, int8_t t57, int8_t t58, int8_t t59,
+                         int8_t t60, int8_t t61, int8_t t62, int8_t t63)
         {
-            // funny, this instruction is not yet implemented in clang or gcc (will come in future versions)
 #if defined(__clang__) || __GNUC__
             return __extension__ (__m512i)(__v64qi)
             {
-                static_cast<char>(std::get<Is>(std::forward<Tup>(t)))...
+              t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15,
+              t16, t17, t18, t19, t20, t21, t22, t23, t24, t25, t26, t27, t28, t29, t30, t31,
+              t32, t33, t34, t35, t36, t37, t38, t39, t40, t41, t42, t43, t44, t45, t46, t47,
+              t48, t49, t50, t51, t52, t53, t54, t55, t56, t57, t58, t59, t60, t61, t62, t63
             };
 #else
-            return _mm512_set_epi8(static_cast<char>(std::get<Is>(std::forward<Tup>(t)))...);
+            return _mm512_set_epi8(
+              t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15,
+              t16, t17, t18, t19, t20, t21, t22, t23, t24, t25, t26, t27, t28, t29, t30, t31,
+              t32, t33, t34, t35, t36, t37, t38, t39, t40, t41, t42, t43, t44, t45, t46, t47,
+              t48, t49, t50, t51, t52, t53, t54, t55, t56, t57, t58, t59, t60, t61, t62, t63);
 #endif
         }
 
-        template<class Tup, std::size_t... Is>
-        __m512i revert_args_set_epi16(Tup&& t, detail::index_sequence<Is...>)
+        inline __m512i int_init(std::integral_constant<std::size_t, 2>,
+                         int16_t t0, int16_t t1, int16_t t2, int16_t t3,
+                         int16_t t4, int16_t t5, int16_t t6, int16_t t7,
+                         int16_t t8, int16_t t9, int16_t t10, int16_t t11,
+                         int16_t t12, int16_t t13, int16_t t14, int16_t t15,
+                         int16_t t16, int16_t t17, int16_t t18, int16_t t19,
+                         int16_t t20, int16_t t21, int16_t t22, int16_t t23,
+                         int16_t t24, int16_t t25, int16_t t26, int16_t t27,
+                         int16_t t28, int16_t t29, int16_t t30, int16_t t31)
         {
 #if defined(__clang__) || __GNUC__
             return __extension__ (__m512i)(__v32hi)
             {
-                static_cast<short>(std::get<Is>(std::forward<Tup>(t)))...
+              t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15,
+              t16, t17, t18, t19, t20, t21, t22, t23, t24, t25, t26, t27, t28, t29, t30, t31
             };
 #else
-            return _mm512_set_epi16(static_cast<short>(std::get<Is>(std::forward<Tup>(t)))...);
+            return _mm512_set_epi16(
+              t0, t1, t2, t3, t4, t5, t6, t7, t8, t9, t10, t11, t12, t13, t14, t15,
+              t16, t17, t18, t19, t20, t21, t22, t23, t24, t25, t26, t27, t28, t29, t30, t31);
 #endif
         }
 
-        template <class... Args>
-        __m512i int_init(std::integral_constant<std::size_t, 1>, Args... args)
-        {
-            return revert_args_set_epi8(std::forward_as_tuple(args...), detail::make_index_sequence<sizeof...(Args)>{});
-        }
-
-        template <class... Args>
-        __m512i int_init(std::integral_constant<std::size_t, 2>, Args... args)
-        {
-            return revert_args_set_epi16(std::forward_as_tuple(args...), detail::make_index_sequence<sizeof...(Args)>{});
-        }
-
         inline __m512i int_init(std::integral_constant<std::size_t, 4>,
                                 int32_t t0, int32_t t1, int32_t t2, int32_t t3,
                                 int32_t t4, int32_t t5, int32_t t6, int32_t t7,

include/xsimd/types/xsimd_avx_double.hpp

@@ -658,6 +658,13 @@ namespace xsimd
                 return _mm256_blendv_pd(b, a, cond);
             }
 
+            template<bool... Values>
+            static batch_type select(const batch_bool_constant<value_type, Values...>&, const batch_type& a, const batch_type& b)
+            {
+                constexpr int mask = batch_bool_constant<value_type, Values...>::mask();
+                return _mm256_blend_pd(b, a, mask);
+            }
+
             static batch_bool_type isnan(const batch_type& x)
             {
                 return _mm256_cmp_pd(x, x, _CMP_UNORD_Q);

include/xsimd/types/xsimd_avx_float.hpp

@@ -706,6 +706,13 @@ namespace xsimd
                 return _mm256_blendv_ps(b, a, cond);
             }
 
+            template<bool... Values>
+            static batch_type select(const batch_bool_constant<value_type, Values...>& cond, const batch_type& a, const batch_type& b)
+            {
+                constexpr int mask = batch_bool_constant<value_type, Values...>::mask();
+                return _mm256_blend_ps(b, a, mask);
+            }
+
             static batch_bool_type isnan(const batch_type& x)
             {
                 return _mm256_cmp_ps(x, x, _CMP_UNORD_Q);

include/xsimd/types/xsimd_base.hpp

@@ -25,6 +25,7 @@
 #include "../memory/xsimd_alignment.hpp"
 #include "xsimd_utils.hpp"
 #include "xsimd_base_bool.hpp"
+#include "xsimd_base_constant.hpp"
 
 namespace xsimd
 {
@@ -220,7 +221,7 @@ namespace xsimd
         simd_batch(simd_batch&&) = default;
         simd_batch& operator=(simd_batch&&) = default;
 
-        simd_batch(storage_type value);
+        constexpr simd_batch(storage_type value);
 
         using char_itype =
             typename std::conditional<std::is_signed<char>::value, int8_t, uint8_t>::type;
@@ -679,7 +680,7 @@ namespace xsimd
      *****************************/
 
     template <class X>
-    inline simd_batch<X>::simd_batch(storage_type value)
+    constexpr inline simd_batch<X>::simd_batch(storage_type value)
         : m_value(value)
     {
     }
@@ -1754,6 +1755,28 @@ namespace xsimd
         return kernel::select(cond(), a(), b());
     }
 
+    /**
+     * @ingroup simd_batch_miscellaneous
+     *
+     * Ternary operator for batches: selects values from the batches \c a or \c b
+     * depending on the boolean values in the constant batch \c cond. Equivalent to
+     * \code{.cpp}
+     * for(std::size_t i = 0; i < N; ++i)
+     *     res[i] = cond[i] ? a[i] : b[i];
+     * \endcode
+     * @param cond constant batch condition.
+     * @param a batch values for truthy condition.
+     * @param b batch value for falsy condition.
+     * @return the result of the selection.
+     */
+    template <class X, bool... Masks>
+    inline batch_type_t<X> select(const batch_bool_constant<typename simd_batch_traits<X>::value_type, Masks...>& cond, const simd_base<X>& a, const simd_base<X>& b)
+    {
+        using value_type = typename simd_batch_traits<X>::value_type;
+        using kernel = detail::batch_kernel<value_type, simd_batch_traits<X>::size>;
+        return kernel::select(cond, a(), b());
+    }
+
     /**
      * Determines if the scalars in the given batch \c x are NaN values.
      * @param x batch of floating point values.

include/xsimd/types/xsimd_base_constant.hpp

@@ -0,0 +1,112 @@
+/***************************************************************************
+* Copyright (c) Serge Guelton                                              *
+* Copyright (c) QuantStack                                                 *
+*                                                                          *
+* Distributed under the terms of the BSD 3-Clause License.                 *
+*                                                                          *
+* The full license is in the file LICENSE, distributed with this software. *
+****************************************************************************/
+
+#ifndef XSIMD_BASE_CONSTANT_HPP
+#define XSIMD_BASE_CONSTANT_HPP
+
+namespace xsimd
+{
+    template <class X>
+    class simd_base;
+
+    template <class T, bool... Values>
+    struct batch_bool_constant
+    {
+        static constexpr std::size_t size = sizeof...(Values);
+        using value_type = bool;
+        using batch_type =
+            typename simd_batch_traits<batch<T, size>>::batch_bool_type;
+
+        batch_type operator()() const { return *this; }
+
+        operator batch_type() const { return {Values...}; }
+
+        bool operator[](size_t i) const
+        {
+            return std::array<value_type, size>{Values...}[i];
+        }
+        static constexpr int mask()
+        {
+            return mask_helper(0, static_cast<int>(Values)...);
+        }
+
+      private:
+        static constexpr int mask_helper(int acc) { return acc; }
+        template <class... Tys>
+        static constexpr int mask_helper(int acc, int mask, Tys... masks)
+        {
+            return mask_helper(acc | mask, (masks << 1)...);
+        }
+    };
+
+    template <class T, T... Values>
+    struct batch_constant
+    {
+        static constexpr std::size_t size = sizeof...(Values);
+        using value_type = T;
+        using batch_type = batch<T, size>;
+
+        batch_type operator()() const { return *this; }
+
+        operator batch_type() const { return {Values...}; }
+
+        constexpr T operator[](size_t i) const
+        {
+            return std::array<value_type, size>{Values...}[i];
+        }
+    };
+
+    namespace detail
+    {
+        template <class G, std::size_t... Is>
+        constexpr auto make_batch_constant(detail::index_sequence<Is...>)
+            -> batch_constant<decltype(G::get(0, 0)),
+                              G::get(Is, sizeof...(Is))...>
+        {
+            return {};
+        }
+        template <class T, class G, std::size_t... Is>
+        constexpr auto make_batch_bool_constant(detail::index_sequence<Is...>)
+            -> batch_bool_constant<T, G::get(Is, sizeof...(Is))...>
+        {
+            return {};
+        }
+        template <class T, T value, std::size_t... Is>
+        constexpr auto make_batch_constant(detail::index_sequence<Is...>)
+            -> batch_constant<T, (Is, value)...>
+        {
+            return {};
+        }
+        template <class T, T value, std::size_t... Is>
+        constexpr auto make_batch_bool_constant(detail::index_sequence<Is...>)
+            -> batch_bool_constant<T, (Is, value)...>
+        {
+            return {};
+        }
+    } // namespace detail
+
+    template <class G, std::size_t N>
+    constexpr auto make_batch_constant() -> decltype(
+        detail::make_batch_constant<G>(detail::make_index_sequence<N>()))
+    {
+        return detail::make_batch_constant<G>(detail::make_index_sequence<N>());
+    }
+
+    template <class T, class G, std::size_t N>
+    constexpr auto make_batch_bool_constant()
+        -> decltype(detail::make_batch_bool_constant<T, G>(
+            detail::make_index_sequence<N>()))
+    {
+        return detail::make_batch_bool_constant<T, G>(
+            detail::make_index_sequence<N>());
+    }
+
+} // namespace xsimd
+
+#endif

include/xsimd/types/xsimd_fallback.hpp

@@ -917,6 +917,12 @@ namespace xsimd
                 XSIMD_FALLBACK_MAPPING_LOOP(batch, (cond[i] ? a[i] : b[i]))
             }
 
+            template<bool... Values>
+            static batch_type select(const batch_bool_constant<value_type, Values...>& cond, const batch_type& a, const batch_type& b)
+            {
+                XSIMD_FALLBACK_MAPPING_LOOP(batch, (cond[i] ? a[i] : b[i]))
+            }
+
             static batch_bool_type isnan(const batch_type& x)
             {
                 XSIMD_FALLBACK_MAPPING_LOOP(batch_bool, std::isnan(x[i]))

include/xsimd/types/xsimd_sse_double.hpp

@@ -629,6 +629,17 @@ namespace xsimd
 #endif
             }
 
+            template<bool... Values>
+            static batch_type select(const batch_bool_constant<value_type, Values...>& cond, const batch_type& a, const batch_type& b)
+            {
+#if XSIMD_X86_INSTR_SET >= XSIMD_X86_SSE4_1_VERSION
+                constexpr int mask = batch_bool_constant<value_type, Values...>::mask();
+                return _mm_blend_pd(b, a, mask);
+#else
+                return select(cond(), a, b);
+#endif
+            }
+
             static batch_bool_type isnan(const batch_type& x)
             {
                 return _mm_cmpunord_pd(x, x);

include/xsimd/types/xsimd_sse_float.hpp

@@ -692,6 +692,18 @@ namespace xsimd
 #endif
             }
 
+            template<bool... Values>
+            static batch_type select(const batch_bool_constant<value_type, Values...>& cond, const batch_type& a, const batch_type& b)
+            {
+#if XSIMD_X86_INSTR_SET >= XSIMD_X86_SSE4_1_VERSION
+                (void)cond;
+                constexpr int mask = batch_bool_constant<value_type, Values...>::mask();
+                return _mm_blend_ps(b, a, mask);
+#else
+                return select((batch_bool_type)cond, a, b);
+#endif
+            }
+
             static batch_bool_type isnan(const batch_type& x)
             {
                 return _mm_cmpunord_ps(x, x);

include/xsimd/types/xsimd_sse_int_base.hpp

@@ -70,7 +70,7 @@ namespace xsimd
         sse_int_batch();
         explicit sse_int_batch(T i);
         template <class... Args, class Enable = detail::is_array_initializer_t<T, N, Args...>>
-        sse_int_batch(Args... args);
+        constexpr sse_int_batch(Args... args);
         explicit sse_int_batch(const T* src);
         sse_int_batch(const T* src, aligned_mode);
         sse_int_batch(const T* src, unaligned_mode);
@@ -324,7 +324,7 @@ namespace xsimd
 
     template <class T, std::size_t N>
     template <class... Args, class>
-    inline sse_int_batch<T, N>::sse_int_batch(Args... args)
+    constexpr inline sse_int_batch<T, N>::sse_int_batch(Args... args)
         : base_type(sse_detail::int_init(std::integral_constant<std::size_t, sizeof(T)>{}, args...))
     {
     }