ENH: add support for python3.6 memory tracing

Python 3.6 added a private API for tracing arbitrary memory allocations
so the tracemalloc module [0] can be used with numpy.

closes gh-4663

[0] https://docs.python.org/3/library/tracemalloc.html

5 files changed, 40 insertions, 1 deletions

diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
index e5690ae4b..e145e3404 100644
--- a/numpy/core/src/multiarray/alloc.c
+++ b/numpy/core/src/multiarray/alloc.c
@@ -2,12 +2,25 @@
 #include <Python.h>
 #include "structmember.h"
 
+#if PY_VERSION_HEX >= 0x03060000
+#include <pymem.h>
+/* public api in 3.7 */
+#if PY_VERSION_HEX < 0x03070000
+#define PyTraceMalloc_Track _PyTraceMalloc_Track
+#define PyTraceMalloc_Untrack _PyTraceMalloc_Untrack
+#endif
+#else
+#define PyTraceMalloc_Track(...)
+#define PyTraceMalloc_Untrack(...)
+#endif
+
 #define NPY_NO_DEPRECATED_API NPY_API_VERSION
 #define _MULTIARRAYMODULE
 #include <numpy/ndarraytypes.h>
 #include "numpy/arrayobject.h"
 #include <numpy/npy_common.h>
 #include "npy_config.h"
+#include "alloc.h"
 
 #include <assert.h>
 
@@ -192,6 +205,7 @@ PyDataMem_NEW(size_t size)
         }
         NPY_DISABLE_C_API
     }
+    PyTraceMalloc_Track(NPY_TRACE_DOMAIN, (npy_uintp)result, size);
     return result;
 }
 
@@ -213,6 +227,7 @@ PyDataMem_NEW_ZEROED(size_t size, size_t elsize)
         }
         NPY_DISABLE_C_API
     }
+    PyTraceMalloc_Track(NPY_TRACE_DOMAIN, (npy_uintp)result, size);
     return result;
 }
 
@@ -222,6 +237,7 @@ PyDataMem_NEW_ZEROED(size_t size, size_t elsize)
 NPY_NO_EXPORT void
 PyDataMem_FREE(void *ptr)
 {
+    PyTraceMalloc_Untrack(NPY_TRACE_DOMAIN, (npy_uintp)ptr);
     free(ptr);
     if (_PyDataMem_eventhook != NULL) {
         NPY_ALLOW_C_API_DEF
@@ -243,6 +259,10 @@ PyDataMem_RENEW(void *ptr, size_t size)
     void *result;
 
     result = realloc(ptr, size);
+    if (result != ptr) {
+	PyTraceMalloc_Untrack(NPY_TRACE_DOMAIN, (npy_uintp)ptr);
+    }
+    PyTraceMalloc_Track(NPY_TRACE_DOMAIN, (npy_uintp)result, size);
     if (_PyDataMem_eventhook != NULL) {
         NPY_ALLOW_C_API_DEF
         NPY_ALLOW_C_API
diff --git a/numpy/core/src/multiarray/alloc.h b/numpy/core/src/multiarray/alloc.h
index 8f6b167d0..39eb99544 100644
--- a/numpy/core/src/multiarray/alloc.h
+++ b/numpy/core/src/multiarray/alloc.h
@@ -4,6 +4,8 @@
 #define _MULTIARRAYMODULE
 #include <numpy/ndarraytypes.h>
 
+#define NPY_TRACE_DOMAIN 389047
+
 NPY_NO_EXPORT void *
 npy_alloc_cache(npy_uintp sz);
 
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 73ba8c5c4..b6b134f2d 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -60,6 +60,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0;
 #include "templ_common.h" /* for npy_mul_with_overflow_intp */
 #include "compiled_base.h"
 #include "mem_overlap.h"
+#include "alloc.h"
 
 #include "get_attr_string.h"
 
@@ -4624,6 +4625,10 @@ PyMODINIT_FUNC initmultiarray(void) {
      */
     PyDict_SetItemString (d, "error", PyExc_Exception);
 
+    s = PyInt_FromLong(NPY_TRACE_DOMAIN);
+    PyDict_SetItemString(d, "tracemalloc_domain", s);
+    Py_DECREF(s);
+
     s = PyUString_FromString("3.1");
     PyDict_SetItemString(d, "__version__", s);
     Py_DECREF(s);
diff --git a/numpy/lib/__init__.py b/numpy/lib/__init__.py
index 4cdb76b20..847a3e896 100644
--- a/numpy/lib/__init__.py
+++ b/numpy/lib/__init__.py
@@ -25,8 +25,9 @@ from .financial import *
 from .arrayterator import Arrayterator
 from .arraypad import *
 from ._version import *
+from numpy.core.multiarray import tracemalloc_domain
 
-__all__ = ['emath', 'math']
+__all__ = ['emath', 'math', 'tracemalloc_domain']
 __all__ += type_check.__all__
 __all__ += index_tricks.__all__
 __all__ += function_base.__all__
diff --git a/tools/allocation_tracking/README.md b/tools/allocation_tracking/README.md
new file mode 100644
index 000000000..fd4f2c871
--- /dev/null
+++ b/tools/allocation_tracking/README.md
@@ -0,0 +1,11 @@
+Example for using the `PyDataMem_SetEventHook` to track allocations inside numpy.
+
+`alloc_hook.pyx` implements a hook in Cython that calls back into a python
+function. `track_allocations.py` uses it for a simple listing of allocations.
+It can be built with the `setup.py` file in this folder.
+
+Note that since Python 3.6 the builtin tracemalloc module can be used to
+track allocations inside numpy.
+Numpy places its CPU memory allocations into the `np.lib.tracemalloc_domain`
+domain.
+See https://docs.python.org/3/library/tracemalloc.html.