Skip to content

Commit a73612b

Browse files
stuartarchibaldstuartarchibald
committed
Make the NRT stats counters optional.
As title. They can be switched on/off with env var: `NUMBA_NRT_STATS`. Testing is updated to ensure that they are on when running CI/conda build test phase. Fixes #8156
1 parent c316c5a commit a73612b

11 files changed

Lines changed: 150 additions & 18 deletions

File tree

‎buildscripts/condarecipe.local/run_test.bat‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
set NUMBA_DEVELOPER_MODE=1
22
set NUMBA_DISABLE_ERROR_MESSAGE_HIGHLIGHTING=1
33
set NUMBA_CAPTURED_ERRORS=new_style
4+
set NUMBA_NRT_STATS=1
45
set PYTHONFAULTHANDLER=1
56

67
@rem Check Numba executables are there

‎buildscripts/condarecipe.local/run_test.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ set -e
55
export NUMBA_DEVELOPER_MODE=1
66
export NUMBA_DISABLE_ERROR_MESSAGE_HIGHLIGHTING=1
77
export NUMBA_CAPTURED_ERRORS="new_style"
8+
export NUMBA_NRT_STATS=1
89
export PYTHONFAULTHANDLER=1
910

1011
# Disable NumPy dispatching to AVX512_SKX feature extensions if the chip is

‎buildscripts/incremental/test.cmd‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ set NUMBA_DEVELOPER_MODE=1
1616
set PYTHONFAULTHANDLER=1
1717
@rem enable new style error handling
1818
set NUMBA_CAPTURED_ERRORS=new_style
19+
@rem switch on NRT stats
20+
set NUMBA_NRT_STATS=1
1921

2022
@rem First check that the test discovery works
2123
python -m numba.tests.test_runtests

‎buildscripts/incremental/test.sh‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ export PYTHONFAULTHANDLER=1
2525

2626
# enable new style error handling
2727
export NUMBA_CAPTURED_ERRORS="new_style"
28+
# switch on NRT stats
29+
export NUMBA_NRT_STATS=1
30+
2831

2932
# Disable NumPy dispatching to AVX512_SKX feature extensions if the chip is
3033
# reported to support the feature and NumPy >= 1.22 as this results in the use

‎docs/source/reference/envvars.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,12 @@ These variables influence what is printed out during compilation of
148148
allocation and ``0xDE`` on deallocation, both to help with debugging memory
149149
leaks.
150150

151+
.. envvar:: NUMBA_NRT_STATS
152+
153+
If set to non-zero, enable the
154+
:ref:`Numba run time (NRT) <arch-numba-runtime>` statistics counters. These
155+
counters are enabled process wide on import of Numba and are atomic.
156+
151157
.. envvar:: NUMBA_DEBUGINFO
152158

153159
If set to non-zero, enable debug for the full application by setting

‎numba/core/config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,9 @@ def optional_str(x):
223223
# Enable debug prints in nrtdynmod and use of "safe" API functions
224224
DEBUG_NRT = _readenv("NUMBA_DEBUG_NRT", int, 0)
225225

226+
# Enable NRT statistics counters
227+
NRT_STATS = _readenv("NUMBA_NRT_STATS", int, 0)
228+
226229
# How many recently deserialized functions to retain regardless
227230
# of external references
228231
FUNCTION_CACHE_SIZE = _readenv("NUMBA_FUNCTION_CACHE_SIZE", int, 128)

‎numba/core/runtime/_nrt_pythonmod.c‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ memsys_get_stats_mi_free(PyObject *self, PyObject *args) {
3737
return PyLong_FromSize_t(NRT_MemSys_get_stats_mi_free());
3838
}
3939

40+
static PyObject *
41+
memsys_enable_stats(PyObject *self, PyObject *args) {
42+
NRT_MemSys_enable_stats();
43+
Py_RETURN_NONE;
44+
}
4045

4146
/*
4247
* Create a new MemInfo with a owner PyObject
@@ -95,6 +100,7 @@ static PyMethodDef ext_methods[] = {
95100
declmethod_noargs(memsys_get_stats_free),
96101
declmethod_noargs(memsys_get_stats_mi_alloc),
97102
declmethod_noargs(memsys_get_stats_mi_free),
103+
declmethod_noargs(memsys_enable_stats),
98104
declmethod(meminfo_new),
99105
declmethod(meminfo_alloc),
100106
declmethod(meminfo_alloc_safe),

‎numba/core/runtime/nrt.cpp‎

Lines changed: 71 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,13 @@ struct NRT_MemSys {
4848
/* Shutdown flag */
4949
int shutting;
5050
/* Stats */
51-
std::atomic_size_t stats_alloc, stats_free, stats_mi_alloc, stats_mi_free;
51+
struct {
52+
bool enabled;
53+
std::atomic_size_t alloc;
54+
std::atomic_size_t free;
55+
std::atomic_size_t mi_alloc;
56+
std::atomic_size_t mi_free;
57+
} stats;
5258
/* System allocation functions */
5359
struct {
5460
NRT_malloc_func malloc;
@@ -64,10 +70,12 @@ static NRT_MemSys TheMSys;
6470

6571
extern "C" void NRT_MemSys_init(void) {
6672
TheMSys.shutting = 0;
67-
TheMSys.stats_alloc = 0;
68-
TheMSys.stats_free = 0;
69-
TheMSys.stats_mi_alloc = 0;
70-
TheMSys.stats_mi_free = 0;
73+
// Stats are off by default, call NRT_MemSys_enable_stats to enable
74+
TheMSys.stats.enabled = false;
75+
TheMSys.stats.alloc = 0;
76+
TheMSys.stats.free = 0;
77+
TheMSys.stats.mi_alloc = 0;
78+
TheMSys.stats.mi_free = 0;
7179
/* Bind to libc allocator */
7280
TheMSys.allocator.malloc = malloc;
7381
TheMSys.allocator.realloc = realloc;
@@ -78,36 +86,68 @@ extern "C" void NRT_MemSys_shutdown(void) {
7886
TheMSys.shutting = 1;
7987
}
8088

89+
extern "C" void NRT_MemSys_enable_stats(void) {
90+
TheMSys.stats.enabled = true;
91+
}
92+
8193
extern "C" void NRT_MemSys_set_allocator(NRT_malloc_func malloc_func,
8294
NRT_realloc_func realloc_func,
8395
NRT_free_func free_func)
8496
{
97+
bool stats_cond = false;
98+
if (TheMSys.stats.enabled)
99+
{
100+
stats_cond = (TheMSys.stats.alloc != TheMSys.stats.free ||
101+
TheMSys.stats.mi_alloc != TheMSys.stats.mi_free);
102+
}
85103
if ((malloc_func != TheMSys.allocator.malloc ||
86104
realloc_func != TheMSys.allocator.realloc ||
87105
free_func != TheMSys.allocator.free) &&
88-
(TheMSys.stats_alloc != TheMSys.stats_free ||
89-
TheMSys.stats_mi_alloc != TheMSys.stats_mi_free)) {
106+
stats_cond) {
90107
nrt_fatal_error("cannot change allocator while blocks are allocated");
91108
}
92109
TheMSys.allocator.malloc = malloc_func;
93110
TheMSys.allocator.realloc = realloc_func;
94111
TheMSys.allocator.free = free_func;
95112
}
96113

114+
/* This value is used as a marker for "stats are disabled", it's ASCII "AAAA" */
115+
static size_t _DISABLED_STATS_VALUE = 0x41414141;
116+
97117
extern "C" size_t NRT_MemSys_get_stats_alloc() {
98-
return TheMSys.stats_alloc;
118+
if (TheMSys.stats.enabled)
119+
{
120+
return TheMSys.stats.alloc.load();
121+
} else {
122+
return _DISABLED_STATS_VALUE;
123+
}
99124
}
100125

101126
extern "C" size_t NRT_MemSys_get_stats_free() {
102-
return TheMSys.stats_free;
127+
if (TheMSys.stats.enabled)
128+
{
129+
return TheMSys.stats.free.load();
130+
} else {
131+
return _DISABLED_STATS_VALUE;
132+
}
103133
}
104134

105135
extern "C" size_t NRT_MemSys_get_stats_mi_alloc() {
106-
return TheMSys.stats_mi_alloc;
136+
if (TheMSys.stats.enabled)
137+
{
138+
return TheMSys.stats.mi_alloc.load();
139+
} else {
140+
return _DISABLED_STATS_VALUE;
141+
}
107142
}
108143

109144
extern "C" size_t NRT_MemSys_get_stats_mi_free() {
110-
return TheMSys.stats_mi_free;
145+
if (TheMSys.stats.enabled)
146+
{
147+
return TheMSys.stats.mi_free.load();
148+
} else {
149+
return _DISABLED_STATS_VALUE;
150+
}
111151
}
112152

113153
/*
@@ -126,7 +166,10 @@ extern "C" void NRT_MemInfo_init(NRT_MemInfo *mi,void *data, size_t size,
126166
mi->external_allocator = external_allocator;
127167
NRT_Debug(nrt_debug_print("NRT_MemInfo_init mi=%p external_allocator=%p\n", mi, external_allocator));
128168
/* Update stats */
129-
TheMSys.stats_mi_alloc++;
169+
if (TheMSys.stats.enabled)
170+
{
171+
TheMSys.stats.mi_alloc++;
172+
}
130173
}
131174

132175
NRT_MemInfo *NRT_MemInfo_new(void *data, size_t size,
@@ -319,15 +362,21 @@ extern "C" void NRT_dealloc(NRT_MemInfo *mi) {
319362
NRT_Debug(nrt_debug_print("NRT_dealloc meminfo: %p external_allocator: %p\n", mi, mi->external_allocator));
320363
if (mi->external_allocator) {
321364
mi->external_allocator->free(mi, mi->external_allocator->opaque_data);
322-
TheMSys.stats_free++;;
365+
if (TheMSys.stats.enabled)
366+
{
367+
TheMSys.stats.free++;
368+
}
323369
} else {
324370
NRT_Free(mi);
325371
}
326372
}
327373

328374
extern "C" void NRT_MemInfo_destroy(NRT_MemInfo *mi) {
329375
NRT_dealloc(mi);
330-
TheMSys.stats_mi_free++;
376+
if (TheMSys.stats.enabled)
377+
{
378+
TheMSys.stats.mi_free++;
379+
}
331380
}
332381

333382
extern "C" void NRT_MemInfo_acquire(NRT_MemInfo *mi) {
@@ -471,7 +520,10 @@ extern "C" void* NRT_Allocate_External(size_t size, NRT_ExternalAllocator *alloc
471520
ptr = TheMSys.allocator.malloc(size);
472521
NRT_Debug(nrt_debug_print("NRT_Allocate_External bytes=%zu ptr=%p\n", size, ptr));
473522
}
474-
TheMSys.stats_alloc++;
523+
if (TheMSys.stats.enabled)
524+
{
525+
TheMSys.stats.alloc++;
526+
}
475527
return ptr;
476528
}
477529

@@ -485,7 +537,10 @@ extern "C" void *NRT_Reallocate(void *ptr, size_t size) {
485537
extern "C" void NRT_Free(void *ptr) {
486538
NRT_Debug(nrt_debug_print("NRT_Free %p\n", ptr));
487539
TheMSys.allocator.free(ptr);
488-
TheMSys.stats_free++;;
540+
if (TheMSys.stats.enabled)
541+
{
542+
TheMSys.stats.free++;
543+
}
489544
}
490545

491546
/*

‎numba/core/runtime/nrt.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,13 @@ void NRT_MemSys_shutdown(void);
4848
VISIBILITY_HIDDEN
4949
void NRT_MemSys_set_allocator(NRT_malloc_func, NRT_realloc_func, NRT_free_func);
5050

51+
/*
52+
* Enable the internal statistics counters.
53+
*/
54+
VISIBILITY_HIDDEN
55+
void NRT_MemSys_enable_stats(void);
56+
57+
5158
/*
5259
* The following functions get internal statistics of the memory subsystem.
5360
*/

‎numba/core/runtime/nrt.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
from numba.core.compiler_lock import global_compiler_lock
88
from numba.core.typing.typeof import typeof_impl
9-
from numba.core import types
9+
from numba.core import types, config
1010
from numba.core.runtime import _nrt_python as _nrt
1111

1212
_nrt_mstats = namedtuple("nrt_mstats", ["alloc", "free", "mi_alloc", "mi_free"])
@@ -27,6 +27,10 @@ def initialize(self, ctx):
2727
# Already initialized
2828
return
2929

30+
# Switch stats on if the config requests them.
31+
if config.NRT_STATS:
32+
_nrt.memsys_enable_stats()
33+
3034
# Register globals into the system
3135
for py_name in _nrt.c_helpers:
3236
if py_name.startswith("_"):

0 commit comments

Comments
 (0)