为何Numba中默认参数用np.empty或全局数组性能骤降？

问题描述 #

主要问题 #

使用Numba的@nb.njit装饰器编写函数时，将可选参数设为任意dtype数组，直接设置默认数组时，发现np.empty或全局数组作为默认参数的速度比np.zeros/np.array慢约80倍，但这些函数在普通Python或非默认参数的njit函数中速度正常，原因是什么？

次要问题 #

在VSCode的Jupyter Notebook中，首次运行代码时性能差异如上述，但第二次运行时所有使用数组作为默认参数的函数速度都变慢，这是真实性能还是有其他原因？

实验代码 #

import numba as nb
import numpy as np

@nb.njit
def func_zeros(b=np.zeros((1,), dtype=np.bool_)):
    return b

@nb.njit
def func_array(b=np.array(0, dtype=np.bool_)):
    return b

@nb.njit
def func_empty(b=np.empty((1,), dtype=np.bool_)):
    return b

_global_array = np.zeros((1,), dtype=np.bool_)

@nb.njit
def func_global(b=_global_array):
    return b

@nb.njit
def func_None(b=None):
    if b is None:
        b = np.empty((1,), dtype=np.bool_)
    return b

# 预热函数
func_zeros(), func_array(), func_empty(), func_global(), func_None()
func_zeros(np.zeros((1,), dtype=np.bool_))
func_array(np.array(0, dtype=np.bool_))
func_empty(np.empty((1,), dtype=np.bool_))
func_global(np.zeros((1,), dtype=np.bool_))
func_None(np.empty((1,), dtype=np.bool_))

# 性能测试
%timeit -n 10000 func_zeros()
%timeit -n 10000 func_array()
%timeit -n 10000 func_empty()
%timeit -n 10000 func_global()
%timeit -n 10000 func_None()
%timeit -n 10000 func_zeros(np.zeros((1,), dtype=np.bool_))
%timeit -n 10000 func_array(np.array(0, dtype=np.bool_))
%timeit -n 10000 func_empty(np.empty((1,), dtype=np.bool_))
%timeit -n 10000 func_global(np.zeros((1,), dtype=np.bool_))
%timeit -n 10000 func_None(np.empty((1,), dtype=np.bool_))

首次运行（重启内核后）结果 #

296 ns ± 2.73 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
291 ns ± 5.45 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
24.4 µs ± 166 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each) #!!!
24.4 µs ± 262 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each) #!!!
357 ns ± 5.69 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
463 ns ± 3.96 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
532 ns ± 4.17 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
468 ns ± 5.55 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
472 ns ± 6.65 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
475 ns ± 4.68 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

问题解答 #

1. np.empty/全局数组作为默认参数速度慢的原因 #

核心是Numba对默认参数的常量折叠优化差异：

• np.zeros和np.array(0)生成的数组内容是确定的（全0），Numba在JIT编译阶段会将这些默认值直接折叠为编译期常量，调用时无需执行数组生成逻辑，直接返回预编译的常量数组，因此速度极快。
• np.empty生成的数组内存是未初始化的随机值，无法被Numba作为编译期常量折叠；全局数组是运行时创建的外部对象，同样无法被编译期优化。调用这类函数时，Numba每次都要从Python层获取默认数组对象，涉及Python与Numba编译代码的桥接开销，这部分开销远大于直接返回常量的开销，导致速度慢一个数量级。
• func_None模式中，np.empty在Numba编译的函数内部执行，完全在JIT代码中运行，避免了Python层的桥接开销，因此速度回归正常。

2. 第二次运行所有数组默认参数函数变慢的原因 #

这不是真实性能退化，而是Jupyter和Numba的缓存与对象状态变化导致的：

• 首次运行时，Numba对函数进行JIT编译，部分默认参数的常量折叠优化生效。但内核保持运行时，Numba的函数缓存保留编译后的代码，后续运行时默认数组对象可能因Python引用计数或内存管理发生细微变化，导致Numba无法复用之前的常量折叠优化，转而每次走Python层的对象获取逻辑，最终所有数组默认参数的函数性能下降。
• 这种现象是环境层面的缓存和对象状态变化导致的，并非Numba本身的性能问题，重启内核后性能会回到首次运行状态。

最佳实践 #

推荐使用b=None的模式（即func_None的写法），在函数内部初始化默认数组。这种写法既能避免Python层的桥接开销，保证类型灵活性，又不会受外部数组对象状态的影响。

内容的提问来源于stack exchange，提问作者Nin17