Merge pull request numpy#20343 from mattip/alloc-version

ENH: add a 'version' field to PyDataMem_Handler

Author: seberg

doc/neps/nep-0049.rst

@@ -109,6 +109,15 @@ The name of the handler will be exposed on the python level via a
 ``numpy.core.multiarray.get_handler_name()`` it will return the name of the
 handler that will be used to allocate data for the next new `ndarrray`.
 
+The version of the handler will be exposed on the python level via a
+``numpy.core.multiarray.get_handler_version(arr)`` function. If called as
+``numpy.core.multiarray.get_handler_version()`` it will return the version of the
+handler that will be used to allocate data for the next new `ndarrray`.
+
+The version, currently 1, allows for future enhancements to the
+``PyDataMemAllocator``. If fields are added, they must be added to the end.
+
+
 NumPy C-API functions
 =====================
 
@@ -119,7 +128,8 @@ NumPy C-API functions
     .. code-block:: c
 
         typedef struct {
-            char name[128];  /* multiple of 64 to keep the struct aligned */
+            char name[127];  /* multiple of 64 to keep the struct aligned */
+            uint8_t version; /* currently 1 */
             PyDataMemAllocator allocator;
         } PyDataMem_Handler;
 
@@ -279,6 +289,7 @@ the ``sz`` argument is correct.
 
     static PyDataMem_Handler new_handler = {
         "secret_data_allocator",
+        1,
         {
             &new_handler_ctx,
             shift_alloc,      /* malloc */

doc/source/reference/c-api/data_memory.rst

@@ -62,7 +62,8 @@ reallocate or free the data memory of the instance.
     .. code-block:: c
 
         typedef struct {
-            char name[128];  /* multiple of 64 to keep the struct aligned */
+            char name[127];  /* multiple of 64 to keep the struct aligned */
+            uint8_t version; /* currently 1 */
             PyDataMemAllocator allocator;
         } PyDataMem_Handler;
 

numpy/core/_add_newdocs.py

@@ -4759,6 +4759,16 @@
     memory, in which case you can traverse ``a.base`` for a memory handler.
     """)
 
+add_newdoc('numpy.core.multiarray', 'get_handler_version',
+    """
+    get_handler_version(a: ndarray) -> int,None
+
+    Return the version of the memory handler used by `a`. If not provided,
+    return the version of the memory handler that will be used to allocate data
+    for the next `ndarray` in this context. May return None if `a` does not own
+    its memory, in which case you can traverse ``a.base`` for a memory handler.
+    """)
+
 add_newdoc('numpy.core.multiarray', '_set_madvise_hugepage',
     """
     _set_madvise_hugepage(enabled: bool) -> bool

numpy/core/include/numpy/ndarraytypes.h

@@ -674,10 +674,15 @@ typedef struct {
     void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
     void* (*realloc) (void *ctx, void *ptr, size_t new_size);
     void (*free) (void *ctx, void *ptr, size_t size);
+    /*
+     * This is the end of the version=1 struct. Only add new fields after
+     * this line
+     */
 } PyDataMemAllocator;
 
 typedef struct {
-    char name[128];  /* multiple of 64 to keep the struct aligned */
+    char name[127];  /* multiple of 64 to keep the struct aligned */
+    uint8_t version; /* currently 1 */
     PyDataMemAllocator allocator;
 } PyDataMem_Handler;
 

numpy/core/multiarray.py

@@ -33,14 +33,14 @@
     'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype',
     'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat',
     'frombuffer', 'fromfile', 'fromiter', 'fromstring',
-    'get_handler_name', 'inner', 'interp', 'interp_complex', 'is_busday',
-    'lexsort', 'matmul', 'may_share_memory', 'min_scalar_type', 'ndarray',
-    'nditer', 'nested_iters', 'normalize_axis_index', 'packbits',
-    'promote_types', 'putmask', 'ravel_multi_index', 'result_type', 'scalar',
-    'set_datetimeparse_function', 'set_legacy_print_mode', 'set_numeric_ops',
-    'set_string_function', 'set_typeDict', 'shares_memory',
-    'tracemalloc_domain', 'typeinfo', 'unpackbits', 'unravel_index', 'vdot',
-    'where', 'zeros']
+    'get_handler_name', 'get_handler_version', 'inner', 'interp',
+    'interp_complex', 'is_busday', 'lexsort', 'matmul', 'may_share_memory',
+    'min_scalar_type', 'ndarray', 'nditer', 'nested_iters',
+    'normalize_axis_index', 'packbits', 'promote_types', 'putmask',
+    'ravel_multi_index', 'result_type', 'scalar', 'set_datetimeparse_function',
+    'set_legacy_print_mode', 'set_numeric_ops', 'set_string_function',
+    'set_typeDict', 'shares_memory', 'tracemalloc_domain', 'typeinfo',
+    'unpackbits', 'unravel_index', 'vdot', 'where', 'zeros']
 
 # For backward compatibility, make sure pickle imports these functions from here
 _reconstruct.__module__ = 'numpy.core.multiarray'

