Merge pull request #724 from IntelPython/fix-slow-copy

Fixes #723: slow item assignment from numpy array

Author: oleksandr-pavlyk

dpctl/tensor/_copy_utils.py

@@ -117,6 +117,12 @@ def _copy_from_numpy_into(dst, np_ary):
     if not isinstance(np_ary, np.ndarray):
         raise TypeError("Expected numpy.ndarray, got {}".format(type(np_ary)))
     src_ary = np.broadcast_to(np.asarray(np_ary, dtype=dst.dtype), dst.shape)
+    if src_ary.size and (dst.flags & 1) and src_ary.flags["C"]:
+        dpm.as_usm_memory(dst).copy_from_host(src_ary.reshape((-1,)).view("u1"))
+        return
+    if src_ary.size and (dst.flags & 2) and src_ary.flags["F"]:
+        dpm.as_usm_memory(dst).copy_from_host(src_ary.reshape((-1,)).view("u1"))
+        return
     for i in range(dst.size):
         mi = np.unravel_index(i, dst.shape)
         host_buf = np.array(src_ary[mi], ndmin=1).view("u1")

dpctl/tests/test_tensor_asarray.py

@@ -67,7 +67,19 @@ def test_asarray_from_numpy():
     Xnp = np.arange(10)
     Y = dpt.asarray(Xnp, usm_type="device")
     assert type(Y) is dpt.usm_ndarray
-    assert Y.shape == (10,)
+    assert Y.shape == Xnp.shape
+    assert Y.dtype == Xnp.dtype
+    # Fortan contiguous case
+    Xnp = np.array([[1, 2, 3], [4, 5, 6]], dtype="f4", order="F")
+    Y = dpt.asarray(Xnp, usm_type="shared")
+    assert type(Y) is dpt.usm_ndarray
+    assert Y.shape == Xnp.shape
+    assert Y.dtype == Xnp.dtype
+    # general strided case
+    Xnp = np.array([[1, 2, 3], [4, 5, 6]], dtype="i8")
+    Y = dpt.asarray(Xnp[::-1, ::-1], usm_type="host")
+    assert type(Y) is dpt.usm_ndarray
+    assert Y.shape == Xnp.shape
     assert Y.dtype == Xnp.dtype