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 files changed, 251 insertions, 0 deletions

diff --git a/numpy/testing/_private/extbuild.py b/numpy/testing/_private/extbuild.py
new file mode 100644
index 000000000..8b3a438dd
--- /dev/null
+++ b/numpy/testing/_private/extbuild.py
@@ -0,0 +1,251 @@
+"""
+Build a c-extension module on-the-fly in tests.
+See build_and_import_extensions for usage hints
+
+"""
+
+import os
+import pathlib
+import sys
+import sysconfig
+from numpy.distutils.ccompiler import new_compiler
+from distutils.errors import CompileError
+
+__all__ = ['build_and_import_extension', 'compile_extension_module']
+
+
+def build_and_import_extension(
+        modname, functions, *, prologue="", build_dir=None,
+        include_dirs=[], more_init=""):
+    """
+    Build and imports a c-extension module `modname` from a list of function
+    fragments `functions`.
+
+
+    Parameters
+    ----------
+    functions : list of fragments
+        Each fragment is a sequence of func_name, calling convention, snippet.
+    prologue : string
+        Code to preceed the rest, usually extra ``#include`` or ``#define``
+        macros.
+    build_dir : pathlib.Path
+        Where to build the module, usually a temporary directory
+    include_dirs : list
+        Extra directories to find include files when compiling
+    more_init : string
+        Code to appear in the module PyMODINIT_FUNC
+
+    Returns
+    -------
+    out: module
+        The module will have been loaded and is ready for use
+
+    Examples
+    --------
+    >>> functions = [("test_bytes", "METH_O", \"\"\"
+        if ( !PyBytesCheck(args)) {
+            Py_RETURN_FALSE;
+        }
+        Py_RETURN_TRUE;
+    \"\"\")]
+    >>> mod = build_and_import_extension("testme", functions)
+    >>> assert not mod.test_bytes(u'abc')
+    >>> assert mod.test_bytes(b'abc')
+    """
+
+    body = prologue + _make_methods(functions, modname)
+    init = """PyObject *mod = PyModule_Create(&moduledef);
+           """
+    if not build_dir:
+        build_dir = pathlib.Path('.')
+    if more_init:
+        init += """#define INITERROR return NULL
+                """
+        init += more_init
+    init += "\nreturn mod;"
+    source_string = _make_source(modname, init, body)
+    try:
+        mod_so = compile_extension_module(
+            modname, build_dir, include_dirs, source_string)
+    except CompileError as e:
+        # shorten the exception chain
+        raise RuntimeError(f"could not compile in {build_dir}:") from e
+    import importlib.util
+    spec = importlib.util.spec_from_file_location(modname, mod_so)
+    foo = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(foo)
+    return foo
+
+
+def compile_extension_module(
+        name, builddir, include_dirs,
+        source_string, libraries=[], library_dirs=[]):
+    """
+    Build an extension module and return the filename of the resulting
+    native code file.
+
+    Parameters
+    ----------
+    name : string
+        name of the module, possibly including dots if it is a module inside a
+        package.
+    builddir : pathlib.Path
+        Where to build the module, usually a temporary directory
+    include_dirs : list
+        Extra directories to find include files when compiling
+    libraries : list
+        Libraries to link into the extension module
+    library_dirs: list
+        Where to find the libraries, ``-L`` passed to the linker
+    """
+    modname = name.split('.')[-1]
+    dirname = builddir / name
+    dirname.mkdir(exist_ok=True)
+    cfile = _convert_str_to_file(source_string, dirname)
+    include_dirs = [sysconfig.get_config_var('INCLUDEPY')] + include_dirs
+
+    return _c_compile(
+        cfile, outputfilename=dirname / modname,
+        include_dirs=include_dirs, libraries=[], library_dirs=[],
+        )
+
+
+def _convert_str_to_file(source, dirname):
+    """Helper function to create a file ``source.c`` in `dirname` that contains
+    the string in `source`. Returns the file name
+    """
+    filename = dirname / 'source.c'
+    with filename.open('w') as f:
+        f.write(str(source))
+    return filename
+
+
+def _make_methods(functions, modname):
+    """ Turns the name, signature, code in functions into complete functions
+    and lists them in a methods_table. Then turns the methods_table into a
+    ``PyMethodDef`` structure and returns the resulting code fragment ready
+    for compilation
+    """
+    methods_table = []
+    codes = []
+    for funcname, flags, code in functions:
+        cfuncname = "%s_%s" % (modname, funcname)
+        if 'METH_KEYWORDS' in flags:
+            signature = '(PyObject *self, PyObject *args, PyObject *kwargs)'
+        else:
+            signature = '(PyObject *self, PyObject *args)'
+        methods_table.append(
+            "{\"%s\", (PyCFunction)%s, %s}," % (funcname, cfuncname, flags))
+        func_code = """
+        static PyObject* {cfuncname}{signature}
+        {{
+        {code}
+        }}
+        """.format(cfuncname=cfuncname, signature=signature, code=code)
+        codes.append(func_code)
+
+    body = "\n".join(codes) + """
+    static PyMethodDef methods[] = {
+    %(methods)s
+    { NULL }
+    };
+    static struct PyModuleDef moduledef = {
+        PyModuleDef_HEAD_INIT,
+        "%(modname)s",  /* m_name */
+        NULL,           /* m_doc */
+        -1,             /* m_size */
+        methods,        /* m_methods */
+    };
+    """ % dict(methods='\n'.join(methods_table), modname=modname)
+    return body
+
+
+def _make_source(name, init, body):
+    """ Combines the code fragments into source code ready to be compiled
+    """
+    code = """
+    #include <Python.h>
+
+    %(body)s
+
+    PyMODINIT_FUNC
+    PyInit_%(name)s(void) {
+    %(init)s
+    }
+    """ % dict(
+        name=name, init=init, body=body,
+    )
+    return code
+
+
+def _c_compile(cfile, outputfilename, include_dirs=[], libraries=[],
+               library_dirs=[]):
+    if sys.platform == 'win32':
+        compile_extra = ["/we4013"]
+        link_extra = ["/LIBPATH:" + os.path.join(sys.exec_prefix, 'libs')]
+    elif sys.platform.startswith('linux'):
+        compile_extra = [
+            "-O0", "-g", "-Werror=implicit-function-declaration", "-fPIC"]
+        link_extra = None
+    else:
+        compile_extra = link_extra = None
+        pass
+    if sys.platform == 'win32':
+        link_extra = link_extra + ['/DEBUG']  # generate .pdb file
+    if sys.platform == 'darwin':
+        # support Fink & Darwinports
+        for s in ('/sw/', '/opt/local/'):
+            if (s + 'include' not in include_dirs
+                    and os.path.exists(s + 'include')):
+                include_dirs.append(s + 'include')
+            if s + 'lib' not in library_dirs and os.path.exists(s + 'lib'):
+                library_dirs.append(s + 'lib')
+
+    outputfilename = outputfilename.with_suffix(get_so_suffix())
+    saved_environ = os.environ.copy()
+    try:
+        build(
+            cfile, outputfilename,
+            compile_extra, link_extra,
+            include_dirs, libraries, library_dirs)
+    finally:
+        # workaround for a distutils bugs where some env vars can
+        # become longer and longer every time it is used
+        for key, value in saved_environ.items():
+            if os.environ.get(key) != value:
+                os.environ[key] = value
+    return outputfilename
+
+
+def build(cfile, outputfilename, compile_extra, link_extra,
+          include_dirs, libraries, library_dirs):
+    "cd into the directory where the cfile is, use distutils to build"
+
+    compiler = new_compiler(force=1, verbose=2)
+    compiler.customize('')
+    objects = []
+
+    old = os.getcwd()
+    os.chdir(cfile.parent)
+    try:
+        res = compiler.compile(
+            [str(cfile.name)],
+            include_dirs=include_dirs,
+            extra_preargs=compile_extra
+            )
+        objects += [str(cfile.parent / r) for r in res]
+    finally:
+        os.chdir(old)
+
+    compiler.link_shared_object(
+        objects, str(outputfilename),
+        libraries=libraries,
+        extra_preargs=link_extra,
+        library_dirs=library_dirs)
+
+
+def get_so_suffix():
+    ret = sysconfig.get_config_var('EXT_SUFFIX')
+    assert ret
+    return ret