Avoid temporary arrays for independent numeric generators

Author: rithikasridharbabu22-sys

docs/source/builder.rst

@@ -26,6 +26,22 @@ Numerical ranges
 - :cpp:func:`xt::linspace(start, stop, num_samples) <xt::linspace>`: generates num_samples evenly spaced numbers over given interval.
 - :cpp:func:`xt::logspace(start, stop, num_samples) <xt::logspace>`: generates num_samples evenly spaced on a log scale over given interval
 
+Assignment of independent generators
+-----------------------------------
+
+Assigning ``arange`` or ``linspace`` with arithmetic types to a contiguous
+container reuses its storage when the size is unchanged, without creating a
+temporary array. Views retain their existing assignment path, which can benefit
+from generating into a contiguous temporary first.
+
+Custom generator functors conservatively retain overlap protection. A functor
+which cannot read array storage can opt in by specializing
+``xt::is_memory_independent<Functor>`` to inherit from ``std::true_type``.
+``xgenerator`` forwards this trait from its functor, including reference-qualified
+functors. Do not opt in for functors that read other expressions, since those
+expressions can refer to the assignment destination. The option
+``XTENSOR_FORCE_TEMPORARY_MEMORY_IN_ASSIGNMENTS`` still forces temporary evaluation.
+
 Joining expressions
 -------------------
 

include/xtensor/core/xsemantic.hpp

@@ -682,7 +682,20 @@ namespace xt
     template <class E>
     inline auto xcontainer_semantic<D>::operator=(const xexpression<E>& e) -> derived_type&
     {
-        return base_type::operator=(e);
+#ifndef XTENSOR_FORCE_TEMPORARY_MEMORY_IN_ASSIGNMENTS
+        // Keep this optimization on containers: a view may benefit from first
+        // generating into a contiguous temporary with a specialized assign_to.
+        if constexpr (is_memory_independent<E>::value && D::contiguous_layout)
+        {
+            // Validate the shape before entering a generator's noexcept assign_to.
+            this->derived_cast().resize(e.derived_cast().shape());
+            return this->assign(e);
+        }
+        else
+#endif
+        {
+            return base_type::operator=(e);
+        }
     }
 
     /*********************************

include/xtensor/generators/xbuilder.hpp

@@ -390,6 +390,12 @@ namespace xt
         };
     }
 
+    template <class T, class R, class S>
+    struct is_memory_independent<detail::arange_generator<T, R, S>>
+        : std::conjunction<std::is_arithmetic<T>, std::is_arithmetic<R>, std::is_arithmetic<S>>
+    {
+    };
+
     /**
      * Generates an array with ones on the diagonal.
      * @param shape shape of the resulting expression

include/xtensor/generators/xgenerator.hpp

@@ -58,6 +58,11 @@ namespace xt
     template <class F, class R, class S>
     class xgenerator;
 
+    template <class F, class R, class S>
+    struct is_memory_independent<xgenerator<F, R, S>> : is_memory_independent<std::decay_t<F>>
+    {
+    };
+
     template <typename T>
     concept xgenerator_concept = is_specialization_of<xgenerator, std::decay_t<T>>::value;
 

include/xtensor/utils/xutils.hpp

@@ -672,6 +672,16 @@ namespace xt
         return !(a == b);
     }
 
+    /**
+     * Whether evaluating a functor or expression is independent of array storage.
+     * Specialize only for types which cannot read from an assignment destination.
+     * Unknown types conservatively retain overlap protection.
+     */
+    template <class T>
+    struct is_memory_independent : std::false_type
+    {
+    };
+
     /*****************
      * assignable_to *
      ******************/

test/test_xassign.cpp

@@ -8,18 +8,66 @@
  ****************************************************************************/
 
 #include <iterator>
+#include <memory>
 #include <type_traits>
 #include <vector>
 
 #include "xtensor/containers/xarray.hpp"
 #include "xtensor/containers/xtensor.hpp"
 #include "xtensor/core/xassign.hpp"
 #include "xtensor/core/xnoalias.hpp"
