sklam · Jul 30, 2022 · Jun 28, 2022 · Jul 20, 2022 · Jul 20, 2022 · Jul 20, 2022
diff --git a/docs/source/reference/envvars.rst b/docs/source/reference/envvars.rst
@@ -148,6 +148,12 @@ These variables influence what is printed out during compilation of
    allocation and ``0xDE`` on deallocation, both to help with debugging memory
    leaks.
 
+.. envvar:: NUMBA_NRT_STATS
+
+   If set to non-zero, enable the
+   :ref:`Numba run time (NRT) <arch-numba-runtime>` statistics counters. These
+   counters are enabled process wide on import of Numba and are atomic.
+
 .. envvar:: NUMBA_DEBUGINFO
 
    If set to non-zero, enable debug for the full application by setting

diff --git a/numba/core/config.py b/numba/core/config.py
@@ -223,6 +223,9 @@ def optional_str(x):
         # Enable debug prints in nrtdynmod and use of "safe" API functions
         DEBUG_NRT = _readenv("NUMBA_DEBUG_NRT", int, 0)
 
+        # Enable NRT statistics counters
+        NRT_STATS = _readenv("NUMBA_NRT_STATS", int, 0)
+
         # How many recently deserialized functions to retain regardless
         # of external references
         FUNCTION_CACHE_SIZE = _readenv("NUMBA_FUNCTION_CACHE_SIZE", int, 128)

