dparray integration

numba/_typeof.c

@@ -768,6 +768,7 @@ int
 typeof_typecode(PyObject *dispatcher, PyObject *val)
 {
     PyTypeObject *tyobj = Py_TYPE(val);
+    int no_subtype_attr;
     /* This needs to be kept in sync with Dispatcher.typeof_pyval(),
      * otherwise funny things may happen.
      */
@@ -794,9 +795,19 @@ typeof_typecode(PyObject *dispatcher, PyObject *val)
         return typecode_arrayscalar(dispatcher, val);
     }
     /* Array handling */
-    else if (PyType_IsSubtype(tyobj, &PyArray_Type)) {
+    else if (tyobj == &PyArray_Type) {
         return typecode_ndarray(dispatcher, (PyArrayObject*)val);
     }
+    /* Subtypes of Array handling */
+    else if (PyType_IsSubtype(tyobj, &PyArray_Type)) {
+        /* If the class has an attribute named __numba_no_subtype_ndarray then
+           don't treat it as a normal variant of a Numpy ndarray but as it's own
+           separate type. */
+        no_subtype_attr = PyObject_HasAttrString(val, "__numba_no_subtype_ndarray__");
+        if (!no_subtype_attr) {
+            return typecode_ndarray(dispatcher, (PyArrayObject*)val);
+        }
+    }
 
     return typecode_using_fingerprint(dispatcher, val);
 }

numba/core/base.py

@@ -435,9 +435,7 @@ def get_external_function_type(self, fndesc):
         return fnty
 
     def declare_function(self, module, fndesc):
-        #print("base.py: declare_function", module, "\n\targs:", fndesc.args, "\n\trestype:", fndesc.restype, "\n\targtypes:", fndesc.argtypes, fndesc.mangled_name, fndesc.noalias)
         fnty = self.call_conv.get_function_type(fndesc.restype, fndesc.argtypes)
-        #print("fnty:", fnty)
         fn = module.get_or_insert_function(fnty, name=fndesc.mangled_name)
         self.call_conv.decorate_function(fn, fndesc.args, fndesc.argtypes, noalias=fndesc.noalias)
         if fndesc.inline:
