Initial support for Numpy subclasses (numba#6148)

Co-authored-by: Todd A. Anderson <[email protected]>

Author: PokhodenkoSA

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/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,26 @@ 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);
+    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 +158,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 +315,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 +353,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 +368,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);

numba/core/runtime/nrt.c

@@ -19,6 +19,7 @@ struct MemInfo {
     void              *dtor_info;
     void              *data;
     size_t            size;    /* only used for NRT allocated memory */
+    NRT_ExternalAllocator *external_allocator;
 };
 
 
@@ -170,13 +171,16 @@ void NRT_MemSys_set_atomic_cas_stub(void) {
  */
 
 void NRT_MemInfo_init(NRT_MemInfo *mi,void *data, size_t size,
-                      NRT_dtor_function dtor, void *dtor_info)
+                      NRT_dtor_function dtor, void *dtor_info,
+                      NRT_ExternalAllocator *external_allocator)
 {
     mi->refct = 1;  /* starts with 1 refct */
     mi->dtor = dtor;
     mi->dtor_info = dtor_info;
     mi->data = data;
     mi->size = size;
+    mi->external_allocator = external_allocator;
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_init mi=%p external_allocator=%p\n", mi, external_allocator));
     /* Update stats */
     TheMSys.atomic_inc(&TheMSys.stats_mi_alloc);
 }
@@ -185,7 +189,8 @@ NRT_MemInfo *NRT_MemInfo_new(void *data, size_t size,
                              NRT_dtor_function dtor, void *dtor_info)
 {
     NRT_MemInfo *mi = NRT_Allocate(sizeof(NRT_MemInfo));
-    NRT_MemInfo_init(mi, data, size, dtor, dtor_info);
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_new mi=%p\n", mi));
+    NRT_MemInfo_init(mi, data, size, dtor, dtor_info, NULL);
     return mi;
 }
 
@@ -206,9 +211,10 @@ void nrt_internal_dtor_safe(void *ptr, size_t size, void *info) {
 }
 
 static
-void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out) {
+void *nrt_allocate_meminfo_and_data(size_t size, NRT_MemInfo **mi_out, NRT_ExternalAllocator *allocator) {
     NRT_MemInfo *mi;
-    char *base = NRT_Allocate(sizeof(NRT_MemInfo) + size);
+    NRT_Debug(nrt_debug_print("nrt_allocate_meminfo_and_data %p\n", allocator));
+    char *base = NRT_Allocate_External(sizeof(NRT_MemInfo) + size, allocator);
     mi = (NRT_MemInfo *) base;
     *mi_out = mi;
     return base + sizeof(NRT_MemInfo);
@@ -230,9 +236,17 @@ void nrt_internal_custom_dtor_safe(void *ptr, size_t size, void *info) {
 
 NRT_MemInfo *NRT_MemInfo_alloc(size_t size) {
     NRT_MemInfo *mi;
-    void *data = nrt_allocate_meminfo_and_data(size, &mi);
+    void *data = nrt_allocate_meminfo_and_data(size, &mi, NULL);
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc %p\n", data));
-    NRT_MemInfo_init(mi, data, size, NULL, NULL);
+    NRT_MemInfo_init(mi, data, size, NULL, NULL, NULL);
+    return mi;
+}
+
+NRT_MemInfo *NRT_MemInfo_alloc_external(size_t size, NRT_ExternalAllocator *allocator) {
+    NRT_MemInfo *mi;
+    void *data = nrt_allocate_meminfo_and_data(size, &mi, allocator);
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc %p\n", data));
+    NRT_MemInfo_init(mi, data, size, NULL, NULL, allocator);
     return mi;
 }
 
@@ -242,22 +256,23 @@ NRT_MemInfo *NRT_MemInfo_alloc_safe(size_t size) {
 
 NRT_MemInfo* NRT_MemInfo_alloc_dtor_safe(size_t size, NRT_dtor_function dtor) {
     NRT_MemInfo *mi;
-    void *data = nrt_allocate_meminfo_and_data(size, &mi);
+    void *data = nrt_allocate_meminfo_and_data(size, &mi, NULL);
     /* Only fill up a couple cachelines with debug markers, to minimize
        overhead. */
     memset(data, 0xCB, MIN(size, 256));
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_dtor_safe %p %zu\n", data, size));
-    NRT_MemInfo_init(mi, data, size, nrt_internal_custom_dtor_safe, dtor);
+    NRT_MemInfo_init(mi, data, size, nrt_internal_custom_dtor_safe, dtor, NULL);
     return mi;
 }
 
 
 static
 void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align,
-                                          NRT_MemInfo **mi)
+                                          NRT_MemInfo **mi, NRT_ExternalAllocator *allocator)
 {
     size_t offset, intptr, remainder;
-    char *base = nrt_allocate_meminfo_and_data(size + 2 * align, mi);
+    NRT_Debug(nrt_debug_print("nrt_allocate_meminfo_and_data_align %p\n", allocator));
+    char *base = nrt_allocate_meminfo_and_data(size + 2 * align, mi, allocator);
     intptr = (size_t) base;
     /* See if we are aligned */
     remainder = intptr % align;
@@ -271,26 +286,48 @@ void *nrt_allocate_meminfo_and_data_align(size_t size, unsigned align,
 
 NRT_MemInfo *NRT_MemInfo_alloc_aligned(size_t size, unsigned align) {
     NRT_MemInfo *mi;
-    void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi);
+    void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi, NULL);
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_aligned %p\n", data));
-    NRT_MemInfo_init(mi, data, size, NULL, NULL);
+    NRT_MemInfo_init(mi, data, size, NULL, NULL, NULL);
     return mi;
 }
 
 NRT_MemInfo *NRT_MemInfo_alloc_safe_aligned(size_t size, unsigned align) {
     NRT_MemInfo *mi;
-    void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi);
+    void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi, NULL);
     /* Only fill up a couple cachelines with debug markers, to minimize
        overhead. */
     memset(data, 0xCB, MIN(size, 256));
     NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_safe_aligned %p %zu\n",
                               data, size));
