Better performance for numpy_iexpr update operator

serge-sans-paille · May 31, 2017 · 2b151e8 · 2b151e8
1 parent 0f5f10c
commit 2b151e8
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 12 deletions.
diff --git a/pythran/pythonic/include/types/numpy_broadcast.hpp b/pythran/pythonic/include/types/numpy_broadcast.hpp
@@ -218,6 +218,10 @@ namespace pythonic
       static constexpr size_t value = 0;
 
       broadcast_base<dtype, is_vectorizable> _base;
+      operator dtype() const
+      {
+        return _base._value;
+      }
 
       broadcast() = default;
       template <class V>

diff --git a/pythran/pythonic/include/types/numpy_iexpr.hpp b/pythran/pythonic/include/types/numpy_iexpr.hpp
@@ -76,6 +76,9 @@ namespace pythonic
       numpy_iexpr &operator=(E const &expr);
       numpy_iexpr &operator=(numpy_iexpr const &expr);
 
+      template <class Op, class E>
+      numpy_iexpr &update_(E const &expr);
+
       template <class E>
       numpy_iexpr &operator+=(E const &expr);
       numpy_iexpr &operator+=(numpy_iexpr const &expr);

diff --git a/pythran/pythonic/types/numpy_iexpr.hpp b/pythran/pythonic/types/numpy_iexpr.hpp
@@ -6,6 +6,13 @@
 #include "pythonic/types/nditerator.hpp"
 #include "pythonic/types/tuple.hpp"
 
+#include "pythonic/operator_/iadd.hpp"
+#include "pythonic/operator_/iand.hpp"
+#include "pythonic/operator_/idiv.hpp"
+#include "pythonic/operator_/imul.hpp"
+#include "pythonic/operator_/ior.hpp"
+#include "pythonic/operator_/isub.hpp"
+
 #include <numeric>
 
 namespace pythonic
@@ -70,81 +77,100 @@ namespace pythonic
                                    false /*NIY*/>(*this, expr);
     }
 
+    template <class Arg>
+    template <class Op, class Expr>
+    numpy_iexpr<Arg> &numpy_iexpr<Arg>::update_(Expr const &expr)
+    {
+      using BExpr =
+          typename std::conditional<std::is_scalar<Expr>::value,
+                                    broadcast<Expr, dtype>, Expr const &>::type;
+      BExpr bexpr = expr;
+      utils::broadcast_update<
+          Op, numpy_iexpr &, BExpr, value,
+          value - (std::is_scalar<Expr>::value + utils::dim_of<Expr>::value),
+          is_vectorizable and
+              std::remove_reference<BExpr>::type::is_vectorizable and
+              std::is_same<dtype,
+                           typename std::decay<BExpr>::type::dtype>::value>(
+          *this, bexpr);
+      return *this;
+    }
+
     template <class Arg>
     template <class E>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator+=(E const &expr)
     {
-      return (*this) = (*this) + expr;
+      return update_<pythonic::operator_::functor::iadd>(expr);
     }
 
     template <class Arg>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator+=(numpy_iexpr<Arg> const &expr)
     {
-      return (*this) = (*this) + expr;
+      return update_<pythonic::operator_::functor::iadd>(expr);
     }
 
     template <class Arg>
     template <class E>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator-=(E const &expr)
     {
-      return (*this) = (*this) - expr;
+      return update_<pythonic::operator_::functor::isub>(expr);
     }
 
     template <class Arg>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator-=(numpy_iexpr<Arg> const &expr)
     {
-      return (*this) = (*this) - expr;
+      return update_<pythonic::operator_::functor::isub>(expr);
     }
 
     template <class Arg>
     template <class E>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator*=(E const &expr)
     {
-      return (*this) = (*this) * expr;
+      return update_<pythonic::operator_::functor::imul>(expr);
     }
 
     template <class Arg>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator*=(numpy_iexpr<Arg> const &expr)
     {
-      return (*this) = (*this) * expr;
+      return update_<pythonic::operator_::functor::imul>(expr);
     }
 
     template <class Arg>
     template <class E>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator/=(E const &expr)
     {
-      return (*this) = (*this) / expr;
+      return update_<pythonic::operator_::functor::idiv>(expr);
     }
 
     template <class Arg>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator/=(numpy_iexpr<Arg> const &expr)
     {
-      return (*this) = (*this) / expr;
+      return update_<pythonic::operator_::functor::idiv>(expr);
     }
 
     template <class Arg>
     template <class E>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator&=(E const &expr)
     {
-      return (*this) = (*this) & expr;
+      return update_<pythonic::operator_::functor::iand>(expr);
     }
 
     template <class Arg>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator&=(numpy_iexpr<Arg> const &expr)
     {
-      return (*this) = (*this) & expr;
+      return update_<pythonic::operator_::functor::iand>(expr);
     }
     template <class Arg>
     template <class E>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator|=(E const &expr)
     {
-      return (*this) = (*this) | expr;
+      return update_<pythonic::operator_::functor::ior>(expr);
     }
 
     template <class Arg>
     numpy_iexpr<Arg> &numpy_iexpr<Arg>::operator|=(numpy_iexpr<Arg> const &expr)
     {
-      return (*this) = (*this) | expr;
+      return update_<pythonic::operator_::functor::ior>(expr);
     }
 
     template <class Arg>