Avoid extra allocation when evaluation complex gexpr

Improve the simple runtime alias analysis to prevent some memory copies, based on the origin of each subexpression.
serge-sans-paille · Nov 28, 2023 · f65953b · f65953b
1 parent 0528c2e
commit f65953b
Show file tree

Hide file tree

Showing 5 changed files with 90 additions and 16 deletions.
diff --git a/pythran/pythonic/include/types/ndarray.hpp b/pythran/pythonic/include/types/ndarray.hpp
@@ -677,6 +677,10 @@ namespace types
     ndarray<T, pshape<long>> flat() const;
     ndarray<T, pS> copy() const;
     intptr_t id() const;
+    intptr_t baseid() const
+    {
+      return id();
+    }
     template <size_t I>
     auto shape() const -> decltype(std::get<I>(_shape))
     {

diff --git a/pythran/pythonic/include/types/numpy_gexpr.hpp b/pythran/pythonic/include/types/numpy_gexpr.hpp
@@ -921,6 +921,11 @@ namespace types
       return {*this};
     }
 
+    intptr_t baseid() const
+    {
+      return arg.baseid();
+    }
+
     template <class Tp, size_t... Is>
     auto recast(utils::index_sequence<Is...>) -> decltype(make_gexpr(
         arg.template recast<Tp>(),

diff --git a/pythran/pythonic/include/types/numpy_iexpr.hpp b/pythran/pythonic/include/types/numpy_iexpr.hpp
@@ -306,6 +306,11 @@ namespace types
       return {*this};
     }
 
+    intptr_t baseid() const
+    {
+      return arg.baseid();
+    }
+
     template <typename Tp>
     auto recast() -> decltype(arg.template recast<Tp>().fast(0))
     {

diff --git a/pythran/pythonic/include/types/numpy_texpr.hpp b/pythran/pythonic/include/types/numpy_texpr.hpp
@@ -233,6 +233,10 @@ namespace types
     explicit operator bool() const;
     long flat_size() const;
     intptr_t id() const;
+    intptr_t baseid() const
+    {
+      return arg.baseid();
+    }
     template <class Expr>
     numpy_texpr_2 &operator=(Expr const &expr);
     template <class Expr>

diff --git a/pythran/pythonic/types/numpy_gexpr.hpp b/pythran/pythonic/types/numpy_gexpr.hpp
@@ -16,6 +16,15 @@
 
 PYTHONIC_NS_BEGIN
 
+namespace numpy
+{
+  namespace details
+  {
+    template <class T>
+    struct arange_index;
+  }
+} // namespace numpy
+
 namespace types
 {
   template <class S0, class S1>
@@ -36,37 +45,84 @@ namespace types
   {
     return s.lower <= i && i < s.upper;
   }
+
+  template <class T>
+  constexpr intptr_t baseid_helper(T const &, std::true_type)
+  {
+    return 0;
+  }
+
+  template <class T>
+  intptr_t baseid_helper(T const &e, std::false_type)
+  {
+    return e.baseid();
+  }
+
+  template <class T>
+  intptr_t baseid(T const &e)
+  {
+    return baseid_helper(
+        e, std::integral_constant<bool, types::is_dtype<T>::value>{});
+  }
+
   template <class E0, class E1>
-  bool may_overlap(E0 const &, E1 const &)
+  bool may_overlap(E0 const &e0, E1 const &e1)
   {
-    return true;
+    return baseid(e0) == baseid(e1);
   }
-  template <class E0, class T0, class T1>
-  bool may_overlap(E0 const &, broadcast<T0, T1> const &)
+
+  template <class E0, class T1>
+  bool may_overlap(E0 const &e0, types::list<T1> const &e1)
   {
     return false;
   }
-  template <class Arg, class E1, class... S>
-  typename std::enable_if<std::is_scalar<E1>::value, bool>::type
-  may_overlap(numpy_gexpr<Arg, S...> const &gexpr, E1 const &)
+
+  template <class E0, class T1>
+  bool may_overlap(E0 const &e0, types::broadcasted<T1> const &e1)
+  {
+    return may_overlap(e0, e1.ref);
+  }
+
+  template <class E0, class T1, size_t N, class S>
+  bool may_overlap(E0 const &e0, types::array_base<T1, N, S> const &e1)
   {
     return false;
   }
-  template <class Arg, class Tuple, class... S, size_t... I>
-  bool may_overlap_helper(numpy_gexpr<Arg, S...> const &gexpr,
-                          Tuple const &args, utils::index_sequence<I...>)
+
+  template <class E0, class Op, class... Args, size_t... Is>
+  bool may_overlap(E0 const &e0, types::numpy_expr<Op, Args...> const &e1,
+                   utils::index_sequence<Is...>)
   {
-    bool overlaps[] = {may_overlap(gexpr, std::get<I>(args))...};
+    bool overlaps[] = {may_overlap(e0, std::get<Is>(e1.args))...};
     return std::any_of(std::begin(overlaps), std::end(overlaps),
                        [](bool b) { return b; });
   }
-  template <class Arg, class... E, class... S>
-  bool may_overlap(numpy_gexpr<Arg, S...> const &gexpr,
-                   numpy_expr<E...> const &expr)
+
+  template <class E0, class Op, class... Args>
+  bool may_overlap(E0 const &e0, types::numpy_expr<Op, Args...> const &e1)
+  {
+    return may_overlap(e0, e1, utils::make_index_sequence<sizeof...(Args)>());
+  }
+
+  template <class E0, class T1>
+  bool may_overlap(E0 const &e0,
+                   pythonic::numpy::details::arange_index<T1> const &e1)
+  {
+    return false;
+  }
+
+  template <class E0, class T0, class T1>
+  bool may_overlap(E0 const &, broadcast<T0, T1> const &)
+  {
+    return false;
+  }
+  template <class Arg, class E1, class... S>
+  typename std::enable_if<std::is_scalar<E1>::value, bool>::type
+  may_overlap(numpy_gexpr<Arg, S...> const &gexpr, E1 const &)
   {
-    return may_overlap_helper(gexpr, expr.args,
-                              utils::make_index_sequence<sizeof...(E) - 1>{});
+    return false;
   }
+
   template <class T, class pS, class Tp, class pSp, class E0, class E1>
   bool may_gexpr_overlap(E0 const &gexpr, E1 const &expr)
   {