-    NRT_MemInfo_init(mi, data, size, nrt_internal_dtor_safe, (void*)size);
+    NRT_MemInfo_init(mi, data, size, nrt_internal_dtor_safe, (void*)size, NULL);
     return mi;
 }
 
+NRT_MemInfo *NRT_MemInfo_alloc_safe_aligned_external(size_t size, unsigned align, NRT_ExternalAllocator *allocator) {
+    NRT_MemInfo *mi;
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_safe_aligned_external %p\n", allocator));
+    void *data = nrt_allocate_meminfo_and_data_align(size, align, &mi, allocator);
+    /* Only fill up a couple cachelines with debug markers, to minimize
+       overhead. */
+    memset(data, 0xCB, MIN(size, 256));
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_alloc_safe_aligned %p %zu\n",
+                              data, size));
+    NRT_MemInfo_init(mi, data, size, nrt_internal_dtor_safe, (void*)size, allocator);
+    return mi;
+}
+
+void NRT_dealloc(NRT_MemInfo *mi) {
+    NRT_Debug(nrt_debug_print("NRT_dealloc meminfo: %p external_allocator: %p\n", mi, mi->external_allocator));
+    if (mi->external_allocator) {
+        mi->external_allocator->free(mi, mi->external_allocator->opaque_data);
+    } else {
+        NRT_Free(mi);
+    }
+}
+
 void NRT_MemInfo_destroy(NRT_MemInfo *mi) {
-    NRT_Free(mi);
+    NRT_dealloc(mi);
     TheMSys.atomic_inc(&TheMSys.stats_mi_free);
 }
 
@@ -328,6 +365,14 @@ size_t NRT_MemInfo_size(NRT_MemInfo* mi) {
     return mi->size;
 }
 
+void * NRT_MemInfo_external_allocator(NRT_MemInfo *mi) {
+    NRT_Debug(nrt_debug_print("NRT_MemInfo_external_allocator meminfo: %p external_allocator: %p\n", mi, mi->external_allocator));
+    return mi->external_allocator;
+}
+
+void *NRT_MemInfo_parent(NRT_MemInfo *mi) {
+    return mi->dtor_info;
+}
 
 void NRT_MemInfo_dump(NRT_MemInfo *mi, FILE *out) {
     fprintf(out, "MemInfo %p refcount %zu\n", mi, mi->refct);
@@ -414,8 +459,18 @@ void NRT_MemInfo_varsize_free(NRT_MemInfo *mi, void *ptr)
  */
 
 void* NRT_Allocate(size_t size) {
-    void *ptr = TheMSys.allocator.malloc(size);
-    NRT_Debug(nrt_debug_print("NRT_Allocate bytes=%zu ptr=%p\n", size, ptr));
+    return NRT_Allocate_External(size, NULL);
+}
+
+void* NRT_Allocate_External(size_t size, NRT_ExternalAllocator *allocator) {
+    void *ptr;
+    if (allocator) {
+        ptr = allocator->malloc(size, allocator->opaque_data);
+        NRT_Debug(nrt_debug_print("NRT_Allocate custom bytes=%zu ptr=%p\n", size, ptr));
+    } else {
+        ptr = TheMSys.allocator.malloc(size);
+        NRT_Debug(nrt_debug_print("NRT_Allocate bytes=%zu ptr=%p\n", size, ptr));
+    }
     TheMSys.atomic_inc(&TheMSys.stats_alloc);
     return ptr;
 }
@@ -460,6 +515,7 @@ NRT_MemInfo* nrt_manage_memory(void *data, NRT_managed_dtor dtor) {
 static const
 NRT_api_functions nrt_functions_table = {
     NRT_MemInfo_alloc,
+    NRT_MemInfo_alloc_external,
     nrt_manage_memory,
     NRT_MemInfo_acquire,
     NRT_MemInfo_release,