Implements in-place addition

Author: ndgrigorian

dpctl/tensor/_elementwise_common.py

@@ -31,6 +31,7 @@
     _empty_like_pair_orderK,
     _find_buf_dtype,
     _find_buf_dtype2,
+    _find_inplace_dtype,
     _to_device_supported_dtype,
 )
 
@@ -331,11 +332,19 @@ class BinaryElementwiseFunc:
     Class that implements binary element-wise functions.
     """
 
-    def __init__(self, name, result_type_resolver_fn, binary_dp_impl_fn, docs):
+    def __init__(
+        self,
+        name,
+        result_type_resolver_fn,
+        binary_dp_impl_fn,
+        docs,
+        binary_inplace_fn=None,
+    ):
         self.__name__ = "BinaryElementwiseFunc"
         self.name_ = name
         self.result_type_resolver_fn_ = result_type_resolver_fn
         self.binary_fn_ = binary_dp_impl_fn
+        self.binary_inplace_fn_ = binary_inplace_fn
         self.__doc__ = docs
 
     def __str__(self):
@@ -631,3 +640,113 @@ def __call__(self, o1, o2, out=None, order="K"):
         )
         dpctl.SyclEvent.wait_for([ht_copy1_ev, ht_copy2_ev, ht_])
         return out
+
+    def inplace(self, lhs, val):
+        if self.binary_inplace_fn_ is None:
+            raise ValueError(
+                f"In-place operation not supported for ufunc '{self.name_}'"
+            )
+        if not isinstance(lhs, dpt.usm_ndarray):
+            raise TypeError(
+                f"Expected dpctl.tensor.usm_ndarray, got {type(lhs)}"
+            )
+        q1, lhs_usm_type = _get_queue_usm_type(lhs)
+        q2, val_usm_type = _get_queue_usm_type(val)
+        if q2 is None:
+            exec_q = q1
+            usm_type = lhs_usm_type
+        else:
+            exec_q = dpctl.utils.get_execution_queue((q1, q2))
+            if exec_q is None:
+                raise ExecutionPlacementError(
+                    "Execution placement can not be unambiguously inferred "
+                    "from input arguments."
+                )
+            usm_type = dpctl.utils.get_coerced_usm_type(
+                (
+                    lhs_usm_type,
+                    val_usm_type,
+                )
+            )
+        dpctl.utils.validate_usm_type(usm_type, allow_none=False)
+        lhs_shape = _get_shape(lhs)
+        val_shape = _get_shape(val)
+        if not all(
+            isinstance(s, (tuple, list))
+            for s in (
+                lhs_shape,
+                val_shape,
+            )
+        ):
+            raise TypeError(
+                "Shape of arguments can not be inferred. "
+                "Arguments are expected to be "
+            )
+        try:
+            res_shape = _broadcast_shape_impl(
+                [
+                    lhs_shape,
+                    val_shape,
+                ]
+            )
+        except ValueError:
+            raise ValueError(
+                "operands could not be broadcast together with shapes "
+                f"{lhs_shape} and {val_shape}"
+            )
+        if res_shape != lhs_shape:
+            raise ValueError(
+                f"output shape {lhs_shape} does not match "
+                f"broadcast shape {res_shape}"
+            )
+        sycl_dev = exec_q.sycl_device
+        lhs_dtype = lhs.dtype
+        val_dtype = _get_dtype(val, sycl_dev)
+        if not _validate_dtype(val_dtype):
+            raise ValueError("Input operand of unsupported type")
+
+        lhs_dtype, val_dtype = _resolve_weak_types(
+            lhs_dtype, val_dtype, sycl_dev
+        )
+
+        buf_dt = _find_inplace_dtype(
+            lhs_dtype, val_dtype, self.result_type_resolver_fn_, sycl_dev
+        )
+
+        if buf_dt is None:
+            raise TypeError(
+                f"function '{self.name_}' does not support input types "
+                f"({lhs_dtype}, {val_dtype}), "
+                "and the inputs could not be safely coerced to any "
+                "supported types according to the casting rule ''safe''."
+            )
+
+        if isinstance(val, dpt.usm_ndarray):
+            rhs = val
+        else:
+            rhs = dpt.asarray(val, dtype=val_dtype, sycl_queue=exec_q)
+
+        if buf_dt == val_dtype:
+            rhs = dpt.broadcast_to(rhs, res_shape)
+            ht_, _ = self.binary_inplace_fn_(
+                lhs=lhs, rhs=rhs, sycl_queue=exec_q
+            )
+            ht_.wait()
+
+        else:
+            buf = dpt.empty_like(rhs, dtype=buf_dt)
+            ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
+                src=rhs, dst=buf, sycl_queue=exec_q
+            )
+
+            buf = dpt.broadcast_to(buf, res_shape)
+            ht_, _ = self.binary_inplace_fn_(
+                lhs=lhs,
+                rhs=buf,
+                sycl_queue=exec_q,
+                depends=[copy_ev],
+            )
+            ht_copy_ev.wait()
+            ht_.wait()
+
+        return lhs

dpctl/tensor/_elementwise_funcs.py

@@ -73,7 +73,11 @@
         returned array is determined by the Type Promotion Rules.
 """
 add = BinaryElementwiseFunc(
-    "add", ti._add_result_type, ti._add, _add_docstring_
+    "add",
+    ti._add_result_type,
+    ti._add,
+    _add_docstring_,
+    binary_inplace_fn=ti._add_inplace,
 )
 
 # U04: ===== ASIN  (x)