+#include "xtensor/generators/xbuilder.hpp"
 #include "xtensor/misc/xmanipulation.hpp"
 
 #include "test_common.hpp"
 #include "test_common_macros.hpp"
 
+namespace
+{
+    template <class T>
+    struct counting_allocator : std::allocator<T>
+    {
+        using value_type = T;
+        inline static std::size_t allocations = 0;
+
+        template <class U>
+        struct rebind
+        {
+            using other = counting_allocator<U>;
+        };
+
+        counting_allocator() = default;
+
+        template <class U>
+        counting_allocator(const counting_allocator<U>&)
+        {
+        }
+
+        T* allocate(std::size_t n)
+        {
+            ++allocations;
+            return std::allocator<T>::allocate(n);
+        }
+    };
+
+    struct reverse_reader
+    {
+        using value_type = int;
+        const xt::xarray<int>& source;
+
+        int operator()(std::size_t i) const
+        {
+            return source(source.size() - 1 - i);
+        }
+
+        template <class It>
+        int element(It first, It) const
+        {
+            return (*this)(*first);
+        }
+    };
+}
+
 // a dummy shape *not derived* from std::vector but compatible
 template <class T>
 class my_vector
@@ -137,6 +185,57 @@ class my_vector
 namespace xt
 {
 
+#ifndef XTENSOR_FORCE_TEMPORARY_MEMORY_IN_ASSIGNMENTS
+    TEST(xassign, independent_generator_reuses_storage)
+    {
+        using allocator = counting_allocator<double>;
+        using array = xtensor<double, 1, layout_type::row_major, allocator>;
+        array a = array::from_shape({16});
+
+        allocator::allocations = 0;
+        a = arange<double>(16);
+        EXPECT_EQ(allocator::allocations, 0u);
+        for (std::size_t i = 0; i < a.size(); ++i)
+        {
+            EXPECT_EQ(a(i), double(i));
+        }
+
+        allocator::allocations = 0;
+        a = linspace<double>(0., 15., 16);
+        EXPECT_EQ(allocator::allocations, 0u);
+        for (std::size_t i = 0; i < a.size(); ++i)
+        {
+            EXPECT_EQ(a(i), double(i));
+        }
+    }
+#endif
+
+    TEST(xassign, independent_generator_resizes)
+    {
+        xarray<double> a = {{-1., -1.}, {-1., -1.}};
+        a = arange<double>(5);
+        EXPECT_EQ(a.dimension(), 1u);
+        EXPECT_EQ(a.size(), 5u);
+        EXPECT_EQ(a(4), 4.);
+
+        a = arange<double>(0);
+        EXPECT_EQ(a.size(), 0u);
+        a = linspace<double>(2., 4., 3);
+        EXPECT_EQ(a, (xarray<double>{2., 3., 4.}));
+    }
+
+    TEST(xassign, reading_generators_keep_overlap_protection)
+    {
+        xarray<int> a = {1, 2, 3, 4};
+        auto reverse = detail::make_xgenerator(reverse_reader{a}, {4});
+        static_assert(!is_memory_independent<decltype(reverse)>::value);
+        a = reverse;
+        EXPECT_EQ(a, (xarray<int>{4, 3, 2, 1}));
+
+        a = concatenate(xtuple(a, a));
+        EXPECT_EQ(a, (xarray<int>{4, 3, 2, 1, 4, 3, 2, 1}));
+    }
+
     template <class T, class C_T>
     struct rebind_container<T, my_vector<C_T>>
     {
@@ -173,5 +272,7 @@ namespace xt
     {
         auto a = xt::xtensor<float, 1>::from_shape({2});
         XT_ASSERT_THROW(a = xt::expand_dims(a, 0), std::runtime_error);
+        auto b = xt::xtensor<float, 2>::from_shape({2, 2});
+        XT_ASSERT_THROW(b = arange<float>(4), std::runtime_error);
     }
 }