Merged 192 200

dpctl/__init__.pxd

@@ -28,4 +28,3 @@
 # cython: language_level=3
 
 from dpctl._sycl_core cimport *
-from dpctl._memory import *

dpctl/dptensor/__init__.py

@@ -0,0 +1 @@
+import dpctl.dptensor.numpy_usm_shared

dpctl/dptensor/numpy_usm_shared.py

@@ -0,0 +1,304 @@
+##===---------- dparray.py - dpctl  -------*- Python -*----===##
+##
+##                      Data Parallel Control (dpCtl)
+##
+## Copyright 2020 Intel Corporation
+##
+## Licensed under the Apache License, Version 2.0 (the "License");
+## you may not use this file except in compliance with the License.
+## You may obtain a copy of the License at
+##
+##    http://www.apache.org/licenses/LICENSE-2.0
+##
+## Unless required by applicable law or agreed to in writing, software
+## distributed under the License is distributed on an "AS IS" BASIS,
+## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+## See the License for the specific language governing permissions and
+## limitations under the License.
+##
+##===----------------------------------------------------------------------===##
+###
+### \file
+### This file implements a dparray - USM aware implementation of ndarray.
+##===----------------------------------------------------------------------===##
+
+import numpy as np
+from inspect import getmembers, isfunction, isclass, isbuiltin
+from numbers import Number
+import sys
+import inspect
+import dpctl
+from dpctl.memory import MemoryUSMShared
+
+debug = False
+
+
+def dprint(*args):
+    if debug:
+        print(*args)
+        sys.stdout.flush()
+
+
+functions_list = [o[0] for o in getmembers(np) if isfunction(o[1]) or isbuiltin(o[1])]
+class_list = [o for o in getmembers(np) if isclass(o[1])]
+
+array_interface_property = "__sycl_usm_array_interface__"
+
+
+def has_array_interface(x):
+    return hasattr(x, array_interface_property)
+
+
+def _get_usm_base(ary):
+    ob = ary
+    while True:
+        if ob is None:
+            return None
+        elif hasattr(ob, "__sycl_usm_array_interface__"):
+            return ob
+        elif isinstance(ob, np.ndarray):
+            ob = ob.base
+        elif isinstance(ob, memoryview):
+            ob = ob.obj
+        else:
+            return None
+
+
+class ndarray(np.ndarray):
+    """
+    numpy.ndarray subclass whose underlying memory buffer is allocated
+    with a foreign allocator.
+    """
+
+    def __new__(
+        subtype, shape, dtype=float, buffer=None, offset=0, strides=None, order=None
+    ):
+        # Create a new array.
+        if buffer is None:
+            dprint("dparray::ndarray __new__ buffer None")
+            nelems = np.prod(shape)
+            dt = np.dtype(dtype)
+            isz = dt.itemsize
+            nbytes = int(isz * max(1, nelems))
+            buf = MemoryUSMShared(nbytes)
+            new_obj = np.ndarray.__new__(
+                subtype,
+                shape,
+                dtype=dt,
+                buffer=buf,
+                offset=0,
+                strides=strides,
+                order=order,
+            )
+            if hasattr(new_obj, array_interface_property):
+                dprint("buffer None new_obj already has sycl_usm")
+            else:
+                dprint("buffer None new_obj will add sycl_usm")
+                setattr(
+                    new_obj,
+                    array_interface_property,
+                    new_obj._getter_sycl_usm_array_interface_(),
+                )
+            return new_obj
+        # zero copy if buffer is a usm backed array-like thing
+        elif hasattr(buffer, array_interface_property):
+            dprint("dparray::ndarray __new__ buffer", array_interface_property)
+            # also check for array interface
+            new_obj = np.ndarray.__new__(
+                subtype,
+                shape,
+                dtype=dtype,
+                buffer=buffer,
+                offset=offset,
+                strides=strides,
+                order=order,
+            )
+            if hasattr(new_obj, array_interface_property):
+                dprint("buffer None new_obj already has sycl_usm")
+            else:
+                dprint("buffer None new_obj will add sycl_usm")
+                setattr(
+                    new_obj,
+                    array_interface_property,
+                    new_obj._getter_sycl_usm_array_interface_(),
+                )
+            return new_obj
+        else:
+            dprint("dparray::ndarray __new__ buffer not None and not sycl_usm")
+            nelems = np.prod(shape)
+            # must copy
+            ar = np.ndarray(
+                shape,
+                dtype=dtype,
+                buffer=buffer,
+                offset=offset,
+                strides=strides,
+                order=order,
+            )
+            nbytes = int(ar.nbytes)
+            buf = MemoryUSMShared(nbytes)
+            new_obj = np.ndarray.__new__(
+                subtype,
+                shape,
+                dtype=dtype,
+                buffer=buf,
+                offset=0,
+                strides=strides,
+                order=order,
+            )
+            np.copyto(new_obj, ar, casting="no")
+            if hasattr(new_obj, array_interface_property):
+                dprint("buffer None new_obj already has sycl_usm")
+            else:
+                dprint("buffer None new_obj will add sycl_usm")
+                setattr(
+                    new_obj,
+                    array_interface_property,
+                    new_obj._getter_sycl_usm_array_interface_(),
+                )
+            return new_obj
+
+    def _getter_sycl_usm_array_interface_(self):
+        ary_iface = self.__array_interface__
+        _base = _get_usm_base(self)
+        if _base is None:
+            raise TypeError
+
+        usm_iface = getattr(_base, "__sycl_usm_array_interface__", None)
+        if usm_iface is None:
+            raise TypeError
+
+        if ary_iface["data"][0] == usm_iface["data"][0]:
+            ary_iface["version"] = usm_iface["version"]
+            ary_iface["syclobj"] = usm_iface["syclobj"]
+        else:
+            raise TypeError
+        return ary_iface
+
+    def __array_finalize__(self, obj):
+        dprint("__array_finalize__:", obj, hex(id(obj)), type(obj))
+        # When called from the explicit constructor, obj is None
+        if obj is None:
+            return
+        # When called in new-from-template, `obj` is another instance of our own
+        # subclass, that we might use to update the new `self` instance.
+        # However, when called from view casting, `obj` can be an instance of any
+        # subclass of ndarray, including our own.
+        if hasattr(obj, array_interface_property):
+            return
+        if isinstance(obj, np.ndarray):
+            ob = self
+            while isinstance(ob, np.ndarray):
+                if hasattr(ob, array_interface_property):
+                    return
+                ob = ob.base
+
+        # Just raise an exception since __array_ufunc__ makes all reasonable cases not
+        # need the code below.
+        raise ValueError(
+            "Non-USM allocated ndarray can not viewed as a USM-allocated one without a copy"
+        )
+
+    # Tell Numba to not treat this type just like a NumPy ndarray but to propagate its type.
+    # This way it will use the custom dparray allocator.
+    __numba_no_subtype_ndarray__ = True
+
+    # Convert to a NumPy ndarray.
+    def as_ndarray(self):
+        return np.copy(self)
+
+    def __array__(self):
+        return self
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        if method == "__call__":
+            N = None
+            scalars = []
+            typing = []
+            for inp in inputs:
+                if isinstance(inp, Number):
+                    scalars.append(inp)
+                    typing.append(inp)
+                elif isinstance(inp, (self.__class__, np.ndarray)):
+                    if isinstance(inp, self.__class__):
+                        scalars.append(np.ndarray(inp.shape, inp.dtype, inp))
+                        typing.append(np.ndarray(inp.shape, inp.dtype))
+                    else:
+                        scalars.append(inp)
+                        typing.append(inp)
+                    if N is not None:
+                        if N != inp.shape:
+                            raise TypeError("inconsistent sizes")
+                    else:
+                        N = inp.shape
+                else:
+                    return NotImplemented
+            # Have to avoid recursive calls to array_ufunc here.
+            # If no out kwarg then we create a dparray out so that we get
+            # USM memory.  However, if kwarg has dparray-typed out then
+            # array_ufunc is called recursively so we cast out as regular
+            # NumPy ndarray (having a USM data pointer).
+            if kwargs.get("out", None) is None:
+                # maybe copy?
+                # deal with multiple returned arrays, so kwargs['out'] can be tuple
+                res_type = np.result_type(*typing)
+                out = empty(inputs[0].shape, dtype=res_type)
+                out_as_np = np.ndarray(out.shape, out.dtype, out)
+                kwargs["out"] = out_as_np
+            else:
+                # If they manually gave dparray as out kwarg then we have to also
+                # cast as regular NumPy ndarray to avoid recursion.
+                if isinstance(kwargs["out"], ndarray):
+                    out = kwargs["out"]
+                    kwargs["out"] = np.ndarray(out.shape, out.dtype, out)
+                else:
+                    out = kwargs["out"]
+            ret = ufunc(*scalars, **kwargs)
+            return out
+        else:
+            return NotImplemented
+
+
+def isdef(x):
+    try:
+        eval(x)
+        return True
+    except NameError:
+        return False
+
+
+for c in class_list:
+    cname = c[0]
+    if isdef(cname):
+        continue
+    # For now we do the simple thing and copy the types from NumPy module into dparray module.
+    new_func = "%s = np.%s" % (cname, cname)
+    try:
+        the_code = compile(new_func, "__init__", "exec")
+        exec(the_code)
+    except:
+        print("Failed to exec type propagation", cname)
+        pass
+
+# Redefine all Numpy functions in this module and if they
+# return a Numpy array, transform that to a USM-backed array
+# instead.  This is a stop-gap.  We should eventually find a
+# way to do the allocation correct to start with.
+for fname in functions_list:
+    if isdef(fname):
+        continue
+    new_func = "def %s(*args, **kwargs):\n" % fname
+    new_func += "    ret = np.%s(*args, **kwargs)\n" % fname
+    new_func += "    if type(ret) == np.ndarray:\n"
+    new_func += "        ret = ndarray(ret.shape, ret.dtype, ret)\n"
+    new_func += "    return ret\n"
+    the_code = compile(new_func, "__init__", "exec")
+    exec(the_code)
+
+
+def from_ndarray(x):
+    return copy(x)
+
+
+def as_ndarray(x):
+    return np.copy(x)

dpctl/tests/__init__.py

@@ -22,6 +22,7 @@
 # Top-level module of all dpctl Python unit test cases.
 # ===-----------------------------------------------------------------------===#
 
+from .test_dparray import *
 from .test_dump_functions import *
 from .test_sycl_device import *
 from .test_sycl_kernel_submit import *
@@ -30,3 +31,4 @@
 from .test_sycl_queue_manager import *
 from .test_sycl_queue_memcpy import *
 from .test_sycl_usm import *
+from .test_dparray import *