dpctl/tensor/_type_utils.py

@@ -294,9 +294,27 @@ def _find_buf_dtype2(arg1_dtype, arg2_dtype, query_fn, sycl_dev):
     return None, None, None
 
 
+def _find_inplace_dtype(lhs_dtype, rhs_dtype, query_fn, sycl_dev):
+    res_dt = query_fn(lhs_dtype, rhs_dtype)
+    if res_dt and res_dt == lhs_dtype:
+        return rhs_dtype
+
+    _fp16 = sycl_dev.has_aspect_fp16
+    _fp64 = sycl_dev.has_aspect_fp64
+    all_dts = _all_data_types(_fp16, _fp64)
+    for buf_dt in all_dts:
+        if _can_cast(rhs_dtype, buf_dt, _fp16, _fp64):
+            res_dt = query_fn(lhs_dtype, buf_dt)
+            if res_dt and res_dt == lhs_dtype:
+                return buf_dt
+
+    return None
+
+
 __all__ = [
     "_find_buf_dtype",
     "_find_buf_dtype2",
+    "_find_inplace_dtype",
     "_empty_like_orderK",
     "_empty_like_pair_orderK",
     "_to_device_supported_dtype",

dpctl/tensor/_usmarray.pyx

@@ -1245,11 +1245,8 @@ cdef class usm_ndarray:
         return _dispatch_binary_elementwise2(other, "logical_xor", self)
 
     def __iadd__(self, other):
-        res = self.__add__(other)
-        if res is NotImplemented:
-            return res
-        self.__setitem__(Ellipsis, res)
-        return self
+        from ._elementwise_funcs import add
+        return add.inplace(self, other)
 
     def __iand__(self, other):
         res = self.__and__(other)

dpctl/tensor/libtensor/include/kernels/elementwise_functions/add.hpp

@@ -34,6 +34,7 @@
 #include "utils/type_utils.hpp"
 
 #include "kernels/elementwise_functions/common.hpp"
+#include "kernels/elementwise_functions/common_inplace.hpp"
 #include <pybind11/pybind11.h>
 
 namespace dpctl
@@ -212,7 +213,6 @@ template <typename fnT, typename T1, typename T2> struct AddTypeMapFactory
     std::enable_if_t<std::is_same<fnT, int>::value, int> get()
     {
         using rT = typename AddOutputType<T1, T2>::value_type;
-        ;
         return td_ns::GetTypeid<rT>{}.get();
     }
 };
