Add select for batch_bool

Author: kalenedrael

include/xsimd/arch/common/xsimd_common_logical.hpp

@@ -212,6 +212,35 @@ namespace xsimd
                     res |= 1ul << i;
             return res;
         }
+
+        // select
+        namespace detail
+        {
+            template <typename T, typename A>
+            using is_batch_bool_register_same = std::is_same<typename batch_bool<T, A>::register_type, typename batch<T, A>::register_type>;
+        }
+
+        template <class A, class T, typename std::enable_if<detail::is_batch_bool_register_same<T, A>::value, int>::type = 3>
+        XSIMD_INLINE batch_bool<T, A> select(batch_bool<T, A> const& cond, batch_bool<T, A> const& true_br, batch_bool<T, A> const& false_br, requires_arch<common>)
+        {
+            using register_type = typename batch_bool<T, A>::register_type;
+            // Do not cast, but rather reinterpret the masks as batches.
+            const auto true_v = batch<T, A> { static_cast<register_type>(true_br) };
+            const auto false_v = batch<T, A> { static_cast<register_type>(false_br) };
+            return batch_bool<T, A> { select(cond, true_v, false_v) };
+        }
+
+        template <class A, class T, typename std::enable_if<!detail::is_batch_bool_register_same<T, A>::value, int>::type = 3>
+        XSIMD_INLINE batch_bool<T, A> select(batch_bool<T, A> const& cond, batch_bool<T, A> const& true_br, batch_bool<T, A> const& false_br, requires_arch<common>)
+        {
+            return (true_br & cond) | (bitwise_andnot(false_br, cond));
+        }
+
+        template <class A, class T, bool... Values>
+        XSIMD_INLINE batch_bool<T, A> select(batch_bool_constant<T, A, Values...> const& cond, batch_bool<T, A> const& true_br, batch_bool<T, A> const& false_br, requires_arch<common>)
+        {
+            return (true_br & cond) | (false_br & ~cond);
+        }
     }
 }
 

include/xsimd/types/xsimd_api.hpp

@@ -2099,6 +2099,27 @@ namespace xsimd
         return kernel::select<A>(cond, true_br, false_br, A {});
     }
 
