8000 ENH: Configurable allocator (#17582) · numpy/numpy@84e0707 · GitHub
[go: up one dir, main page]
Skip to content

Commit 84e0707

Browse files
mattipeliaskoromilas
mattip
and
eliaskoromilas
authored
ENH: Configurable allocator (#17582)
Fixes gh-17467. Adds a public struct to hold memory manipulation routines PyDataMem_Handler and two new API functions PyDataMem_SetHandler to replace the current routines with the new ones, and PyDataMem_GetHandlerName to get the string name of the current routines (either globally or for a specific ndarray object). This also changes the size of the ndarray object to hold the PyDataMem_Handler active when it was created so subsequent actions on its data memory will remain consistent. Tests and documentation are included. Along the way, I found some places in the code where the current policy is inconsistent (all data memory handling should have gone through npy_*_cache not PyDataMem_*) so even if this is rejected it might improve the cache handling. The PyDataMem_Handler has fields to override memcpy, these are currently not implemented: memcpy in the code base is untouched. I think this PR is invasive enough as-is, if desired memcpy can be handled in a follow-up PR. * ENH: add and use global configurable memory routines * ENH: add tests and a way to compile c-extensions from tests * fix allocation/free exposed by tests * DOC: document the new APIs (and some old ones too) * BUG: return void from FREE, also some cleanup * MAINT: changes from review * fixes from linter * setting ndarray->descr on 0d or scalars mess with FREE * make scalar allocation more consistent wrt np_alloc_cache * change formatting for sphinx * remove memcpy variants * update to match NEP 49 * ENH: add a python-level get_handler_name * ENH: add core.multiarray.get_handler_name * Allow closure-like definition of the data mem routines * Fix incompatible pointer warnings * Note PyDataMemAllocator and PyMemAllocatorEx differentiation Co-authored-by: Matti Picus <matti.picus@gmail.com> * Redefine default allocator handling * Always allocate new arrays using the current_handler * Search for the mem_handler name of the data owner * Sub-comparisons don't need a local mem_handler * Make the default_handler a valid PyDataMem_Handler * Fix PyDataMem_SetHandler description (NEP discussion) * Pass the allocators by reference * Implement allocator context-locality * Fix documentation, make PyDataMem_GetHandler return const * remove import of setuptools==49.1.3, doesn't work on python3.10 * Fix refcount leaks * fix function signatures in test * Return early on PyDataMem_GetHandler error (VOID_compare) * Add context/thread-locality tests, allow testing custom policies * ENH: add and use global configurable memory routines * ENH: add tests and a way to compile c-extensions from tests * fix allocation/free exposed by tests * DOC: document the new APIs (and some old ones too) * BUG: return void from FREE, also some cleanup * MAINT: changes from review * fixes from linter * setting ndarray->descr on 0d or scalars mess with FREE * make scalar allocation more consistent wrt np_alloc_cache * change formatting for sphinx * remove memcpy variants * update to match NEP 49 * ENH: add a python-level get_handler_name * ENH: add core.multiarray.get_handler_name * Allow closure-like definition of the data mem routines * Fix incompatible pointer warnings * Note PyDataMemAllocator and PyMemAllocatorEx differentiation Co-authored-by: Matti Picus <matti.picus@gmail.com> * Redefine default allocator handling * Always allocate new arrays using the current_handler * Search for the mem_handler name of the data owner * Sub-comparisons don't need a local mem_handler * Make the default_handler a valid PyDataMem_Handler * Fix PyDataMem_SetHandler description (NEP discussion) * Pass the allocators by reference * remove import of setuptools==49.1.3, doesn't work on python3.10 * fix function signatures in test * try to fix cygwin extension building * YAPF mem_policy test * Less empty lines, more comments (tests) * Apply suggestions from code review (set an exception and) Co-authored-by: Matti Picus <matti.picus@gmail.com> * skip test on cygwin * update API hash for changed signature * TST: add gc.collect to make sure cycles are broken * Implement thread-locality for PyPy Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net> * Update numpy/core/tests/test_mem_policy.py Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net> * fixes from review * update circleci config * fix test * make the connection between OWNDATA and having a allocator handle more explicit * improve docstring, fix flake8 for tests * update PyDataMem_GetHandler() from review * Implement allocator lifetime management * update NEP and add best-effort handling of error in PyDataMem_UserFREE * ENH: fix and test for blindly taking ownership of data * Update doc/neps/nep-0049.rst Co-authored-by: Elias Koromilas <elias.koromilas@gmail.com>
1 parent 48e6ac6 commit 84e0707

28 files changed

+1351
-79
lines changed

‎doc/TESTS.rst.txt

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,21 @@ originally written without unit tests, there are still several modules
139139
that don't have tests yet. Please feel free to choose one of these
140140
modules and develop tests for it.
141141

142+
Using C code in tests
143+
---------------------
144+
145+
NumPy exposes a rich :ref:`C-API<c-api>` . These are tested using c-extension
146+
modules written "as-if" they know nothing about the internals of NumPy, rather
147+
using the official C-API interfaces only. Examples of such modules are tests
148+
for a user-defined ``rational`` dtype in ``_rational_tests`` or the ufunc
149+
machinery tests in ``_umath_tests`` which are part of the binary distribution.
150+
Starting from version 1.21, you can also write snippets of C code in tests that
151+
will be compiled locally into c-extension modules and loaded into python.
152+
153+
.. currentmodule:: numpy.testing.extbuild
154+
155+
.. autofunction:: build_and_import_extension
156+
142157
Labeling tests
143158
--------------
144159

‎doc/neps/nep-0049.rst

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -93,19 +93,21 @@ High level design
9393

9494
Users who wish to change the NumPy data memory management routines will use
9595
:c:func:`PyDataMem_SetHandler`, which uses a :c:type:`PyDataMem_Handler`
96-
structure to hold pointers to functions used to manage the data memory.
96+
structure to hold pointers to functions used to manage the data memory. In
97+
order to allow lifetime management of the ``context``, the structure is wrapped
98+
in a ``PyCapsule``.
9799

98100
Since a call to ``PyDataMem_SetHandler`` will change the default functions, but
99101
that function may be called during the lifetime of an ``ndarray`` object, each
100-
``ndarray`` will carry with it the ``PyDataMem_Handler`` struct used at the
101-
time of its instantiation, and these will be used to reallocate or free the
102-
data memory of the instance. Internally NumPy may use ``memcpy`` or ``memset``
103-
on the pointer to the data memory.
102+
``ndarray`` will carry with it the ``PyDataMem_Handler``-wrapped PyCapsule used
103+
at the time of its instantiation, and these will be used to reallocate or free
104+
the data memory of the instance. Internally NumPy may use ``memcpy`` or
105+
``memset`` on the pointer to the data memory.
104106

105107
The name of the handler will be exposed on the python level via a
106108
``numpy.core.multiarray.get_handler_name(arr)`` function. If called as
107109
``numpy.core.multiarray.get_handler_name()`` it will return the name of the
108-
global handler that will be used to allocate data for the next new `ndarrray`.
110+
handler that will be used to allocate data for the next new `ndarrray`.
109111

110112
NumPy C-API functions
111113
=====================
@@ -150,20 +152,19 @@ NumPy C-API functions
150152
15780_ and 15788_ but has not yet been resolved. When it is this NEP should
151153
be revisited.
152154

153-
.. c:function:: const PyDataMem_Handler * PyDataMem_SetHandler(PyDataMem_Handler *handler)
155+
.. c:function:: PyObject * PyDataMem_SetHandler(PyObject *handler)
154156
155157
Sets a new allocation policy. If the input value is ``NULL``, will reset
156-
the policy to the default. Returns the previous policy, ``NULL`` if the
157-
previous policy was the default. We wrap the user-provided functions
158+
the policy to the default. Return the previous policy, or
159+
return NULL if an error has occurred. We wrap the user-provided
158160
so they will still call the Python and NumPy memory management callback
159161
hooks. All the function pointers must be filled in, ``NULL`` is not
160162
accepted.
161163
162-
.. c:function:: const PyDataMem_Handler * PyDataMem_GetHandler(PyArrayObject *obj)
164+
.. c:function:: const PyObject * PyDataMem_GetHandler()
163165
164-
Return the ``PyDataMem_Handler`` used by the
165-
``PyArrayObject``. If ``NULL``, return the handler
166-
that will be used to allocate data for the next ``PyArrayObject``.
166+
Return the current policy that will be used to allocate data for the
167+
next ``PyArrayObject``. On failure, return ``NULL``.
167168
168169
``PyDataMem_Handler`` thread safety and lifetime
169170
================================================

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

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
Memory management in NumPy
2+
==========================
3+
4+
The `numpy.ndarray` is a python class. It requires additional memory allocations
5+
to hold `numpy.ndarray.strides`, `numpy.ndarray.shape` and
6+
`numpy.ndarray.data` attributes. These attributes are specially allocated
7+
after creating the python object in `__new__`. The ``strides`` and
8+
``shape`` are stored in a piece of memory allocated internally.
9+
10+
The ``data`` allocation used to store the actual array values (which could be
11+
pointers in the case of ``object`` arrays) can be very large, so NumPy has
12+
provided interfaces to manage its allocation and release. This document details
13+
how those interfaces work.
14+
15+
Historical overview
16+
-------------------
17+
18+
Since version 1.7.0, NumPy has exposed a set of ``PyDataMem_*`` functions
19+
(:c:func:`PyDataMem_NEW`, :c:func:`PyDataMem_FREE`, :c:func:`PyDataMem_RENEW`)
20+
which are backed by `alloc`, `free`, `realloc` respectively. In that version
21+
NumPy also exposed the `PyDataMem_EventHook` function described below, which
22+
wrap the OS-level calls.
23+
24+
Since those early days, Python also improved its memory management
25+
capabilities, and began providing
26+
various :ref:`management policies <memoryoverview>` beginning in version
27+
3.4. These routines are divided into a set of domains, each domain has a
28+
:c:type:`PyMemAllocatorEx` structure of routines for memory management. Python also
29+
added a `tracemalloc` module to trace calls to the various routines. These
30+
tracking hooks were added to the NumPy ``PyDataMem_*`` routines.
31+
32+
NumPy added a small cache of allocated memory in its internal
33+
``npy_alloc_cache``, ``npy_alloc_cache_zero``, and ``npy_free_cache``
34+
functions. These wrap ``alloc``, ``alloc-and-memset(0)`` and ``free``
35+
respectively, but when ``npy_free_cache`` is called, it adds the pointer to a
36+
short list of available blocks marked by size. These blocks can be re-used by
37+
subsequent calls to ``npy_alloc*``, avoiding memory thrashing.
38+
39+
Configurable memory routines in NumPy (NEP 49)
40+
----------------------------------------------
41+
42+
Users may wish to override the internal data memory routines with ones of their
43+
own. Since NumPy does not use the Python domain strategy to manage data memory,
44+
it provides an alternative set of C-APIs to change memory routines. There are
45+
no Python domain-wide strategies for large chunks of object data, so those are
46+
less suited to NumPy's needs. User who wish to change the NumPy data memory
47+
management routines can use :c:func:`PyDataMem_SetHandler`, which uses a
48+
:c:type:`PyDataMem_Handler` structure to hold pointers to functions used to
49+
manage the data memory. The calls are still wrapped by internal routines to
50+
call :c:func:`PyTraceMalloc_Track`, :c:func:`PyTraceMalloc_Untrack`, and will
51+
use the :c:func:`PyDataMem_EventHookFunc` mechanism. Since the functions may
52+
change during the lifetime of the process, each ``ndarray`` carries with it the
53+
functions used at the time of its instantiation, and these will be used to
54+
reallocate or free the data memory of the instance.
55+
56+
.. c:type:: PyDataMem_Handler
57+
58+
A struct to hold function pointers used to manipulate memory
59+
60+
.. code-block:: c
61+
62+
typedef struct {
63+
char name[128]; /* multiple of 64 to keep the struct aligned */
64+
PyDataMemAllocator allocator;
65+
} PyDataMem_Handler;
66+
67+
where the allocator structure is
68+
69+
.. code-block:: c
70+
71+
/* The declaration of free differs from PyMemAllocatorEx */
72+
typedef struct {
73+
void *ctx;
74+
void* (*malloc) (void *ctx, size_t size);
75+
void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
76+
void* (*realloc) (void *ctx, void *ptr, size_t new_size);
77+
void (*free) (void *ctx, void *ptr, size_t size);
78+
} PyDataMemAllocator;
79+
80+
.. c:function:: PyObject * PyDataMem_SetHandler(PyObject *handler)
81+
82+
Set a new allocation policy. If the input value is ``NULL``, will reset the
83+
policy to the default. Return the previous policy, or
84+
return ``NULL`` if an error has occurred. We wrap the user-provided functions
85+
so they will still call the python and numpy memory management callback
86+
hooks.
87+
88+
.. c:function:: PyObject * PyDataMem_GetHandler()
89+
90+
Return the current policy that will be used to allocate data for the
91+
next ``PyArrayObject``. On failure, return ``NULL``.
92+
93+
For an example of setting up and using the PyDataMem_Handler, see the test in
94+
:file:`numpy/core/tests/test_mem_policy.py`
95+
96+
.. c:function:: void PyDataMem_EventHookFunc(void *inp, void *outp, size_t size, void *user_data);
97+
98+
This function will be called during data memory manipulation
99+
100+
.. c:function:: PyDataMem_EventHookFunc * PyDataMem_SetEventHook(PyDataMem_EventHookFunc *newhook, void *user_data, void **old_data)
101+
102+
Sets the allocation event hook for numpy array data.
103+
104+
Returns a pointer to the previous hook or ``NULL``. If old_data is
105+
non-``NULL``, the previous user_data pointer will be copied to it.
106+
107+
If not ``NULL``, hook will be called at the end of each ``PyDataMem_NEW/FREE/RENEW``:
108+
109+
.. code-block:: c
110+
111+
result = PyDataMem_NEW(size) -> (*hook)(NULL, result, size, user_data)
112+
PyDataMem_FREE(ptr) -> (*hook)(ptr, NULL, 0, user_data)
113+
result = PyDataMem_RENEW(ptr, size) -> (*hook)(ptr, result, size, user_data)
114+
115+
When the hook is called, the GIL will be held by the calling
116+
thread. The hook should be written to be reentrant, if it performs
117+
operations that might cause new allocation events (such as the
118+
creation/destruction numpy objects, or creating/destroying Python
119+
objects which might cause a gc)

‎doc/source/reference/c-api/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,3 +49,4 @@ code.
4949
generalized-ufuncs
5050
coremath
5151
deprecations
52+
data_memory

‎numpy/core/_add_newdocs.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4727,6 +4727,16 @@
47274727
and then throwing away the ufunc.
47284728
""")
47294729

4730+
add_newdoc('numpy.core.multiarray', 'get_handler_name',
4731+
"""
4732+
get_handler_name(a: ndarray) -> str,None
4733+
4734+
Return the name of the memory handler used by `a`. If not provided, return
4735+
the name of the memory handler that will be used to allocate data for the
4736+
next `ndarray` in this context. May return None if `a` does not own its
4737+
memory, in which case you can traverse ``a.base`` for a memory handler.
4738+
""")
4739+
47304740
add_newdoc('numpy.core.multiarray', '_set_madvise_hugepage',
47314741
"""
47324742
_set_madvise_hugepage(enabled: bool) -> bool

‎numpy/core/code_generators/cversions.txt

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,5 +56,7 @@
5656
# DType related API additions.
5757
# A new field was added to the end of PyArrayObject_fields.
5858
# Version 14 (NumPy 1.21) No change.
59-
# Version 14 (NumPy 1.22) No change.
6059
0x0000000e = 17a0f366e55ec05e5c5c149123478452
60+
61+
# Version 15 (NumPy 1.22) Configurable memory allocations
62+
0x0000000f = 0c420aed67010594eb81f23ddfb02a88

‎numpy/core/code_generators/numpy_api.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,9 +76,9 @@
7676
# End 1.6 API
7777
}
7878

79-
#define NPY_NUMUSERTYPES (*(int *)PyArray_API[6])
80-
#define PyBoolArrType_Type (*(PyTypeObject *)PyArray_API[7])
81-
#define _PyArrayScalar_BoolValues ((PyBoolScalarObject *)PyArray_API[8])
79+
# define NPY_NUMUSERTYPES (*(int *)PyArray_API[6])
80+
# define PyBoolArrType_Type (*(PyTypeObject *)PyArray_API[7])
81+
# define _PyArrayScalar_BoolValues ((PyBoolScalarObject *)PyArray_API[8])
8282

8383
multiarray_funcs_api = {
8484
'PyArray_GetNDArrayCVersion': (0,),
@@ -350,6 +350,9 @@
350350
'PyArray_ResolveWritebackIfCopy': (302,),
351351
'PyArray_SetWritebackIfCopyBase': (303,),
352352
# End 1.14 API
353+
'PyDataMem_SetHandler': (304,),
354+
'PyDataMem_GetHandler': (305,),
355+
# End 1.21 API
353356
}
354357

355358
ufunc_types_api = {

‎numpy/core/include/numpy/ndarraytypes.h

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -355,12 +355,10 @@ struct NpyAuxData_tag {
355355
#define NPY_ERR(str) fprintf(stderr, #str); fflush(stderr);
356356
#define NPY_ERR2(str) fprintf(stderr, str); fflush(stderr);
357357

358-
/*
359-
* Macros to define how array, and dimension/strides data is
360-
* allocated.
361-
*/
362-
363-
/* Data buffer - PyDataMem_NEW/FREE/RENEW are in multiarraymodule.c */
358+
/*
359+
* Macros to define how array, and dimension/strides data is
360+
* allocated. These should be made private
361+
*/
364362

365363
#define NPY_USE_PYMEM 1
366364

@@ -666,6 +664,24 @@ typedef struct _arr_descr {
666664
PyObject *shape; /* a tuple */
667665
} PyArray_ArrayDescr;
668666

667+
/*
668+
* Memory handler structure for array data.
669+
*/
670+
/* The declaration of free differs from PyMemAllocatorEx */
671+
typedef struct {
672+
void *ctx;
673+
void* (*malloc) (void *ctx, size_t size);
674+
void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
675+
void* (*realloc) (void *ctx, void *ptr, size_t new_size);
676+
void (*free) (void *ctx, void *ptr, size_t size);
677+
} PyDataMemAllocator;
678+
679+
typedef struct {
680+
char name[128]; /* multiple of 64 to keep the struct aligned */
681+
PyDataMemAllocator allocator;
682+
} PyDataMem_Handler;
683+
684+
669685
/*
670686
* The main array object structure.
671687
*
@@ -716,6 +732,10 @@ typedef struct tagPyArrayObject_fields {
716732
/* For weak references */
717733
PyObject *weakreflist;
718734
void *_buffer_info; /* private buffer info, tagged to allow warning */
735+
/*
736+
* For malloc/calloc/realloc/free per object
737+
*/
738+
PyObject *mem_handler;
719739
} PyArrayObject_fields;
720740

721741
/*
@@ -1659,6 +1679,12 @@ PyArray_CLEARFLAGS(PyArrayObject *arr, int flags)
16591679
((PyArrayObject_fields *)arr)->flags &= ~flags;
16601680
}
16611681

1682+
static NPY_INLINE NPY_RETURNS_BORROWED_REF PyObject *
1683+
PyArray_HANDLER(PyArrayObject *arr)
1684+
{
1685+
return ((PyArrayObject_fields *)arr)->mem_handler;
1686+
}
1687+
16621688
#define PyTypeNum_ISBOOL(type) ((type) == NPY_BOOL)
16631689

16641690
#define PyTypeNum_ISUNSIGNED(type) (((type) == NPY_UBYTE) || \

‎numpy/core/multiarray.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@
3131
'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data',
3232
'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype',
3333
'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat',
34-
'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'inner',
35-
'interp', 'interp_complex', 'is_busday', 'lexsort',
34+
'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'get_handler_name',
35+
'inner', 'interp', 'interp_complex', 'is_busday', 'lexsort',
3636
'matmul', 'may_share_memory', 'min_scalar_type', 'ndarray', 'nditer',
3737
'nested_iters', 'normalize_axis_index', 'packbits',
3838
'promote_types', 'putmask', 'ravel_multi_index', 'result_type', 'scalar',

‎numpy/core/setup_common.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,8 @@
4343
# 0x0000000d - 1.19.x
4444
# 0x0000000e - 1.20.x
4545
# 0x0000000e - 1.21.x
46-
# 0x0000000e - 1.22.x
47-
C_API_VERSION = 0x0000000e
46+
# 0x0000000f - 1.22.x
47+
C_API_VERSION = 0x0000000f
4848

4949
class MismatchCAPIWarning(Warning):
5050
pass

0 commit comments

Comments
 (0)
0