@@ -364,6 +364,122 @@ struct AddContigRowContigMatrixBroadcastFactory
     }
 };
 
+template <typename argT, typename resT> struct AddInplaceFunctor
+{
+
+    using supports_sg_loadstore = std::negation<
+        std::disjunction<tu_ns::is_complex<argT>, tu_ns::is_complex<resT>>>;
+    using supports_vec = std::negation<
+        std::disjunction<tu_ns::is_complex<argT>, tu_ns::is_complex<resT>>>;
+
+    void operator()(resT &res, const argT &in)
+    {
+        res += in;
+    }
+
+    template <int vec_sz>
+    void operator()(sycl::vec<resT, vec_sz> &res,
+                    const sycl::vec<argT, vec_sz> &in)
+    {
+        res += in;
+    }
+};
+
+template <typename argT,
+          typename resT,
+          unsigned int vec_sz = 4,
+          unsigned int n_vecs = 2>
+using AddInplaceContigFunctor = elementwise_common::BinaryInplaceContigFunctor<
+    argT,
+    resT,
+    AddInplaceFunctor<argT, resT>,
+    vec_sz,
+    n_vecs>;
+
+template <typename argT, typename resT, typename IndexerT>
+using AddInplaceStridedFunctor =
+    elementwise_common::BinaryInplaceStridedFunctor<
+        argT,
+        resT,
+        IndexerT,
+        AddInplaceFunctor<argT, resT>>;
+
+template <typename argT,
+          typename resT,
+          unsigned int vec_sz,
+          unsigned int n_vecs>
+class add_inplace_contig_kernel;
+
+template <typename argTy, typename resTy>
+sycl::event
+add_inplace_contig_impl(sycl::queue exec_q,
+                        size_t nelems,
+                        const char *arg_p,
+                        py::ssize_t arg_offset,
+                        char *res_p,
+                        py::ssize_t res_offset,
+                        const std::vector<sycl::event> &depends = {})
+{
+    return elementwise_common::binary_inplace_contig_impl<
+        argTy, resTy, AddInplaceContigFunctor, add_inplace_contig_kernel>(
+        exec_q, nelems, arg_p, arg_offset, res_p, res_offset, depends);
+}
+
+template <typename fnT, typename T1, typename T2> struct AddInplaceContigFactory
+{
+    fnT get()
+    {
+        if constexpr (std::is_same_v<typename AddOutputType<T1, T2>::value_type,
+                                     void>) {
+            fnT fn = nullptr;
+            return fn;
+        }
+        else {
+            fnT fn = add_inplace_contig_impl<T1, T2>;
+            return fn;
+        }
+    }
+};
+
+template <typename resT, typename argT, typename IndexerT>
+class add_inplace_strided_kernel;
+
+template <typename argTy, typename resTy>
+sycl::event
+add_inplace_strided_impl(sycl::queue exec_q,
+                         size_t nelems,
+                         int nd,
+                         const py::ssize_t *shape_and_strides,
+                         const char *arg_p,
+                         py::ssize_t arg_offset,
+                         char *res_p,
+                         py::ssize_t res_offset,
+                         const std::vector<sycl::event> &depends,
+                         const std::vector<sycl::event> &additional_depends)
+{
+    return elementwise_common::binary_inplace_strided_impl<
+        argTy, resTy, AddInplaceStridedFunctor, add_inplace_strided_kernel>(
+        exec_q, nelems, nd, shape_and_strides, arg_p, arg_offset, res_p,
+        res_offset, depends, additional_depends);
+}
+
+template <typename fnT, typename T1, typename T2>
+struct AddInplaceStridedFactory
+{
+    fnT get()
+    {
+        if constexpr (std::is_same_v<typename AddOutputType<T1, T2>::value_type,
+                                     void>) {
+            fnT fn = nullptr;
+            return fn;
+        }
+        else {
+            fnT fn = add_inplace_strided_impl<T1, T2>;
+            return fn;
+        }
+    }
+};
+
 } // namespace add
 } // namespace kernels
 } // namespace tensor