@@ -553,21 +551,14 @@ def get_function(self, fn, sig, _firstcall=True):
         The return value is a callable with the signature (builder, args).
         """
         assert sig is not None
-#        print("get_function", fn, sig, type(fn), type(sig))
         sig = sig.as_function()
-#        print("get_function", sig, type(sig))
         if isinstance(fn, (types.Function, types.BoundFunction,
                            types.Dispatcher)):
             key = fn.get_impl_key(sig)
             overloads = self._defns[key]
-#            print("function or boundfunction or dispatcher")
         else:
             key = fn
             overloads = self._defns[key]
-#            print("other")
-#        print("overloads", overloads)
-#        print_overloads(overloads.versions)
-#        print_defns(self._defns)
 
         try:
             return _wrap_impl(overloads.find(sig.args), self, sig)

numba/core/cpu.py

@@ -66,6 +66,7 @@ def load_additional_registries(self):
         from numba.np import npyimpl
         from numba.cpython import cmathimpl, mathimpl, printimpl, randomimpl
         from numba.misc import cffiimpl
+        from numba.dppl.dparray import numba_register as dparray_register
         self.install_registry(cmathimpl.registry)
         self.install_registry(cffiimpl.registry)
         self.install_registry(mathimpl.registry)
@@ -75,6 +76,7 @@ def load_additional_registries(self):
 
         # load 3rd party extensions
         numba.core.entrypoints.init_all()
+        dparray_register()
 
     @property
     def target_data(self):

numba/core/extending.py

@@ -14,7 +14,7 @@
     lower_setattr, lower_setattr_generic, lower_cast)  # noqa: F401
 from numba.core.datamodel import models   # noqa: F401
 from numba.core.datamodel import register_default as register_model  # noqa: F401, E501
-from numba.core.pythonapi import box, unbox, reflect, NativeValue  # noqa: F401
+from numba.core.pythonapi import box, unbox, reflect, NativeValue, allocator  # noqa: F401
 from numba._helperlib import _import_cython_function  # noqa: F401
 from numba.core.serialize import ReduceMixin
 

numba/core/ir_utils.py

@@ -64,6 +64,8 @@ def mk_alloc(typemap, calltypes, lhs, size_var, dtype, scope, loc):
     out = []
     ndims = 1
     size_typ = types.intp
+    # Get the type of the array being allocated.
+    arr_typ = typemap[lhs.name]
     if isinstance(size_var, tuple):
         if len(size_var) == 1:
             size_var = size_var[0]
@@ -108,11 +110,13 @@ def mk_alloc(typemap, calltypes, lhs, size_var, dtype, scope, loc):
     typ_var_assign = ir.Assign(np_typ_getattr, typ_var, loc)
     alloc_call = ir.Expr.call(attr_var, [size_var, typ_var], (), loc)
     if calltypes:
-        calltypes[alloc_call] = typemap[attr_var.name].get_call_type(
+        cac = typemap[attr_var.name].get_call_type(
             typing.Context(), [size_typ, types.functions.NumberClass(dtype)], {})
-    # signature(
-    #    types.npytypes.Array(dtype, ndims, 'C'), size_typ,
-    #    types.functions.NumberClass(dtype))
+        # By default, all calls to "empty" are typed as returning a standard
+        # Numpy ndarray.  If we are allocating a ndarray subclass here then
+        # just change the return type to be that of the subclass.
+        cac._return_type = arr_typ
+        calltypes[alloc_call] = cac
     alloc_assign = ir.Assign(alloc_call, lhs, loc)
 
     out.extend([g_np_assign, attr_assign, typ_var_assign, alloc_assign])

numba/core/pythonapi.py

@@ -45,10 +45,13 @@ def lookup(self, typeclass, default=None):
 _boxers = _Registry()
 _unboxers = _Registry()
 _reflectors = _Registry()
+# Registry of special allocators for types.
+_allocators = _Registry()
 
 box = _boxers.register
 unbox = _unboxers.register
 reflect = _reflectors.register
+allocator = _allocators.register
 
 class _BoxContext(namedtuple("_BoxContext",
                   ("context", "builder", "pyapi", "env_manager"))):
@@ -1186,8 +1189,11 @@ def nrt_adapt_ndarray_to_python(self, aryty, ary, dtypeptr):
         assert self.context.enable_nrt, "NRT required"
 
         intty = ir.IntType(32)
+        # Embed the Python type of the array (maybe subclass) in the LLVM.
+        serial_aryty_pytype = self.unserialize(self.serialize_object(aryty.py_type))
+
         fnty = Type.function(self.pyobj,
-                             [self.voidptr, intty, intty, self.pyobj])
+                             [self.voidptr, self.pyobj, intty, intty, self.pyobj])
         fn = self._get_function(fnty, name="NRT_adapt_ndarray_to_python")
         fn.args[0].add_attribute(lc.ATTR_NO_CAPTURE)
 
@@ -1197,6 +1203,7 @@ def nrt_adapt_ndarray_to_python(self, aryty, ary, dtypeptr):
         aryptr = cgutils.alloca_once_value(self.builder, ary)
         return self.builder.call(fn, [self.builder.bitcast(aryptr,
                                                            self.voidptr),
+                                      serial_aryty_pytype,
                                       ndim, writable, dtypeptr])
 
     def nrt_meminfo_new_from_pyobject(self, data, pyobj):

numba/core/runtime/_nrt_python.c

@@ -55,6 +55,8 @@ int MemInfo_init(MemInfoObject *self, PyObject *args, PyObject *kwds) {
         return -1;
     }
     raw_ptr = PyLong_AsVoidPtr(raw_ptr_obj);
+    NRT_Debug(nrt_debug_print("MemInfo_init self=%p raw_ptr=%p\n", self, raw_ptr));
+
     if(PyErr_Occurred()) return -1;
     self->meminfo = (NRT_MemInfo *)raw_ptr;
     assert (NRT_MemInfo_refcount(self->meminfo) > 0 && "0 refcount");
@@ -109,6 +111,27 @@ MemInfo_get_refcount(MemInfoObject *self, void *closure) {
     return PyLong_FromSize_t(refct);
 }
 
+static
+PyObject*
+MemInfo_get_external_allocator(MemInfoObject *self, void *closure) {
+    void *p = NRT_MemInfo_external_allocator(self->meminfo);
+    printf("MemInfo_get_external_allocator %p\n", p);
+    return PyLong_FromVoidPtr(p);
+}
+
+static
+PyObject*
+MemInfo_get_parent(MemInfoObject *self, void *closure) {
+    void *p = NRT_MemInfo_parent(self->meminfo);
+    if (p) {
+        Py_INCREF(p);
+        return (PyObject*)p;
+    } else {
+        Py_INCREF(Py_None);
+        return Py_None;
+    }
+}
+
 static void
 MemInfo_dealloc(MemInfoObject *self)
 {
@@ -136,6 +159,13 @@ static PyGetSetDef MemInfo_getsets[] = {
      (getter)MemInfo_get_refcount, NULL,
      "Get the refcount",
      NULL},
+    {"external_allocator",
+     (getter)MemInfo_get_external_allocator, NULL,
+     "Get the external allocator",
+     NULL},
+    {"parent",
+     (getter)MemInfo_get_parent, NULL,
+     NULL},
     {NULL}  /* Sentinel */
 };
 
@@ -286,7 +316,7 @@ PyObject* try_to_return_parent(arystruct_t *arystruct, int ndim,
 }
 
 NUMBA_EXPORT_FUNC(PyObject *)
-NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
+NRT_adapt_ndarray_to_python(arystruct_t* arystruct, PyTypeObject *retty, int ndim,
                             int writeable, PyArray_Descr *descr)
 {
     PyArrayObject *array;
@@ -324,10 +354,13 @@ NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
         args = PyTuple_New(1);
         /* SETITEM steals reference */
         PyTuple_SET_ITEM(args, 0, PyLong_FromVoidPtr(arystruct->meminfo));
+        NRT_Debug(nrt_debug_print("NRT_adapt_ndarray_to_python arystruct->meminfo=%p\n", arystruct->meminfo));
         /*  Note: MemInfo_init() does not incref.  This function steals the
          *        NRT reference.
          */
+        NRT_Debug(nrt_debug_print("NRT_adapt_ndarray_to_python created MemInfo=%p\n", miobj));
         if (MemInfo_init(miobj, args, NULL)) {
+            NRT_Debug(nrt_debug_print("MemInfo_init returned 0.\n"));
             return NULL;
         }
         Py_DECREF(args);
@@ -336,7 +369,7 @@ NRT_adapt_ndarray_to_python(arystruct_t* arystruct, int ndim,
     shape = arystruct->shape_and_strides;
     strides = shape + ndim;
     Py_INCREF((PyObject *) descr);
-    array = (PyArrayObject *) PyArray_NewFromDescr(&PyArray_Type, descr, ndim,
+    array = (PyArrayObject *) PyArray_NewFromDescr(retty, descr, ndim,
                                                    shape, strides, arystruct->data,
                                                    flags, (PyObject *) miobj);
 

numba/core/runtime/_nrt_pythonmod.c

@@ -163,6 +163,7 @@ declmethod(MemInfo_alloc);
 declmethod(MemInfo_alloc_safe);
 declmethod(MemInfo_alloc_aligned);
 declmethod(MemInfo_alloc_safe_aligned);
+declmethod(MemInfo_alloc_safe_aligned_external);
 declmethod(MemInfo_alloc_dtor_safe);
 declmethod(MemInfo_call_dtor);
 declmethod(MemInfo_new_varsize);