diff --git a/numba/core/runtime/_nrt_pythonmod.c b/numba/core/runtime/_nrt_pythonmod.c
@@ -19,24 +19,60 @@ memsys_use_cpython_allocator(PyObject *self, PyObject *args) {
 
 static PyObject *
 memsys_get_stats_alloc(PyObject *self, PyObject *args) {
+    if(!NRT_MemSys_stats_enabled()) {
+        PyErr_SetString(PyExc_RuntimeError, "NRT stats are disabled.");
+        return NULL;
+    }
     return PyLong_FromSize_t(NRT_MemSys_get_stats_alloc());
 }
 
 static PyObject *
 memsys_get_stats_free(PyObject *self, PyObject *args) {
+    if(!NRT_MemSys_stats_enabled()) {
+        PyErr_SetString(PyExc_RuntimeError, "NRT stats are disabled.");
+        return NULL;
+    }
     return PyLong_FromSize_t(NRT_MemSys_get_stats_free());
 }
 
 static PyObject *
 memsys_get_stats_mi_alloc(PyObject *self, PyObject *args) {
+    if(!NRT_MemSys_stats_enabled()) {
+        PyErr_SetString(PyExc_RuntimeError, "NRT stats are disabled.");
+        return NULL;
+    }
     return PyLong_FromSize_t(NRT_MemSys_get_stats_mi_alloc());
 }
 
 static PyObject *
 memsys_get_stats_mi_free(PyObject *self, PyObject *args) {
+    if(!NRT_MemSys_stats_enabled()) {
+        PyErr_SetString(PyExc_RuntimeError, "NRT stats are disabled.");
+        return NULL;
+    }
     return PyLong_FromSize_t(NRT_MemSys_get_stats_mi_free());
 }
 
+static PyObject *
+memsys_stats_enabled(PyObject *self, PyObject *args) {
+    if (NRT_MemSys_stats_enabled()) {
+        Py_RETURN_TRUE;
+    } else {
+        Py_RETURN_FALSE;
+    }
+}
+
+static PyObject *
+memsys_enable_stats(PyObject *self, PyObject *args) {
+    NRT_MemSys_enable_stats();
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+memsys_disable_stats(PyObject *self, PyObject *args) {
+    NRT_MemSys_disable_stats();
+    Py_RETURN_NONE;
+}
 
 /*
  * Create a new MemInfo with a owner PyObject
@@ -95,6 +131,9 @@ static PyMethodDef ext_methods[] = {
     declmethod_noargs(memsys_get_stats_free),
     declmethod_noargs(memsys_get_stats_mi_alloc),
     declmethod_noargs(memsys_get_stats_mi_free),
+    declmethod_noargs(memsys_stats_enabled),
+    declmethod_noargs(memsys_enable_stats),
+    declmethod_noargs(memsys_disable_stats),
     declmethod(meminfo_new),
     declmethod(meminfo_alloc),
     declmethod(meminfo_alloc_safe),

diff --git a/numba/core/runtime/nrt.cpp b/numba/core/runtime/nrt.cpp
@@ -48,7 +48,13 @@ struct NRT_MemSys {
     /* Shutdown flag */
     int shutting;
     /* Stats */
-    std::atomic_size_t stats_alloc, stats_free, stats_mi_alloc, stats_mi_free;
+    struct {
+        bool enabled;
+        std::atomic_size_t alloc;
+        std::atomic_size_t free;
+        std::atomic_size_t mi_alloc;
+        std::atomic_size_t mi_free;
+    } stats;
     /* System allocation functions */
     struct {
         NRT_malloc_func malloc;
@@ -64,10 +70,12 @@ static NRT_MemSys TheMSys;
 
 extern "C" void NRT_MemSys_init(void) {
     TheMSys.shutting = 0;
-    TheMSys.stats_alloc = 0;
-    TheMSys.stats_free = 0;
-    TheMSys.stats_mi_alloc = 0;
-    TheMSys.stats_mi_free = 0;
+    // Stats are off by default, call NRT_MemSys_enable_stats to enable
+    TheMSys.stats.enabled = false;
+    TheMSys.stats.alloc = 0;
+    TheMSys.stats.free = 0;
+    TheMSys.stats.mi_alloc = 0;
+    TheMSys.stats.mi_free = 0;
     /* Bind to libc allocator */
     TheMSys.allocator.malloc = malloc;
     TheMSys.allocator.realloc = realloc;
@@ -78,36 +86,76 @@ extern "C" void NRT_MemSys_shutdown(void) {
     TheMSys.shutting = 1;
 }
 
+extern "C" void NRT_MemSys_enable_stats(void) {
+    TheMSys.stats.enabled = true;
+}
+
+extern "C" void NRT_MemSys_disable_stats(void) {
+    TheMSys.stats.enabled = false;
+}
+
+extern "C" size_t NRT_MemSys_stats_enabled(void) {
+    return (size_t)TheMSys.stats.enabled;
+}
+
 extern "C" void NRT_MemSys_set_allocator(NRT_malloc_func malloc_func,
                               NRT_realloc_func realloc_func,
                               NRT_free_func free_func)
 {
+    bool stats_cond = false;
+    if (TheMSys.stats.enabled)
+    {
+        stats_cond = (TheMSys.stats.alloc != TheMSys.stats.free ||
+                      TheMSys.stats.mi_alloc != TheMSys.stats.mi_free);
+    }
     if ((malloc_func != TheMSys.allocator.malloc ||
          realloc_func != TheMSys.allocator.realloc ||
          free_func != TheMSys.allocator.free) &&
-         (TheMSys.stats_alloc != TheMSys.stats_free ||
-          TheMSys.stats_mi_alloc != TheMSys.stats_mi_free)) {
+         stats_cond) {
         nrt_fatal_error("cannot change allocator while blocks are allocated");
     }
     TheMSys.allocator.malloc = malloc_func;
     TheMSys.allocator.realloc = realloc_func;
     TheMSys.allocator.free = free_func;
 }
 
+/* This value is used as a marker for "stats are disabled", it's ASCII "AAAA" */
+static size_t _DISABLED_STATS_VALUE = 0x41414141;
+
 extern "C" size_t NRT_MemSys_get_stats_alloc() {
-    return TheMSys.stats_alloc;
+    if (TheMSys.stats.enabled)
+    {
+        return TheMSys.stats.alloc.load();
+    } else  {
+        return _DISABLED_STATS_VALUE;
+    }
 }
 
 extern "C" size_t NRT_MemSys_get_stats_free() {
-    return TheMSys.stats_free;
+    if (TheMSys.stats.enabled)
+    {
+        return TheMSys.stats.free.load();
+    } else  {
+        return _DISABLED_STATS_VALUE;
+    }
 }
 
 extern "C" size_t NRT_MemSys_get_stats_mi_alloc() {
-    return TheMSys.stats_mi_alloc;
+    if (TheMSys.stats.enabled)
+    {
+        return TheMSys.stats.mi_alloc.load();
+    } else  {
+        return _DISABLED_STATS_VALUE;
+    }
 }
 
 extern "C" size_t NRT_MemSys_get_stats_mi_free() {
-    return TheMSys.stats_mi_free;
+    if (TheMSys.stats.enabled)
+    {
+        return TheMSys.stats.mi_free.load();
+    } else  {
+        return _DISABLED_STATS_VALUE;
+    }
 }
 
 /*
@@ -126,7 +174,10 @@ extern "C" void NRT_MemInfo_init(NRT_MemInfo *mi,void *data, size_t size,
     mi->external_allocator = external_allocator;
     NRT_Debug(nrt_debug_print("NRT_MemInfo_init mi=%p external_allocator=%p\n", mi, external_allocator));
     /* Update stats */
-    TheMSys.stats_mi_alloc++;
+    if (TheMSys.stats.enabled)
+    {
+        TheMSys.stats.mi_alloc++;
+    }
 }
 
 NRT_MemInfo *NRT_MemInfo_new(void *data, size_t size,
@@ -319,15 +370,21 @@ extern "C" void NRT_dealloc(NRT_MemInfo *mi) {
     NRT_Debug(nrt_debug_print("NRT_dealloc meminfo: %p external_allocator: %p\n", mi, mi->external_allocator));
     if (mi->external_allocator) {
         mi->external_allocator->free(mi, mi->external_allocator->opaque_data);
-        TheMSys.stats_free++;;
+        if (TheMSys.stats.enabled)
+        {
+            TheMSys.stats.free++;
+        }
     } else {
         NRT_Free(mi);
     }
 }
 
 extern "C" void NRT_MemInfo_destroy(NRT_MemInfo *mi) {
     NRT_dealloc(mi);
-    TheMSys.stats_mi_free++;
+    if (TheMSys.stats.enabled)
+    {
+        TheMSys.stats.mi_free++;
+    }
 }
 
 extern "C" void NRT_MemInfo_acquire(NRT_MemInfo *mi) {
@@ -471,7 +528,10 @@ extern "C" void* NRT_Allocate_External(size_t size, NRT_ExternalAllocator *alloc
         ptr = TheMSys.allocator.malloc(size);
         NRT_Debug(nrt_debug_print("NRT_Allocate_External bytes=%zu ptr=%p\n", size, ptr));
     }
-    TheMSys.stats_alloc++;
+    if (TheMSys.stats.enabled)
+    {
+        TheMSys.stats.alloc++;
+    }
     return ptr;
 }
 
@@ -485,7 +545,10 @@ extern "C" void *NRT_Reallocate(void *ptr, size_t size) {
 extern "C" void NRT_Free(void *ptr) {
     NRT_Debug(nrt_debug_print("NRT_Free %p\n", ptr));
     TheMSys.allocator.free(ptr);
-    TheMSys.stats_free++;;
+    if (TheMSys.stats.enabled)
+    {
+        TheMSys.stats.free++;
+    }
 }
 
 /*

diff --git a/numba/core/runtime/nrt.h b/numba/core/runtime/nrt.h
@@ -48,6 +48,25 @@ void NRT_MemSys_shutdown(void);
 VISIBILITY_HIDDEN
 void NRT_MemSys_set_allocator(NRT_malloc_func, NRT_realloc_func, NRT_free_func);
 
+/*
+ * Enable the internal statistics counters.
+ */
+VISIBILITY_HIDDEN
+void NRT_MemSys_enable_stats(void);
+
+/*
+ * Disable the internal statistics counters.
+ */
+VISIBILITY_HIDDEN
+void NRT_MemSys_disable_stats(void);
+
+/*
+ * Query whether the internal statistics counters are enabled.
+ * Returns 1 if they are, 0 if they are not.
+ */
+VISIBILITY_HIDDEN
+size_t NRT_MemSys_stats_enabled(void);
+
 /*
  * The following functions get internal statistics of the memory subsystem.
  */

diff --git a/numba/core/runtime/nrt.py b/numba/core/runtime/nrt.py
@@ -6,7 +6,7 @@
 
 from numba.core.compiler_lock import global_compiler_lock
 from numba.core.typing.typeof import typeof_impl
-from numba.core import types
+from numba.core import types, config
 from numba.core.runtime import _nrt_python as _nrt
 
 _nrt_mstats = namedtuple("nrt_mstats", ["alloc", "free", "mi_alloc", "mi_free"])
@@ -27,6 +27,10 @@ def initialize(self, ctx):
             # Already initialized
             return
 
+        # Switch stats on if the config requests them.
+        if config.NRT_STATS:
+            _nrt.memsys_enable_stats()
+
         # Register globals into the system
         for py_name in _nrt.c_helpers:
             if py_name.startswith("_"):

diff --git a/numba/tests/support.py b/numba/tests/support.py
@@ -40,6 +40,7 @@
 from numba.core.runtime import rtsys
 from numba.np import numpy_support
 from numba.pycc.platform import _external_compiler_ok
+from numba.core.runtime import _nrt_python as _nrt
 
 
 try:
@@ -775,6 +776,16 @@ def capture_cache_log():
             yield out
 
 
+class EnableNRTStatsMixin(object):
+    """Mixin to enable the NRT statistics counters."""
+
+    def setUp(self):
+        _nrt.memsys_enable_stats()
+
+    def tearDown(self):
+        _nrt.memsys_disable_stats()
+
+
 class MemoryLeak(object):
 
     __enable_leak_check = True
@@ -803,16 +814,16 @@ def disable_leak_check(self):
         self.__enable_leak_check = False
 
 
-class MemoryLeakMixin(MemoryLeak):
+class MemoryLeakMixin(EnableNRTStatsMixin, MemoryLeak):
 
     def setUp(self):
         super(MemoryLeakMixin, self).setUp()
         self.memory_leak_setup()
 
     def tearDown(self):
-        super(MemoryLeakMixin, self).tearDown()
         gc.collect()
         self.memory_leak_teardown()
+        super(MemoryLeakMixin, self).tearDown()
 
 
 @contextlib.contextmanager