Merge pull request #1530 from IntelPython/dpnp_use_inplace_dpctl

utilize new functionality of dpctl for in-place operators

Author: vtavana

dpnp/dpnp_algo/dpnp_elementwise_common.py

@@ -27,8 +27,6 @@
 # *****************************************************************************
 
 
-import dpctl
-import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
 from dpctl.tensor._elementwise_common import (
     BinaryElementwiseFunc,
@@ -538,32 +536,11 @@ def _call_divide(src1, src2, dst, sycl_queue, depends=None):
     return ti._divide(src1, src2, dst, sycl_queue, depends)
 
 
-def _call_divide_inplace(lhs, rhs, sycl_queue, depends=None):
-    """In place workaround until dpctl.tensor provides the functionality."""
-
-    if depends is None:
-        depends = []
-
-    # allocate temporary memory for out array
-    out = dpt.empty_like(lhs, dtype=dpnp.result_type(lhs.dtype, rhs.dtype))
-
-    # call a general callback
-    div_ht_, div_ev_ = _call_divide(lhs, rhs, out, sycl_queue, depends)
-
-    # store the result into left input array and return events
-    cp_ht_, cp_ev_ = ti._copy_usm_ndarray_into_usm_ndarray(
-        src=out, dst=lhs, sycl_queue=sycl_queue, depends=[div_ev_]
-    )
-    dpctl.SyclEvent.wait_for([div_ht_])
-    return (cp_ht_, cp_ev_)
-
-
 divide_func = BinaryElementwiseFunc(
     "divide",
     ti._divide_result_type,
     _call_divide,
     _divide_docstring_,
-    _call_divide_inplace,
 )
 
 

dpnp/dpnp_array.py

@@ -188,7 +188,10 @@ def __eq__(self, other):
     def __float__(self):
         return self._array_obj.__float__()
 
-    # '__floordiv__',
+    def __floordiv__(self, other):
+        """Return self//value."""
+        return dpnp.floor_divide(self, other)
+
     # '__format__',
 
     def __ge__(self, other):
@@ -227,15 +230,22 @@ def __iand__(self, other):
         dpnp.bitwise_and(self, other, out=self)
         return self
 
-    # '__ifloordiv__',
+    def __ifloordiv__(self, other):
+        """Return self//=value."""
+        dpnp.floor_divide(self, other, out=self)
+        return self
 
     def __ilshift__(self, other):
         """Return self<<=value."""
         dpnp.left_shift(self, other, out=self)
         return self
 
     # '__imatmul__',
-    # '__imod__',
+
+    def __imod__(self, other):
+        """Return self%=value."""
+        dpnp.remainder(self, other, out=self)
+        return self
 
     def __imul__(self, other):
         """Return self*=value."""
@@ -345,7 +355,8 @@ def __rand__(self, other):
     def __repr__(self):
         return dpt.usm_ndarray_repr(self._array_obj, prefix="array")
 
-    # '__rfloordiv__',
+    def __rfloordiv__(self, other):
+        return dpnp.floor_divide(self, other)
 
     def __rlshift__(self, other):
         return dpnp.left_shift(other, self)

dpnp/dpnp_iface_mathematical.py

@@ -901,6 +901,14 @@ def floor_divide(
 
     >>> np.floor_divide(np.array([1., 2., 3., 4.]), 2.5)
     array([ 0.,  0.,  1.,  1.])
+
+    The ``//`` operator can be used as a shorthand for ``floor_divide`` on
+    :class:`dpnp.ndarray`.
+
+    >>> x1 = np.array([1., 2., 3., 4.])
+    >>> x1 // 2.5
+    array([0., 0., 1., 1.])
+
     """
 
     return check_nd_call_func(

tests/test_bitwise.py

@@ -67,50 +67,41 @@ def test_bitwise_and(self, lhs, rhs, dtype):
         )
         assert_array_equal(dp_a & dp_b, np_a & np_b)
 
-        """
-        TODO: unmute once dpctl support that
         if (
             not (inp.isscalar(dp_a) or inp.isscalar(dp_b))
             and dp_a.shape == dp_b.shape
         ):
             dp_a &= dp_b
             np_a &= np_b
             assert_array_equal(dp_a, np_a)
-        """
 
     def test_bitwise_or(self, lhs, rhs, dtype):
         dp_a, dp_b, np_a, np_b = self._test_binary_int(
             "bitwise_or", lhs, rhs, dtype
         )
         assert_array_equal(dp_a | dp_b, np_a | np_b)
 
-        """
-        TODO: unmute once dpctl support that
         if (
             not (inp.isscalar(dp_a) or inp.isscalar(dp_b))
             and dp_a.shape == dp_b.shape
         ):
             dp_a |= dp_b
             np_a |= np_b
             assert_array_equal(dp_a, np_a)
-        """
 
     def test_bitwise_xor(self, lhs, rhs, dtype):
         dp_a, dp_b, np_a, np_b = self._test_binary_int(
             "bitwise_xor", lhs, rhs, dtype
         )
         assert_array_equal(dp_a ^ dp_b, np_a ^ np_b)
 
-        """
-        TODO: unmute once dpctl support that
         if (
             not (inp.isscalar(dp_a) or inp.isscalar(dp_b))
             and dp_a.shape == dp_b.shape
         ):
             dp_a ^= dp_b
             np_a ^= np_b
             assert_array_equal(dp_a, np_a)
-        """
 
     def test_invert(self, lhs, rhs, dtype):
         dp_a, np_a = self._test_unary_int("invert", lhs, dtype)
@@ -122,30 +113,24 @@ def test_left_shift(self, lhs, rhs, dtype):
         )
         assert_array_equal(dp_a << dp_b, np_a << np_b)
 
-        """
-        TODO: unmute once dpctl support that
         if (
             not (inp.isscalar(dp_a) or inp.isscalar(dp_b))
             and dp_a.shape == dp_b.shape
         ):
             dp_a <<= dp_b
             np_a <<= np_b
             assert_array_equal(dp_a, np_a)
-        """
 
     def test_right_shift(self, lhs, rhs, dtype):
         dp_a, dp_b, np_a, np_b = self._test_binary_int(
             "right_shift", lhs, rhs, dtype
         )
         assert_array_equal(dp_a >> dp_b, np_a >> np_b)
 
-        """
-        TODO: unmute once dpctl support that
         if (
             not (inp.isscalar(dp_a) or inp.isscalar(dp_b))
             and dp_a.shape == dp_b.shape
         ):
             dp_a >>= dp_b
             np_a >>= np_b
             assert_array_equal(dp_a, np_a)
-        """

tests/test_mathematical.py

@@ -1177,3 +1177,31 @@ def test_mean_scalar(self):
         result = dp_array.mean()
         expected = np_array.mean()
         assert_allclose(expected, result)
+
+
+@pytest.mark.parametrize(
+    "dtype", get_all_dtypes(no_bool=True, no_none=True, no_complex=True)
+)
+def test_inplace_remainder(dtype):
+    size = 21
+    np_a = numpy.arange(size, dtype=dtype)
+    dp_a = dpnp.arange(size, dtype=dtype)
+
+    np_a %= 4
+    dp_a %= 4
+
+    assert_allclose(dp_a, np_a)
+
+
+@pytest.mark.parametrize(
+    "dtype", get_all_dtypes(no_bool=True, no_none=True, no_complex=True)
+)
+def test_inplace_floor_divide(dtype):
+    size = 21
+    np_a = numpy.arange(size, dtype=dtype)
+    dp_a = dpnp.arange(size, dtype=dtype)
+
+    np_a //= 4
+    dp_a //= 4
+
+    assert_allclose(dp_a, np_a)

tests/test_usm_type.py

@@ -83,8 +83,35 @@ def test_coerced_usm_types_remainder(usm_type_x, usm_type_y):
     y = dp.arange(100, usm_type=usm_type_y).reshape(10, 10)
     y = y.T + 1
 
+    z = 100 % y
+    z = y % 7
     z = x % y
 
+    # inplace remainder
+    z %= y
+    z %= 5
+
+    assert x.usm_type == usm_type_x
+    assert y.usm_type == usm_type_y
+    assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y])
+
+
+@pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("usm_type_y", list_of_usm_types, ids=list_of_usm_types)
+def test_coerced_usm_types_floor_divide(usm_type_x, usm_type_y):
+    x = dp.arange(100, usm_type=usm_type_x).reshape(10, 10)
+    y = dp.arange(100, usm_type=usm_type_y).reshape(10, 10)
+    x = x + 1.5
+    y = y.T + 0.5
+
+    z = 3.4 // y
+    z = y // 2.7
+    z = x // y
+
+    # inplace floor_divide
+    z //= y
+    z //= 2.5
+
     assert x.usm_type == usm_type_x
     assert y.usm_type == usm_type_y
     assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y])