Improve map_unary tail

xtensor-stack · AntoinePrv · Sep 29, 2026 · Sep 10, 2026 · Sep 10, 2026 · Sep 10, 2026
commit bb86eeecbb727504585ff004df3d7fccfe69cbaf
diff --git a/include/xsimd_algorithm/builder.hpp b/include/xsimd_algorithm/builder.hpp
@@ -23,49 +23,55 @@
 
 namespace xsimd::builder
 {
-    struct alignment_options
+    struct alignment
     {
         bool start_aligned = false;
         bool end_aligned = false;
     };
 
     template <typename T>
-    auto prev_aligned(T* ptr, std::size_t alignment) -> T*
+    XSIMD_INLINE auto prev_aligned(T* ptr, std::size_t alignment) -> T*
     {
         assert(std::has_single_bit(alignment));
         auto const address = reinterpret_cast<std::size_t>(ptr);
         return reinterpret_cast<T*>(address & ~(alignment - 1));
     }
 
     template <typename T>
-    auto next_aligned(T* ptr, std::size_t alignment) -> T*
+    XSIMD_INLINE auto next_aligned(T* ptr, std::size_t alignment) -> T*
     {
         assert(std::has_single_bit(alignment));
         auto const address = reinterpret_cast<std::size_t>(ptr);
         return reinterpret_cast<T*>((address + alignment - 1) & ~(alignment - 1));
     }
 
     template <typename T>
-    auto bytes_to_next_aligned(T* ptr, std::size_t alignment) -> std::size_t
+    XSIMD_INLINE auto bytes_to_next_aligned(T* ptr, std::size_t alignment) -> std::size_t
     {
         assert(std::has_single_bit(alignment));
         auto const address = reinterpret_cast<std::uintptr_t>(ptr);
         return (alignment - (address & (alignment - 1))) & (alignment - 1);
     }
 
     template <typename T, typename U>
-    auto are_aliased(std::span<T> lhs, std::span<U> rhs) -> bool
+    XSIMD_INLINE auto are_aliased(std::span<T> lhs, std::span<U> rhs) -> bool
     {
         // Comparing pointers from unrelated objects is unspecified, integers are not.
         auto const lhs_begin = reinterpret_cast<std::uintptr_t>(lhs.data());
         auto const rhs_begin = reinterpret_cast<std::uintptr_t>(rhs.data());
         return (lhs_begin < rhs_begin + rhs.size_bytes()) && (rhs_begin < lhs_begin + lhs.size_bytes());
     }
 
+    struct unary_options
+    {
+        std::size_t unroll_factor = 4;
+        bool pure = false;
+    };
+
     template <
         typename Arch = xsimd::default_arch,
         typename T, typename U, typename Func>
-    void map_unary_batch(
+    XSIMD_INLINE void map_unary_batch(
         T const* XSIMD_RESTRICT begin,
         T const* XSIMD_RESTRICT end,
         U* XSIMD_RESTRICT out,
@@ -91,11 +97,38 @@ namespace xsimd::builder
         std::memcpy(out, output_buffer, in_count * sizeof(T));
     }
 
+    template <typename T, typename A, bool aligned>
+    XSIMD_INLINE xsimd::batch<T, A> load_batch(T const* ptr)
+    {
+        if constexpr (aligned)
+        {
+            return xsimd::batch<T, A>::load_aligned(ptr);
+        }
+        else
+        {
+            return xsimd::batch<T, A>::load_unaligned(ptr);
+        }
+    }
+
+    template <typename T, typename A, bool aligned>
+    XSIMD_INLINE void store_batch(xsimd::batch<T, A> x, T* ptr)
+    {
+        if constexpr (aligned)
+        {
+            x.store_aligned(ptr);
+        }
+        else
+        {
+            x.store_unaligned(ptr);
+        }
+    }
+
     template <
-        alignment_options aligned = {},
+        alignment align = {},
+        unary_options opts = {},
         typename Arch = xsimd::default_arch,
         typename T, typename U, typename Func>
-    void map_unary(std::span<T const> in, std::span<U> out, Func&& func)
+    XSIMD_INLINE void map_unary(std::span<T const> in, std::span<U> out, Func&& func)
     {
         using input_batch = xsimd::batch<T, Arch>;
         using output_batch = xsimd::batch<U, Arch>;
@@ -104,8 +137,8 @@ namespace xsimd::builder
         // If input is not guarenteed aligned, we will try to align preferably the
         // output (more expensive unaligned stores) or otherwise the input.
         constexpr bool align_output = sizeof(U) >= sizeof(T);
-        constexpr bool load_is_aligned = aligned.start_aligned || !align_output;
-        constexpr bool store_is_aligned = aligned.start_aligned || align_output;
+        constexpr bool load_is_aligned = align.start_aligned || !align_output;
+        constexpr bool store_is_aligned = align.start_aligned || align_output;
 
         assert(in.size() * sizeof(T) == out.size() * sizeof(U));
         assert(!are_aliased(in, out));
@@ -117,11 +150,11 @@ namespace xsimd::builder
 
         auto ot = out.data();
         auto it = in.data();
-        auto const end = in.data() + in.size();
+        auto const iend = in.data() + in.size();
 
         // Input and output may not have the same alignment so it may be impossible
         // to get both aligned, so we align a single side.
-        if constexpr (!aligned.start_aligned)
+        if constexpr (!align.start_aligned)
         {
             // The span may be too short to reach the next alignment boundary.
             const auto head_bytes = std::min(
@@ -136,39 +169,49 @@ namespace xsimd::builder
             ot += head_bytes / sizeof(U);
         }
 
-        // No loop-carried dependencies and no aliasing, so we leave the compiler
-        // to unroll the loop.
-        while (static_cast<std::size_t>(end - it) >= input_batch::size)
+        // Unrolled loop processing multiple batches at a time
+        while (static_cast<std::size_t>(iend - it) >= opts.unroll_factor * input_batch::size)
         {
-            input_batch x;
-            if constexpr (load_is_aligned)
+            input_batch x[opts.unroll_factor];
+            for (std::size_t u = 0; u < opts.unroll_factor; ++u)
             {
-                x = input_batch::load_aligned(it);
+                x[u] = load_batch<T, Arch, load_is_aligned>(it + u * input_batch::size);
             }
-            else
+            for (std::size_t u = 0; u < opts.unroll_factor; ++u)
             {
-                x = input_batch::load_unaligned(it);
+                store_batch<U, Arch, store_is_aligned>(func(x[u]), ot + u * output_batch::size);
             }
 
-            const auto y = func(x);
-            if constexpr (store_is_aligned)
-            {
-                y.store_aligned(ot);
-            }
-            else
-            {
-                y.store_unaligned(ot);
-            }
+            it += opts.unroll_factor * input_batch::size;
+            ot += opts.unroll_factor * output_batch::size;
+        }
 
+        while (static_cast<std::size_t>(iend - it) >= input_batch::size)
+        {
+            const auto x = load_batch<T, Arch, load_is_aligned>(it);
+            store_batch<U, Arch, store_is_aligned>(func(x), ot);
             it += input_batch::size;
             ot += output_batch::size;
         }
 
         // Unlikely to be skipped, meant for users that know they allocate
-        // a multiple of the batch size.
-        if constexpr (!aligned.end_aligned)
+        // a multiple of the batch size, such as in a local buffer
+        if constexpr (!align.end_aligned)
         {
-            map_unary_batch<Arch>(it, end, ot, func);
+            auto const oend = out.data() + out.size();
+            // Stepping back shifts the batch boundary, which pairs lanes differently
+            // than starting from the front unless both sides have the same lane count.
+            constexpr bool can_step_back = input_batch::size == output_batch::size;
+            if (can_step_back && it != iend && (in.size() >= input_batch::size)) [[likely]]
+            {
+                // Recompute overlapping data, this time starting from the end.
+                const auto x = load_batch<T, Arch, false>(iend - input_batch::size);
+                store_batch<U, Arch, false>(func(x), oend - output_batch::size);
+            }
+            else
+            {
+                map_unary_batch<Arch>(it, iend, ot, func);
+            }
         }
     }
 }

diff --git a/include/xsimd_algorithm/math.hpp b/include/xsimd_algorithm/math.hpp
@@ -16,23 +16,25 @@
 namespace xsimd::algo
 {
     template <
-        xsimd::builder::alignment_options aligned = {},
+        xsimd::builder::alignment align = {},
         typename Arch = xsimd::default_arch,
         typename T>
     void sqrt(std::span<T const> in, std::span<T> out)
     {
-        return xsimd::builder::map_unary<aligned, Arch>(
+        constexpr builder::unary_options opts = { .unroll_factor = 4, .pure = true };
+        return xsimd::builder::map_unary<align, opts, Arch>(
             in, out, [](auto x)
             { return sqrt(x); });
     }
 
     template <
-        xsimd::builder::alignment_options aligned = {},
+        xsimd::builder::alignment align = {},
         typename Arch = xsimd::default_arch,
         typename T>
     void abs(std::span<T const> in, std::span<T> out)
     {
-        return xsimd::builder::map_unary<aligned, Arch>(
+        constexpr builder::unary_options opts = { .unroll_factor = 4, .pure = true };
+        return xsimd::builder::map_unary<align, opts, Arch>(
             in, out, [](auto x)
             { return abs(x); });
     }

diff --git a/test-utils/include/xsimd_test_utils/math_data.hpp b/test-utils/include/xsimd_test_utils/math_data.hpp
@@ -59,7 +59,7 @@ namespace xsimd::test
             return std::sqrt(x);
         }
 
-        template <xsimd::builder::alignment_options aligned = {}>
+        template <xsimd::builder::alignment aligned = {}>
         static void apply_range_simd(std::span<T const> in, std::span<T> out)
         {
             xsimd::algo::sqrt<aligned>(in, out);
@@ -82,7 +82,7 @@ namespace xsimd::test
             return std::abs(x);
         }
 
-        template <xsimd::builder::alignment_options aligned = {}>
+        template <xsimd::builder::alignment aligned = {}>
         static void apply_range_simd(std::span<T const> in, std::span<T> out)
         {
             xsimd::algo::abs<aligned>(in, out);

diff --git a/test-utils/include/xsimd_test_utils/math_ops.hpp b/test-utils/include/xsimd_test_utils/math_ops.hpp
@@ -0,0 +1,75 @@
+/****************************************************************************
+ * Copyright (c) xsimd-algorithm contributors                               *
+ *                                                                          *
+ * Distributed under the terms of the BSD 3-Clause License.                 *
+ *                                                                          *
+ * The full license is in the file LICENSE, distributed with this software. *
+ ****************************************************************************/
+
+#ifndef XSIMD_ALGORITHM_TEST_UTILS_MATH_OPS_HPP
+#define XSIMD_ALGORITHM_TEST_UTILS_MATH_OPS_HPP
+
+#include <cmath>
+#include <cstddef>
+#include <span>
+#include <vector>
+
+#include <xsimd_algorithm/builder.hpp>
+#include <xsimd_algorithm/math.hpp>
+
+#include "xsimd_test_utils/utils.hpp"
+
+namespace xsimd::test
+{
+    template <typename T>
+    struct sqrt_op
+    {
+        using value_type = T;
+
+        static constexpr auto name = "sqrt";
+
+        template <xsimd::builder::alignment aligned = {}>
+        static void apply(std::span<T const> in, std::span<T> out)
+        {
+            xsimd::algo::sqrt<aligned>(in, out);
+        }
+
+        static T scalar(T x)
+        {
+            return std::sqrt(x);
+        }
+
+        template <typename Alloc>
+        static std::vector<T, Alloc> input(std::size_t size)
+        {
+            return xsimd::test::make_arange<T, Alloc>(size);
+        }
+    };
+
+    template <typename T>
+    struct abs_op
+    {
+        using value_type = T;
+
+        static constexpr auto name = "abs";
+
+        template <xsimd::builder::alignment aligned = {}>
+        static void apply(std::span<T const> in, std::span<T> out)
+        {
+            xsimd::algo::abs<aligned>(in, out);
+        }
+
+        static T scalar(T x)
+        {
+            return std::abs(x);
+        }
+
+        template <typename Alloc>
+        static std::vector<T, Alloc> input(std::size_t size)
+        {
+            return xsimd::test::make_arange<T, Alloc>(size, -static_cast<T>(size) / 2);
+        }
+    };
+}
+
+#endif