In-place operations enabled in standard binary operators

- functionality such as binop(x, y, out=x) now possible, with some edge cases still WIP

Author: ndgrigorian

dpctl/tensor/_elementwise_common.py

@@ -354,6 +354,13 @@ def __repr__(self):
         return f"<BinaryElementwiseFunc '{self.name_}'>"
 
     def __call__(self, o1, o2, out=None, order="K"):
+        # FIXME: replace with check against base array
+        # when views can be identified
+        if o1 is out:
+            return self._inplace(o1, o2)
+        elif o2 is out:
+            return self._inplace(o2, o1)
+
         if order not in ["K", "C", "F", "A"]:
             order = "K"
         q1, o1_usm_type = _get_queue_usm_type(o1)
@@ -397,6 +404,7 @@ def __call__(self, o1, o2, out=None, order="K"):
             raise TypeError(
                 "Shape of arguments can not be inferred. "
                 "Arguments are expected to be "
+                "lists, tuples, or both"
             )
         try:
             res_shape = _broadcast_shape_impl(
@@ -424,7 +432,7 @@ def __call__(self, o1, o2, out=None, order="K"):
 
         if res_dt is None:
             raise TypeError(
-                "function 'add' does not support input types "
+                f"function '{self.name_}' does not support input types "
                 f"({o1_dtype}, {o2_dtype}), "
                 "and the inputs could not be safely coerced to any "
                 "supported types according to the casting rule ''safe''."
@@ -641,7 +649,7 @@ 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):
+    def _inplace(self, lhs, val):
         if self.binary_inplace_fn_ is None:
             raise ValueError(
                 f"In-place operation not supported for ufunc '{self.name_}'"
@@ -681,6 +689,7 @@ def inplace(self, lhs, val):
             raise TypeError(
                 "Shape of arguments can not be inferred. "
                 "Arguments are expected to be "
+                "lists, tuples, or both"
             )
         try:
             res_shape = _broadcast_shape_impl(
@@ -715,7 +724,7 @@ def inplace(self, lhs, val):
 
         if buf_dt is None:
             raise TypeError(
-                f"function '{self.name_}' does not support input types "
+                f"In-place '{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''."

dpctl/tensor/_usmarray.pyx

@@ -1246,7 +1246,7 @@ cdef class usm_ndarray:
 
     def __iadd__(self, other):
         from ._elementwise_funcs import add
-        return add.inplace(self, other)
+        return add._inplace(self, other)
 
     def __iand__(self, other):
         res = self.__and__(other)
@@ -1285,7 +1285,7 @@ cdef class usm_ndarray:
 
     def __imul__(self, other):
         from ._elementwise_funcs import multiply
-        return multiply.inplace(self, other)
+        return multiply._inplace(self, other)
 
     def __ior__(self, other):
         res = self.__or__(other)
@@ -1310,7 +1310,7 @@ cdef class usm_ndarray:
 
     def __isub__(self, other):
         from ._elementwise_funcs import subtract
-        return subtract.inplace(self, other)
+        return subtract._inplace(self, other)
 
     def __itruediv__(self, other):
         res = self.__truediv__(other)

dpctl/tests/elementwise/test_add.py

@@ -294,7 +294,8 @@ def test_add_errors():
 
     ar1 = dpt.ones(2, dtype="float32")
     ar2 = dpt.ones_like(ar1, dtype="int32")
-    y = ar1
+    # identical view but a different object
+    y = ar1[:]
     assert_raises_regex(
         TypeError,
         "Input and output arrays have memory overlap",
@@ -423,7 +424,7 @@ def test_add_inplace_errors():
     ar1 = np.ones(2, dtype="float32")
     ar2 = dpt.ones(2, dtype="float32")
     with pytest.raises(TypeError):
-        dpt.add.inplace(ar1, ar2)
+        ar1 += ar2
 
     ar1 = dpt.ones(2, dtype="float32")
     ar2 = dict()
@@ -434,3 +435,19 @@ def test_add_inplace_errors():
     ar2 = dpt.ones((1, 2), dtype="float32")
     with pytest.raises(ValueError):
         ar1 += ar2
+
+
+def test_add_inplace_overlap():
+    get_queue_or_skip()
+
+    ar1 = dpt.ones(10, dtype="i4")
+    ar1 += ar1
+    assert (dpt.asnumpy(ar1) == np.full(ar1.shape, 2, dtype="i4")).all()
+
+    ar1 = dpt.ones(10, dtype="i4")
+    ar2 = dpt.ones(10, dtype="i4")
+    dpt.add(ar1, ar2, out=ar1)
+    assert (dpt.asnumpy(ar1) == np.full(ar1.shape, 2, dtype="i4")).all()
+
+    dpt.add(ar2, ar1, out=ar2)
+    assert (dpt.asnumpy(ar2) == np.full(ar2.shape, 3, dtype="i4")).all()