[Python] Why does this fast function with Numba JIT slow down if I JIT compile another...

So I have this function:

import numpy as np
import numba as nb


@nb.njit(cache=True, parallel=True, nogil=True)
def triangle_half_UR_LL(size: int, swap: bool = False) -> tuple[np.ndarray, np.ndarray]:
total = (size + 1) * size // 2
x_coords = np.full(total, 0, dtype=np.uint16)
y_coords = np.full(total, 0, dtype=np.uint16)
offset = 0
side = np.arange(size, dtype=np.uint16)
for i in nb.prange(size):
offset = i * size - (i - 1) * i // 2
end = offset + size - i
x_coords[offset:end] = i
y_coords[offset:end] = side[i:]

return (x_coords, y_coords) if not swap else (y_coords, x_coords)


What it does is not important, the point is it is JIT compiled with Numba and therefore very fast:

In [2]: triangle_half_UR_LL(10)
Out[2]:
(array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2,
2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5,
5, 6, 6, 6, 6, 7, 7, 7, 8, 8, 9], dtype=uint16),
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 2, 3, 4,
5, 6, 7, 8, 9, 3, 4, 5, 6, 7, 8, 9, 4, 5, 6, 7, 8, 9, 5, 6, 7, 8,
9, 6, 7, 8, 9, 7, 8, 9, 8, 9, 9], dtype=uint16))

In [3]: %timeit triangle_half_UR_LL(1000)
166 μs ± 489 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [4]: %timeit triangle_half_UR_LL(1000)
166 μs ± 270 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [5]: %timeit triangle_half_UR_LL(1000)
166 μs ± 506 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)


Now if I define another function and JIT compile it with Numba, the performance of the fast function inexplicably drops:

In [6]: @nb.njit(cache=True)
...: def dummy():
...: pass

In [7]: dummy()

In [8]: %timeit triangle_half_UR_LL(1000)
980 μs ± 20 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [9]: %timeit triangle_half_UR_LL(1000)
976 μs ± 9.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [10]: %timeit triangle_half_UR_LL(1000)
974 μs ± 3.11 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)


This is real, I have successfully reproduced this issue many times without fail, I start a new interpreter session, I paste the code, it runs fast. I define the dummy function then call the dummy function, and the fast function inexplicably slows down.

Screenshot as proof:



I am using Windows 11, and I absolutely have no idea what the hell is going on.

Is there an explanation for this? And how can I prevent this issue?


Interestingly if I get rid of nogil parameter and without changing anything else, the problem magically goes away:

In [1]: import numpy as np
...: import numba as nb
...:
...:
...: @nb.njit(cache=True, parallel=True)
...: def triangle_half_UR_LL(size: int, swap: bool = False) -> tuple[np.ndarray, np.ndarray]:
...: total = (size + 1) * size // 2
...: x_coords = np.full(total, 0, dtype=np.uint16)
...: y_coords = np.full(total, 0, dtype=np.uint16)
...: offset = 0
...: side = np.arange(size, dtype=np.uint16)
...: for i in nb.prange(size):
...: offset = i * size - (i - 1) * i // 2
...: end = offset + size - i
...: x_coords[offset:end] = i
...: y_coords[offset:end] = side[i:]
...:
...: return (x_coords, y_coords) if not swap else (y_coords, x_coords)

In [2]: %timeit triangle_half_UR_LL(1000)
186 μs ± 47.9 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: %timeit triangle_half_UR_LL(1000)
167 μs ± 1.61 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [4]: %timeit triangle_half_UR_LL(1000)
166 μs ± 109 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [5]: @nb.njit(cache=True)
...: def dummy():
...: pass

In [6]: dummy()

In [7]: %timeit triangle_half_UR_LL(1000)
167 μs ± 308 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [8]: %timeit triangle_half_UR_LL(1000)
166 μs ± 312 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [9]: %timeit triangle_half_UR_LL(1000)
167 μs ± 624 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)


Why does this happen?


But no, if I define other functions, somehow the first function slows down again. The simplest way to reproduce the issue is just redefining it:

In [7]: dummy()

In [8]: %timeit triangle_half_UR_LL(1000)
168 μs ± 750 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [9]: import numpy as np

In [10]: %timeit triangle_half_UR_LL(1000)
167 μs ± 958 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [11]: import numba as nb

In [12]: %timeit triangle_half_UR_LL(1000)
167 μs ± 311 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [13]: @nb.njit(cache=True, parallel=True)
...: def triangle_half_UR_LL(size: int, swap: bool = False) -> tuple[np.ndarray, np.ndarray]:
...: total = (size + 1) * size // 2
...: x_coords = np.full(total, 0, dtype=np.uint16)
...: y_coords = np.full(total, 0, dtype=np.uint16)
...: offset = 0
...: side = np.arange(size, dtype=np.uint16)
...: for i in nb.prange(size):
...: offset = i * size - (i - 1) * i // 2
...: end = offset + size - i
...: x_coords[offset:end] = i
...: y_coords[offset:end] = side[i:]
...:
...: return (x_coords, y_coords) if not swap else (y_coords, x_coords)

In [14]: %timeit triangle_half_UR_LL(1000)
1.01 ms ± 94.3 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [15]: %timeit triangle_half_UR_LL(1000)
964 μs ± 2.02 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)


The slowdown also happens if I define the following function and call it:

@nb.njit(cache=True)
def Farey_sequence(n: int) -> np.ndarray:
a, b, c, d = 0, 1, 1, n
result = [(a, b)]
while 0 <= c <= n:
k = (n + b) // d
a, b, c, d = c, d, k * c - a, k * d - b
result.append((a, b))

return np.array(result, dtype=np.uint64)


In [6]: %timeit triangle_half_UR_LL(1000)
166 μs ± 296 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

In [7]: %timeit Farey_sequence(16)
The slowest run took 6.25 times longer than the fastest. This could mean that an intermediate result is being cached.
6.03 μs ± 5.72 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [8]: %timeit Farey_sequence(16)
2.77 μs ± 50.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [9]: %timeit triangle_half_UR_LL(1000)
966 μs ± 6.48 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Continue reading...