From 2bbeb9c7a6ebab1dd14347fa6a2f0020be9781a9 Mon Sep 17 00:00:00 2001
From: serge-sans-paille <serge.guelton@telecom-bretagne.eu>
Date: Mon, 26 Feb 2024 19:49:24 +0100
Subject: [PATCH] avg

---
 .../xsimd/arch/generic/xsimd_generic_math.hpp | 34 ++++++++++++++++++-
 include/xsimd/arch/xsimd_avx2.hpp             | 17 ++++++++++
 include/xsimd/arch/xsimd_avx512bw.hpp         | 17 ++++++++++
 include/xsimd/arch/xsimd_sse2.hpp             | 20 +++++++++++
 include/xsimd/types/xsimd_api.hpp             | 15 ++++++++
 test/test_batch.cpp                           | 24 +++++++++++++
 test/test_xsimd_api.cpp                       |  7 ++++
 7 files changed, 133 insertions(+), 1 deletion(-)

diff --git a/include/xsimd/arch/generic/xsimd_generic_math.hpp b/include/xsimd/arch/generic/xsimd_generic_math.hpp
index 8fa887dc5..d99196f61 100644
--- a/include/xsimd/arch/generic/xsimd_generic_math.hpp
+++ b/include/xsimd/arch/generic/xsimd_generic_math.hpp
@@ -26,7 +26,7 @@ namespace xsimd
 
         using namespace types;
         // abs
-        template <class A, class T, class /*=typename std::enable_if<std::is_integral<T>::value, void>::type*/>
+        template <class A, class T, class>
         inline batch<T, A> abs(batch<T, A> const& self, requires_arch<generic>) noexcept
         {
             if (std::is_unsigned<T>::value)
@@ -45,6 +45,38 @@ namespace xsimd
             return hypot(z.real(), z.imag());
         }
 
+        // avg
+        namespace detail {
+          template <class A, class T>
+          inline batch<T, A> avg(batch<T, A> const& x, batch<T, A> const& y, std::true_type, std::false_type) noexcept
+          {
+            return (incr(x) + y) >> 1;
+          }
+
+          template <class A, class T>
+          inline batch<T, A> avg(batch<T, A> const& x, batch<T, A> const& y, std::true_type, std::true_type) noexcept
+          {
+            using uT = typename std::make_unsigned<T>::type;
+            // Inspired by
+            // https://stackoverflow.com/questions/5697500/take-the-average-of-two-signed-numbers-in-c
+            return select((x ^ y) < 0,
+                          (x + y) / 2,
+                          bitwise_cast<T>(avg(decr(bitwise_cast<uT>(x)), bitwise_cast<uT>(y)))
+                          );
+          }
+
+          template <class A, class T>
+          inline batch<T, A> avg(batch<T, A> const& x, batch<T, A> const& y, std::false_type, std::true_type) noexcept
+          {
+              return (x + y) / 2;
+          }
+        }
+
+        template <class A, class T>
+        inline batch<T, A> avg(batch<T, A> const& x, batch<T, A> const& y, requires_arch<generic>) noexcept {
+          return detail::avg(x, y, typename std::is_integral<T>::type{}, typename std::is_signed<T>::type{});
+        }
+
         // batch_cast
         template <class A, class T>
         inline batch<T, A> batch_cast(batch<T, A> const& self, batch<T, A> const&, requires_arch<generic>) noexcept
diff --git a/include/xsimd/arch/xsimd_avx2.hpp b/include/xsimd/arch/xsimd_avx2.hpp
index a5b07ec9d..452ac4855 100644
--- a/include/xsimd/arch/xsimd_avx2.hpp
+++ b/include/xsimd/arch/xsimd_avx2.hpp
@@ -76,6 +76,23 @@ namespace xsimd
             }
         }
 
+        // avg
+        template <class A, class T, class = typename std::enable_if<std::is_unsigned<T>::value, void>::type>
+        inline batch<T, A> avg(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx2>) noexcept
+        {
+            XSIMD_IF_CONSTEXPR(sizeof(T) == 1)
+            {
+                return _mm256_avg_epu8(self, other);
+            }
+            else XSIMD_IF_CONSTEXPR(sizeof(T) == 2)
+            {
+                return _mm256_avg_epu16(self, other);
+            }
+            else {
+              return avg(self, other, generic{});
+            }
+        }
+
         // bitwise_and
         template <class A, class T, class = typename std::enable_if<std::is_integral<T>::value, void>::type>
         inline batch<T, A> bitwise_and(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx2>) noexcept
diff --git a/include/xsimd/arch/xsimd_avx512bw.hpp b/include/xsimd/arch/xsimd_avx512bw.hpp
index 94a194dab..4d113a96c 100644
--- a/include/xsimd/arch/xsimd_avx512bw.hpp
+++ b/include/xsimd/arch/xsimd_avx512bw.hpp
@@ -112,6 +112,23 @@ namespace xsimd
             }
         }
 
+        // avg
+        template <class A, class T, class = typename std::enable_if<std::is_unsigned<T>::value, void>::type>
+        inline batch<T, A> avg(batch<T, A> const& self, batch<T, A> const& other, requires_arch<avx512bw>) noexcept
+        {
+            XSIMD_IF_CONSTEXPR(sizeof(T) == 1)
+            {
+                return _mm512_avg_epu8(self, other);
+            }
+            else XSIMD_IF_CONSTEXPR(sizeof(T) == 2)
+            {
+                return _mm512_avg_epu16(self, other);
+            }
+            else {
+              return avg(self, other, generic{});
+            }
+        }
+
         // bitwise_lshift
         template <class A, class T, class = typename std::enable_if<std::is_integral<T>::value, void>::type>
         inline batch<T, A> bitwise_lshift(batch<T, A> const& self, int32_t other, requires_arch<avx512bw>) noexcept
diff --git a/include/xsimd/arch/xsimd_sse2.hpp b/include/xsimd/arch/xsimd_sse2.hpp
index 0a34cb1e9..0eb4b8809 100644
--- a/include/xsimd/arch/xsimd_sse2.hpp
+++ b/include/xsimd/arch/xsimd_sse2.hpp
@@ -60,6 +60,9 @@ namespace xsimd
         inline batch<T, A> insert(batch<T, A> const& self, T val, index<I>, requires_arch<generic>) noexcept;
         template <class A, typename T, typename ITy, ITy... Indices>
         inline batch<T, A> shuffle(batch<T, A> const& x, batch<T, A> const& y, batch_constant<batch<ITy, A>, Indices...>, requires_arch<generic>) noexcept;
+        template <class A, class T>
+        inline batch<T, A> avg(batch<T, A> const& , batch<T, A> const&, requires_arch<generic>) noexcept;
+
 
         // abs
         template <class A>
@@ -148,6 +151,23 @@ namespace xsimd
             return _mm_movemask_epi8(self) != 0;
         }
 
+        // avg
+        template <class A, class T, class = typename std::enable_if<std::is_unsigned<T>::value, void>::type>
+        inline batch<T, A> avg(batch<T, A> const& self, batch<T, A> const& other, requires_arch<sse2>) noexcept
+        {
+            XSIMD_IF_CONSTEXPR(sizeof(T) == 1)
+            {
+                return _mm_avg_epu8(self, other);
+            }
+            else XSIMD_IF_CONSTEXPR(sizeof(T) == 2)
+            {
+                return _mm_avg_epu16(self, other);
+            }
+            else {
+              return avg(self, other, generic{});
+            }
+        }
+
         // batch_bool_cast
         template <class A, class T_out, class T_in>
         inline batch_bool<T_out, A> batch_bool_cast(batch_bool<T_in, A> const& self, batch_bool<T_out, A> const&, requires_arch<sse2>) noexcept
diff --git a/include/xsimd/types/xsimd_api.hpp b/include/xsimd/types/xsimd_api.hpp
index 0420f0a09..0bedf5fb4 100644
--- a/include/xsimd/types/xsimd_api.hpp
+++ b/include/xsimd/types/xsimd_api.hpp
@@ -202,6 +202,21 @@ namespace xsimd
         return kernel::atanh<A>(x, A {});
     }
 
+    /**
+     * @ingroup batch_math
+     *
+     * Computes the average of batches \c x and \c y
+     * @param x batch of T
+     * @param y batch of T
+     * @return the average of elements between \c x and \c y.
+     */
+    template <class T, class A>
+    inline batch<T, A> avg(batch<T, A> const& x, batch<T, A> const& y) noexcept
+    {
+        detail::static_check_supported_config<T, A>();
+        return kernel::avg<A>(x, y, A {});
+    }
+
     /**
      * @ingroup batch_conversion
      *
diff --git a/test/test_batch.cpp b/test/test_batch.cpp
index 9c3217361..ab058f08d 100644
--- a/test/test_batch.cpp
+++ b/test/test_batch.cpp
@@ -737,6 +737,25 @@ struct batch_test
         }
     }
 
+    void test_avg() const
+    {
+        array_type expected;
+        std::transform(lhs.cbegin(), lhs.cend(), rhs.cbegin(), expected.begin(),
+                       [](const value_type &l, const value_type &r) -> value_type {
+                         if (std::is_integral<value_type>::value) {
+                           if (std::is_signed<value_type>::value)
+                             return ((long long)l + r) / 2;
+                           else
+                             return (l + r + 1) / 2;
+                         } else {
+                           return (l + r) / 2;
+                         }
+                       });
+        batch_type res = avg(batch_lhs(), batch_rhs());
+        INFO("avg");
+        CHECK_BATCH_EQ(res, expected);
+    }
+
     void test_horizontal_operations() const
     {
         // reduce_add
@@ -938,6 +957,11 @@ TEST_CASE_TEMPLATE("[batch]", B, BATCH_TYPES)
         Test.test_abs();
     }
 
+    SUBCASE("avg")
+    {
+        Test.test_avg();
+    }
+
     SUBCASE("horizontal_operations")
     {
         Test.test_horizontal_operations();
diff --git a/test/test_xsimd_api.cpp b/test/test_xsimd_api.cpp
index 440ef015c..f6215d656 100644
--- a/test/test_xsimd_api.cpp
+++ b/test/test_xsimd_api.cpp
@@ -1156,6 +1156,13 @@ struct xsimd_api_all_types_functions
         CHECK_EQ(extract(xsimd::add(T(val0), T(val1))), val0 + val1);
     }
 
+    void test_avg()
+    {
+        value_type val0(1);
+        value_type val1(3);
+        CHECK_EQ(extract(xsimd::avg(T(val0), T(val1))), (val0 + val1) / value_type(2));
+    }
+
     void test_decr()
     {
         value_type val0(1);