3 files changed, 27 insertions, 0 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);