numpy/core/src/multiarray/alloc.c

@@ -370,6 +370,7 @@ default_free(void *NPY_UNUSED(ctx), void *ptr, size_t size)
 /* Memory handler global default */
 PyDataMem_Handler default_handler = {
     "default_allocator",
+    1,
     {
         NULL,            /* ctx */
         default_malloc,  /* malloc */
@@ -395,7 +396,6 @@ PyDataMem_UserNEW(size_t size, PyObject *mem_handler)
     if (handler == NULL) {
         return NULL;
     }
-
     assert(size != 0);
     result = handler->allocator.malloc(handler->allocator.ctx, size);
     if (_PyDataMem_eventhook != NULL) {
@@ -639,3 +639,40 @@ get_handler_name(PyObject *NPY_UNUSED(self), PyObject *args)
     Py_DECREF(mem_handler);
     return name;
 }
+
+NPY_NO_EXPORT PyObject *
+get_handler_version(PyObject *NPY_UNUSED(self), PyObject *args)
+{
+    PyObject *arr=NULL;
+    if (!PyArg_ParseTuple(args, "|O:get_handler_version", &arr)) {
+        return NULL;
+    }
+    if (arr != NULL && !PyArray_Check(arr)) {
+         PyErr_SetString(PyExc_ValueError, "if supplied, argument must be an ndarray");
+         return NULL;
+    }
+    PyObject *mem_handler;
+    PyDataMem_Handler *handler;
+    PyObject *version;
+    if (arr != NULL) {
+        mem_handler = PyArray_HANDLER((PyArrayObject *) arr);
+        if (mem_handler == NULL) {
+            Py_RETURN_NONE;
+        }
+        Py_INCREF(mem_handler);
+    }
+    else {
+        mem_handler = PyDataMem_GetHandler();
+        if (mem_handler == NULL) {
+            return NULL;
+        }
+    }
+    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    if (handler == NULL) {
+        Py_DECREF(mem_handler);
+        return NULL;
+    }
+    version = PyLong_FromLong(handler->version);
+    Py_DECREF(mem_handler);
+    return version;
+}

numpy/core/src/multiarray/alloc.h

@@ -47,5 +47,7 @@ extern PyDataMem_Handler default_handler;
 
 NPY_NO_EXPORT PyObject *
 get_handler_name(PyObject *NPY_UNUSED(self), PyObject *obj);
+NPY_NO_EXPORT PyObject *
+get_handler_version(PyObject *NPY_UNUSED(self), PyObject *obj);
 
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_ALLOC_H_ */

numpy/core/src/multiarray/multiarraymodule.c

@@ -4437,6 +4437,9 @@ static struct PyMethodDef array_module_methods[] = {
     {"get_handler_name",
         (PyCFunction) get_handler_name,
         METH_VARARGS, NULL},
+    {"get_handler_version",
+        (PyCFunction) get_handler_version,
+        METH_VARARGS, NULL},
     {"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc,
         METH_VARARGS, NULL},
     {"_get_sfloat_dtype",

numpy/core/tests/test_mem_policy.py

@@ -179,6 +179,7 @@ def get_module(tmp_path):
         };
         static PyDataMem_Handler secret_data_handler = {
             "secret_data_allocator",
+            1,
             {
                 &secret_data_handler_ctx, /* ctx */
                 shift_alloc,              /* malloc */
@@ -212,17 +213,22 @@ def get_module(tmp_path):
 def test_set_policy(get_module):
 
     get_handler_name = np.core.multiarray.get_handler_name
+    get_handler_version = np.core.multiarray.get_handler_version
     orig_policy_name = get_handler_name()
 
     a = np.arange(10).reshape((2, 5))  # a doesn't own its own data
     assert get_handler_name(a) is None
+    assert get_handler_version(a) is None
     assert get_handler_name(a.base) == orig_policy_name
+    assert get_handler_version(a.base) == 1
 
     orig_policy = get_module.set_secret_data_policy()
 
     b = np.arange(10).reshape((2, 5))  # b doesn't own its own data
     assert get_handler_name(b) is None
+    assert get_handler_version(b) is None
     assert get_handler_name(b.base) == 'secret_data_allocator'
+    assert get_handler_version(b.base) == 1
 
     if orig_policy_name == 'default_allocator':
         get_module.set_old_policy(None)  # tests PyDataMem_SetHandler(NULL)