+    /**
+     * @ingroup batch_bool_logical
+     *
+     * Ternary operator for conditions: selects values from the batches \c true_br or \c false_br
+     * 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] ? true_br[i] : false_br[i];
+     * \endcode
+     * @param cond batch condition.
+     * @param true_br batch values for truthy condition.
+     * @param false_br batch value for falsy condition.
+     * @return the result of the selection.
+     */
+    template <class T, class A>
+    XSIMD_INLINE batch_bool<T, A> select(batch_bool<T, A> const& cond, batch_bool<T, A> const& true_br, batch_bool<T, A> const& false_br) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::select<A>(cond, true_br, false_br, A {});
+    }
+
     /**
      * @ingroup batch_cond
      *
@@ -2141,6 +2162,27 @@ namespace xsimd
         return kernel::select<A>(cond, true_br, false_br, A {});
     }
 
+    /**
+     * @ingroup batch_cond
+     *
+     * Ternary operator for mask batches: selects values from the masks \c true_br or \c false_br
+     * 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] ? true_br[i] : false_br[i];
+     * \endcode
+     * @param cond constant batch condition.
+     * @param true_br batch values for truthy condition.
+     * @param false_br batch value for falsy condition.
+     * @return the result of the selection.
+     */
+    template <class T, class A, bool... Values>
+    XSIMD_INLINE batch_bool<T, A> select(batch_bool_constant<T, A, Values...> const& cond, batch_bool<T, A> const& true_br, batch_bool<T, A> const& false_br) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::select<A>(cond, true_br, false_br, A {});
+    }
+
     /**
      * @ingroup batch_data_transfer
      *

test/test_select.cpp

@@ -18,22 +18,26 @@ template <class B>
 struct select_test
 {
     using batch_type = B;
+    using batch_bool_type = typename B::batch_bool_type;
     using value_type = typename B::value_type;
     using arch_type = typename B::arch_type;
     static constexpr size_t size = B::size;
-    using vector_type = std::vector<value_type>;
+    static constexpr size_t nb_input = size * 10000;
+    using vector_type = std::array<value_type, nb_input>;
+    using vector_bool_type = std::array<bool, nb_input>;
 
-    size_t nb_input;
     vector_type lhs_input;
     vector_type rhs_input;
     vector_type expected;
     vector_type res;
 
+    vector_bool_type lhs_input_b;
+    vector_bool_type rhs_input_b;
+    vector_bool_type expected_b;
+    vector_bool_type res_b;
+
     select_test()
     {
-        nb_input = size * 10000;
-        lhs_input.resize(nb_input);
-        rhs_input.resize(nb_input);
         auto clamp = [](double v)
         {
             return static_cast<value_type>(std::min(v, static_cast<double>(std::numeric_limits<value_type>::max())));
@@ -42,28 +46,37 @@ struct select_test
         {
             lhs_input[i] = clamp(i / 4 + 1.2 * std::sqrt(i + 0.25));
             rhs_input[i] = clamp(10.2 / (i + 2) + 0.25);
+            lhs_input_b[i] = (int)lhs_input[i] % 2;
+            rhs_input_b[i] = (int)rhs_input[i] % 2;
         }
-        expected.resize(nb_input);
-        res.resize(nb_input);
     }
 
     void test_select_dynamic()
     {
         for (size_t i = 0; i < nb_input; ++i)
         {
             expected[i] = lhs_input[i] > value_type(3) ? lhs_input[i] : rhs_input[i];
+            expected_b[i] = lhs_input[i] > value_type(3) ? lhs_input_b[i] : rhs_input_b[i];
         }
 
-        batch_type lhs_in, rhs_in, out;
+        batch_type lhs_in, rhs_in;
+        batch_bool_type lhs_in_b, rhs_in_b;
         for (size_t i = 0; i < nb_input; i += size)
         {
             detail::load_batch(lhs_in, lhs_input, i);
             detail::load_batch(rhs_in, rhs_input, i);
-            out = xsimd::select(lhs_in > value_type(3), lhs_in, rhs_in);
+            const auto out = xsimd::select(lhs_in > value_type(3), lhs_in, rhs_in);
             detail::store_batch(out, res, i);
+
+            detail::load_batch(lhs_in_b, lhs_input_b, i);
+            detail::load_batch(rhs_in_b, rhs_input_b, i);
+            const auto out_b = xsimd::select(lhs_in > value_type(3), lhs_in_b, rhs_in_b);
+            detail::store_batch(out_b, res_b, i);
         }
         size_t diff = detail::get_nb_diff(res, expected);
+        size_t diff_b = detail::get_nb_diff(res_b, expected_b);
         CHECK_EQ(diff, 0);
+        CHECK_EQ(diff_b, 0);
     }
     struct pattern
     {
@@ -77,25 +90,35 @@ struct select_test
         for (size_t i = 0; i < nb_input; ++i)
         {
             expected[i] = mask.get(i % size) ? lhs_input[i] : rhs_input[i];
+            expected_b[i] = mask.get(i % size) ? lhs_input_b[i] : rhs_input_b[i];
         }
 
-        batch_type lhs_in, rhs_in, out;
+        batch_type lhs_in, rhs_in;
+        batch_bool_type lhs_in_b, rhs_in_b;
         for (size_t i = 0; i < nb_input; i += size)
         {
             detail::load_batch(lhs_in, lhs_input, i);
             detail::load_batch(rhs_in, rhs_input, i);
-            out = xsimd::select(mask, lhs_in, rhs_in);
+            const auto out = xsimd::select(mask, lhs_in, rhs_in);
             detail::store_batch(out, res, i);
+
+            detail::load_batch(lhs_in_b, lhs_input_b, i);
+            detail::load_batch(rhs_in_b, rhs_input_b, i);
+            const auto out_b = xsimd::select(mask, lhs_in_b, rhs_in_b);
+            detail::store_batch(out_b, res_b, i);
         }
         size_t diff = detail::get_nb_diff(res, expected);
+        size_t diff_b = detail::get_nb_diff(res_b, expected_b);
         CHECK_EQ(diff, 0);
+        CHECK_EQ(diff_b, 0);
     }
 };
 
 TEST_CASE_TEMPLATE("[select]", B, BATCH_TYPES)
 {
-    select_test<B> Test;
-    SUBCASE("select_dynamic") { Test.test_select_dynamic(); }
-    SUBCASE("select_static") { Test.test_select_static(); }
+    // Allocate on heap to avoid stack overflow from excessively large object.
+    std::unique_ptr<select_test<B>> Test { new select_test<B> };
+    SUBCASE("select_dynamic") { Test->test_select_dynamic(); }
+    SUBCASE("select_static") { Test->test_select_static(); }
 }
 #endif