Add dparray

dpctl/__init__.pxd

@@ -29,4 +29,3 @@
 
 from dpctl._sycl_core cimport *
 from dpctl._memory import *
-

dpctl/dparray.py

@@ -0,0 +1,261 @@
+##===---------- 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
+from types import FunctionType as ftype, BuiltinFunctionType as bftype
+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)
+
+
+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, {})
+            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, {})
+            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, {})
+            return new_obj
+
+    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 *

dpctl/tests/test_dparray.py

@@ -0,0 +1,79 @@
+##===---------- test_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
+### A basic unit test for dpctl.dparray.
+##===----------------------------------------------------------------------===##
+
+import unittest
+from dpctl import dparray
+import numpy
+
+
+class Test_dparray(unittest.TestCase):
+    def setUp(self):
+        self.X = dparray.ndarray((256, 4), dtype="d")
+        self.X.fill(1.0)
+
+    def test_dparray_type(self):
+        self.assertIsInstance(self.X, dparray.ndarray)
+
+    def test_dparray_as_ndarray_self(self):
+        Y = self.X.as_ndarray()
+        self.assertEqual(type(Y), numpy.ndarray)
+
+    def test_dparray_as_ndarray(self):
+        Y = dparray.as_ndarray(self.X)
+        self.assertEqual(type(Y), numpy.ndarray)
+
+    def test_dparray_from_ndarray(self):
+        Y = dparray.as_ndarray(self.X)
+        dp1 = dparray.from_ndarray(Y)
+        self.assertIsInstance(dp1, dparray.ndarray)
+
+    def test_multiplication_dparray(self):
+        C = self.X * 5
+        self.assertIsInstance(C, dparray.ndarray)
+
+    def test_dparray_mixing_dpctl_and_numpy(self):
+        dp_numpy = numpy.ones((256, 4), dtype="d")
+        res = dp_numpy * self.X
+        self.assertIsInstance(res, dparray.ndarray)
+
+    def test_dparray_shape(self):
+        res = self.X.shape
+        self.assertEqual(res, (256, 4))
+
+    def test_dparray_T(self):
+        res = self.X.T
+        self.assertEqual(res.shape, (4, 256))
+
+    def test_numpy_ravel_with_dparray(self):
+        res = numpy.ravel(self.X)
+        self.assertEqual(res.shape, (1024,))
+
+    @unittest.expectedFailure
+    def test_numpy_sum_with_dparray(self):
+        res = numpy.sum(self.X)
+        self.assertEqual(res, 1024.0)
+
+
+if __name__ == "__main__":
+    unittest.main()