Merge pull request #708 from IntelPython/tensor-reshuffling

Exported dpctl.tensor.Device as array-API device class

Author: oleksandr-pavlyk

dpctl/tensor/__init__.py

@@ -29,16 +29,15 @@
 
 """
 
-from dpctl.tensor._copy_utils import astype, copy
-from dpctl.tensor._copy_utils import copy_from_numpy as from_numpy
-from dpctl.tensor._copy_utils import copy_to_numpy as asnumpy
-from dpctl.tensor._copy_utils import copy_to_numpy as to_numpy
+from dpctl.tensor._copy_utils import asnumpy, astype, copy, from_numpy, to_numpy
 from dpctl.tensor._ctors import asarray, empty
+from dpctl.tensor._device import Device
 from dpctl.tensor._dlpack import from_dlpack
 from dpctl.tensor._reshape import reshape
 from dpctl.tensor._usmarray import usm_ndarray
 
 __all__ = [
+    "Device",
     "usm_ndarray",
     "asarray",
     "astype",

dpctl/tensor/_copy_utils.py

@@ -19,6 +19,7 @@
 
 import dpctl.memory as dpm
 import dpctl.tensor as dpt
+from dpctl.tensor._device import normalize_queue_device
 
 
 def contract_iter2(shape, strides1, strides2):
@@ -64,7 +65,7 @@ def contract_iter2(shape, strides1, strides2):
     return (sh, st1, disp1, st2, disp2)
 
 
-def has_memory_overlap(x1, x2):
+def _has_memory_overlap(x1, x2):
     m1 = dpm.as_usm_memory(x1)
     m2 = dpm.as_usm_memory(x2)
     if m1.sycl_device == m2.sycl_device:
@@ -77,7 +78,7 @@ def has_memory_overlap(x1, x2):
         return False
 
 
-def copy_to_numpy(ary):
+def _copy_to_numpy(ary):
     if type(ary) is not dpt.usm_ndarray:
         raise TypeError
     h = ary.usm_data.copy_to_host().view(ary.dtype)
@@ -93,7 +94,7 @@ def copy_to_numpy(ary):
     )
 
 
-def copy_from_numpy(np_ary, usm_type="device", sycl_queue=None):
+def _copy_from_numpy(np_ary, usm_type="device", sycl_queue=None):
     "Copies numpy array `np_ary` into a new usm_ndarray"
     # This may peform a copy to meet stated requirements
     Xnp = np.require(np_ary, requirements=["A", "O", "C", "E"])
@@ -111,7 +112,8 @@ def copy_from_numpy(np_ary, usm_type="device", sycl_queue=None):
     return Xusm
 
 
-def copy_from_numpy_into(dst, np_ary):
+def _copy_from_numpy_into(dst, np_ary):
+    "Copies `np_ary` into `dst` of type :class:`dpctl.tensor.usm_ndarray"
     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)
@@ -122,6 +124,54 @@ def copy_from_numpy_into(dst, np_ary):
         usm_mem.copy_from_host(host_buf)
 
 
+def from_numpy(np_ary, device=None, usm_type="device", sycl_queue=None):
+    """
+    from_numpy(arg, device=None, usm_type="device", sycl_queue=None)
+
+    Creates :class:`dpctl.tensor.usm_ndarray` from instance of
+    `numpy.ndarray`.
+
+    Args:
+        arg: An instance of `numpy.ndarray`
+        device: array API specification of device where the output array
+            is created.
+        sycl_queue: a :class:`dpctl.SyclQueue` used to create the output
+            array is created
+    """
+    q = normalize_queue_device(sycl_queue=sycl_queue, device=device)
+    return _copy_from_numpy(np_ary, usm_type=usm_type, sycl_queue=q)
+
+
+def to_numpy(usm_ary):
+    """
+    to_numpy(usm_ary)
+
+    Copies content of :class:`dpctl.tensor.usm_ndarray` instance `usm_ary`
+    into `numpy.ndarray` instance of the same shape and same data type.
+
+    Args:
+        usm_ary: An instance of :class:`dpctl.tensor.usm_ndarray`
+    Returns:
+        An instance of `numpy.ndarray` populated with content of `usm_ary`.
+    """
+    return _copy_to_numpy(usm_ary)
+
+
+def asnumpy(usm_ary):
+    """
+    asnumpy(usm_ary)
+
+    Copies content of :class:`dpctl.tensor.usm_ndarray` instance `usm_ary`
+    into `numpy.ndarray` instance of the same shape and same data type.
+
+    Args:
+        usm_ary: An instance of :class:`dpctl.tensor.usm_ndarray`
+    Returns:
+        An instance of `numpy.ndarray` populated with content of `usm_ary`.
+    """
+    return _copy_to_numpy(usm_ary)
+
+
 class Dummy:
     def __init__(self, iface):
         self.__sycl_usm_array_interface__ = iface
@@ -138,9 +188,9 @@ def copy_same_dtype(dst, src):
         raise ValueError
 
     # check that memory regions do not overlap
-    if has_memory_overlap(dst, src):
-        tmp = copy_to_numpy(src)
-        copy_from_numpy_into(dst, tmp)
+    if _has_memory_overlap(dst, src):
+        tmp = _copy_to_numpy(src)
+        _copy_from_numpy_into(dst, tmp)
         return
 
     if (dst.flags & 1) and (src.flags & 1):
@@ -184,10 +234,10 @@ def copy_same_shape(dst, src):
         return
 
     # check that memory regions do not overlap
-    if has_memory_overlap(dst, src):
-        tmp = copy_to_numpy(src)
+    if _has_memory_overlap(dst, src):
+        tmp = _copy_to_numpy(src)
         tmp = tmp.astype(dst.dtype)
-        copy_from_numpy_into(dst, tmp)
+        _copy_from_numpy_into(dst, tmp)
         return
 
     # simplify strides
@@ -218,7 +268,7 @@ def copy_same_shape(dst, src):
         mdst.copy_from_host(tmp.view("u1"))
 
 
-def copy_from_usm_ndarray_to_usm_ndarray(dst, src):
+def _copy_from_usm_ndarray_to_usm_ndarray(dst, src):
     if type(dst) is not dpt.usm_ndarray or type(src) is not dpt.usm_ndarray:
         raise TypeError
 
@@ -389,7 +439,7 @@ def astype(usm_ary, newdtype, order="K", casting="unsafe", copy=True):
                 buffer=R.usm_data,
                 strides=new_strides,
             )
-        copy_from_usm_ndarray_to_usm_ndarray(R, usm_ary)
+        _copy_from_usm_ndarray_to_usm_ndarray(R, usm_ary)
         return R
     else:
         return usm_ary

dpctl/tensor/_ctors.py

@@ -20,6 +20,7 @@
 import dpctl.memory as dpm
 import dpctl.tensor as dpt
 import dpctl.utils
+from dpctl.tensor._device import normalize_queue_device
 
 _empty_tuple = tuple()
 _host_set = frozenset([None])
@@ -72,29 +73,6 @@ def _array_info_sequence(li):
     return (n,) + dim, dt, device
 
 
-def _normalize_queue_device(q=None, d=None):
-    if q is None:
-        d = dpt._device.Device.create_device(d)
-        return d.sycl_queue
-    else:
-        if not isinstance(q, dpctl.SyclQueue):
-            raise TypeError(f"Expected dpctl.SyclQueue, got {type(q)}")
-        if d is None:
-            return q
-        d = dpt._device.Device.create_device(d)
-        qq = dpctl.utils.get_execution_queue(
-            (
-                q,
-                d.sycl_queue,
-            )
-        )
-        if qq is None:
-            raise TypeError(
-                "sycl_queue and device keywords can not be both specified"
-            )
-        return qq
-
-
 def _asarray_from_usm_ndarray(
     usm_ndary,
     dtype=None,
@@ -115,7 +93,7 @@ def _asarray_from_usm_ndarray(
         exec_q = dpctl.utils.get_execution_queue(
             [usm_ndary.sycl_queue, sycl_queue]
         )
-        copy_q = _normalize_queue_device(q=sycl_queue, d=exec_q)
+        copy_q = normalize_queue_device(sycl_queue=sycl_queue, device=exec_q)
     else:
         copy_q = usm_ndary.sycl_queue
     # Conditions for zero copy:
@@ -194,7 +172,7 @@ def _asarray_from_numpy_ndarray(
         usm_type = "device"
     if dtype is None:
         dtype = ary.dtype
-    copy_q = _normalize_queue_device(q=None, d=sycl_queue)
+    copy_q = normalize_queue_device(sycl_queue=None, device=sycl_queue)
     f_contig = ary.flags["F"]
     c_contig = ary.flags["C"]
     fc_contig = f_contig or c_contig
@@ -327,7 +305,9 @@ def asarray(
             )
     # 5. Normalize device/sycl_queue [keep it None if was None]
     if device is not None or sycl_queue is not None:
-        sycl_queue = _normalize_queue_device(q=sycl_queue, d=device)
+        sycl_queue = normalize_queue_device(
+            sycl_queue=sycl_queue, device=device
+        )
 
     # handle instance(obj, usm_ndarray)
     if isinstance(obj, dpt.usm_ndarray):
@@ -459,7 +439,7 @@ def empty(
         raise TypeError(
             f"Expected usm_type to be of type str, got {type(usm_type)}"
         )
-    sycl_queue = _normalize_queue_device(q=sycl_queue, d=device)
+    sycl_queue = normalize_queue_device(sycl_queue=sycl_queue, device=device)
     res = dpt.usm_ndarray(
         sh,
         dtype=dtype,

dpctl/tensor/_device.py

@@ -105,3 +105,54 @@ def __repr__(self):
         except TypeError:
             # This is a sub-device
             return repr(self.sycl_queue)
+
+
+def normalize_queue_device(sycl_queue=None, device=None):
+    """
+    normalize_queue_device(sycl_queue=None, device=None)
+
+    Utility to process exclusive keyword arguments 'device'
+    and 'sycl_queue' in functions of `dpctl.tensor`.
+
+    Args:
+        sycl_queue(:class:`dpctl.SyclQueue`, optional):
+            explicitly indicates where USM allocation is done
+            and the population code (if any) is executed.
+            Value `None` is interpreted as get the SYCL queue
+            from `device` keyword, or use default queue.
+            Default: None
+        device (string, :class:`dpctl.SyclDevice`, :class:`dpctl.SyclQueue,
+            :class:`dpctl.tensor.Device`, optional):
+            array-API keyword indicating non-partitioned SYCL device
+            where array is allocated.
+    Returns
+        :class:`dpctl.SyclQueue` object implied by either of provided
+        keywords. If both are None, `dpctl.SyclQueue()` is returned.
+        If both are specified and imply the same queue, `sycl_queue`
+        is returned.
+    Raises:
+        TypeError: if argument is not of the expected type, or keywords
+            imply incompatible queues.
+    """
+    q = sycl_queue
+    d = device
+    if q is None:
+        d = Device.create_device(d)
+        return d.sycl_queue
+    else:
+        if not isinstance(q, dpctl.SyclQueue):
+            raise TypeError(f"Expected dpctl.SyclQueue, got {type(q)}")
+        if d is None:
+            return q
+        d = Device.create_device(d)
+        qq = dpctl.utils.get_execution_queue(
+            (
+                q,
+                d.sycl_queue,
+            )
+        )
+        if qq is None:
+            raise TypeError(
+                "sycl_queue and device keywords can not be both specified"
+            )
+        return qq

dpctl/tensor/_usmarray.pyx

@@ -884,13 +884,13 @@ cdef class usm_ndarray:
         except (ValueError, IndexError) as e:
             raise e
         from ._copy_utils import (
-            copy_from_numpy_into,
-            copy_from_usm_ndarray_to_usm_ndarray,
+            _copy_from_numpy_into,
+            _copy_from_usm_ndarray_to_usm_ndarray,
         )
         if isinstance(val, usm_ndarray):
-            copy_from_usm_ndarray_to_usm_ndarray(Xv, val)
+            _copy_from_usm_ndarray_to_usm_ndarray(Xv, val)
         else:
-            copy_from_numpy_into(Xv, np.asarray(val))
+            _copy_from_numpy_into(Xv, np.asarray(val))
 
     def __sub__(first, other):
         "See comment in __add__"