Repository: numba/numba

Issue:
Performance of Numba Runtime Allocations
Hello,

I have tested the performance creating very short Numpy array in Numba  njit function and I am a bit surprised of the computation time. I found out that a significant part of the time is spent in NRT functions defined in `core/runtime/nrt.c`, and more specifically, in `memset` calls (in `NRT_MemInfo_alloc_safe_aligned` and `nrt_internal_dtor_safe`) that appear to only be useful for debugging purpose. Is it the case? If so, I propose to add a debug boolean parameter to the function so to enable this only in debug mode. A faster alternative solution is to add functions with the _debug prefix that can be called by the JIT only in debug mode. Such `memset` calls takes 20-25% of the time of very short Numpy allocations on my machine (with a i5-9600KF processor).

Additionally, I am wondering why such calls are not inlined by the JIT. Some calls like this could be faster if they were inlined because of constant propagation and the expression optimization of the provided function arguments. Is this behavior intended?

Finally, I found that `nrt_allocate_meminfo_and_data_align` performs a modulus that tends to be quite slow. Since the divider is the memory alignment which is generally a power of two, this modulus can be optimized. We can check if the alignment is a power of two with `if((align & (align - 1)) == 0)`, and if this is the case, then use a simple `remainder = intptr & (align - 1)`. Otherwise, a modulus can be used (rare case so the processor can predict the conditional). If the alignment is guaranteed to be always a power of two, then the conditional can even be removed! Note that this strategy is already used by most compilers for stack alignment of optimized programs and optimized memory allocators.

Here is a simple code to test this:

```python
import numba as nb
import numpy as np

@nb.njit('void()')
def test():
    for i in range(1000):
        np.empty(2)

%timeit  -n 20_000 test()
```


Discussion:
@zephyr111 thanks for the report.

> I have tested the performance creating very short Numpy array in Numba njit function and I am a bit surprised of the computation time. I found out that a significant part of the time is spent in NRT functions defined in `core/runtime/nrt.c`, and more specifically, in `memset` calls (in `NRT_MemInfo_alloc_safe_aligned` and `nrt_internal_dtor_safe`) that appear to only be useful for debugging purpose. Is it the case? If so, I propose to add a debug boolean parameter to the function so to enable this only in debug mode. A faster alternative solution is to add functions with the _debug prefix that can be called by the JIT only in debug mode. Such `memset` calls takes 20-25% of the time of very short Numpy allocations on my machine (with a i5-9600KF processor).

Running something like the above:
```python
import numba as nb
import numpy as np

@nb.njit('void()')
def test():
    for i in range(1000):
        np.empty(2)

n = 20000
for x in range(n):
    test()

```
Removing all those `memset` calls does lead to a faster alloc/dealloc loop, locally 4.1s with `memset`, 3.8s without `memset`. Looking at the machine code, seems like the `memset` is inlined, but it obviously still has a cost. Will raise this at the development meeting tomorrow. Is there a specific use case for a tight loop with small heap allocations (that couldn't just be reused) like in the above?

> Additionally, I am wondering why such calls are not inlined by the JIT. Some calls like this could be faster if they were inlined because of constant propagation and the expression optimization of the provided function arguments. Is this behavior intended?

The JIT is generating LLVM IR which is then compiled just-in-time by LLVM, this is C code compiled to a library, the symbols of which are called in the LLVM IR, so it can't be inlined.

> Finally, I found that nrt_allocate_meminfo_and_data_align performs a modulus that tends to be quite slow. Since the divider is the memory alignment which is generally a power of two, this modulus can be optimized. We can check if the alignment is a power of two with if((align & (align - 1)) == 0), and if this is the case, then use a simple remainder = intptr & (align - 1). Otherwise, a modulus can be used (rare case so the processor can predict the conditional). If the alignment is guaranteed to be always a power of two, then the conditional can even be removed! Note that this strategy is already used by most compilers for stack alignment of optimized programs and optimized memory allocators.

I think that this could be improved as suggested, a patch would be welcomed.
A related request that has come up from time to time (for a long time!) is the desire to be able to stack-allocate short, constant-sized NumPy arrays used for temporary storage in a loop.

The initial thought was to make this an explicit request, but the compiler would still need to verify that no reference to the temporary array escapes the function.  Then one could consider automatically turning regular array allocations into stack allocations when the required conditions are met.
@zephyr111 I wrote up some more information on this here: https://github.com/numba/numba/issues/8101#issuecomment-1135597413

Resolving patch:
diff --git a/numba/core/config.py b/numba/core/config.py
--- a/numba/core/config.py
+++ b/numba/core/config.py
@@ -220,7 +220,7 @@ def optional_str(x):
         # (up to and including IR generation)
         DEBUG_FRONTEND = _readenv("NUMBA_DEBUG_FRONTEND", int, 0)
 
-        # Enable debug prints in nrtdynmod
+        # Enable debug prints in nrtdynmod and use of "safe" API functions
         DEBUG_NRT = _readenv("NUMBA_DEBUG_NRT", int, 0)
 
         # How many recently deserialized functions to retain regardless
diff --git a/numba/core/runtime/context.py b/numba/core/runtime/context.py
--- a/numba/core/runtime/context.py
+++ b/numba/core/runtime/context.py
@@ -1,8 +1,24 @@
 import functools
+from collections import namedtuple
 
 from llvmlite import ir
+from numba.core import types, cgutils, errors, config
 
-from numba.core import types, cgutils, errors
+
+_NRT_Meminfo_Functions = namedtuple("_NRT_Meminfo_Functions",
+                                    ("alloc",
+                                     "alloc_dtor",
+                                     "alloc_aligned"))
+
+
+_NRT_MEMINFO_SAFE_API = _NRT_Meminfo_Functions("NRT_MemInfo_alloc_safe",
+                                               "NRT_MemInfo_alloc_dtor_safe",
+                                               "NRT_MemInfo_alloc_safe_aligned")
+
+
+_NRT_MEMINFO_DEFAULT_API = _NRT_Meminfo_Functions("NRT_MemInfo_alloc",
+                                                  "NRT_MemInfo_alloc_dtor",
+                                                  "NRT_MemInfo_alloc_aligned")
 
 
 class NRTContext(object):
@@ -13,6 +29,12 @@ class NRTContext(object):
     def __init__(self, context, enabled):
         self._context = context
         self._enabled = enabled
+        # If DEBUG_NRT is set, use the safe function variants which use memset
+        # to inject a few known bytes into the start of allocated regions.
+        if config.DEBUG_NRT:
+            self._meminfo_api = _NRT_MEMINFO_SAFE_API
+        else:
+            self._meminfo_api = _NRT_MEMINFO_DEFAULT_API
 
     def _require_nrt(self):
         if not self._enabled:
@@ -84,7 +106,8 @@ def meminfo_alloc_unchecked(self, builder, size):
 
         mod = builder.module
         fnty = ir.FunctionType(cgutils.voidptr_t, [cgutils.intp_t])
-        fn = cgutils.get_or_insert_function(mod, fnty, "NRT_MemInfo_alloc_safe")
+        fn = cgutils.get_or_insert_function(mod, fnty,
+                                            self._meminfo_api.alloc)
         fn.return_value.add_attribute("noalias")
         return builder.call(fn, [size])
 
@@ -116,7 +139,7 @@ def meminfo_alloc_dtor_unchecked(self, builder, size, dtor):
         fnty = ir.FunctionType(cgutils.voidptr_t,
                                [cgutils.intp_t, cgutils.voidptr_t])
         fn = cgutils.get_or_insert_function(mod, fnty,
-                                            "NRT_MemInfo_alloc_dtor_safe")
+                                            self._meminfo_api.alloc_dtor)
         fn.return_value.add_attribute("noalias")
         return builder.call(fn, [size,
                                  builder.bitcast(dtor, cgutils.voidptr_t)])
@@ -151,7 +174,7 @@ def meminfo_alloc_aligned_unchecked(self, builder, size, align):
         u32 = ir.IntType(32)
         fnty = ir.FunctionType(cgutils.voidptr_t, [cgutils.intp_t, u32])
         fn = cgutils.get_or_insert_function(mod, fnty,
-                                            "NRT_MemInfo_alloc_safe_aligned")
+                                            self._meminfo_api.alloc_aligned)
         fn.return_value.add_attribute("noalias")
         if isinstance(align, int):
             align = self._context.get_constant(types.uint32, align)


Test patch: