Numba 和 numpy 数组分配：为什么这么慢？

我最近使用 Cython 和 Numba 来加速进行数值模拟的 python 的小片段。起初，使用 numba 进行开发似乎更容易。然而，我发现很难理解 numba 何时会提供更好的性能，何时不会。

意外性能下降的一个例子是当我使用该函数时np.zeros()在编译函数中分配一个大数组。例如，考虑三个函数定义：

import numpy as np 
from numba import jit 

def pure_python(n):
    mat = np.zeros((n,n), dtype=np.double)
    # do something
    return mat.reshape((n**2))

@jit(nopython=True)
def pure_numba(n):
    mat = np.zeros((n,n), dtype=np.double)
    # do something
    return mat.reshape((n**2))

def mixed_numba1(n):
    return mixed_numba2(np.zeros((n,n)))

@jit(nopython=True)

def mixed_numba2(array):
    n = len(array)
    # do something
    return array.reshape((n,n))

# To compile 
pure_numba(10)
mixed_numba1(10)

自从#do something是空的，我不期望pure_numba功能变得更快。然而，我没想到性能会如此下降：

n=10000
%timeit x = pure_python(n)
%timeit x = pure_numba(n)
%timeit x = mixed_numba1(n)

我获得（Mac上的python 3.7.7，numba 0.48.0）

4.96 µs ± 65.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
344 ms ± 7.76 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
3.8 µs ± 30.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

在这里，当我使用该函数时，numba 代码要慢得多np.zeros()在已编译的函数内。当np.zeros()位于函数之外。

我在这里做错了什么，还是应该总是分配大数组，比如这些由 numba 编译的外部函数？

Update

这似乎与矩阵的延迟初始化有关np.zeros((n,n)) when n足够大（参见Numpy 中的 Zeros 函数的性能 https://stackoverflow.com/questions/44487786/performance-of-zeros-function-in-numpy ).

for n in [1000, 2000, 5000]:
    print('n=',n)
    %timeit x = pure_python(n)
    %timeit x = pure_numba(n)
    %timeit x = mixed_numba1(n)

给我：

n = 1000
468 µs ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
296 µs ± 6.55 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
300 µs ± 2.26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
n = 2000
4.79 ms ± 182 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.45 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.54 ms ± 127 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
n = 5000
270 µs ± 4.66 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
104 ms ± 599 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
119 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

tl;dr Numpy 使用 C 内存函数，而 Numba 必须分配零

我编写了一个脚本来绘制几个选项完成所需的时间，当选项的大小变大时，Numba 的性能似乎会严重下降。np.zeros数组到达2048*2048*8 = 32 MB在我的机器上如下图所示。

Numba 的实施np.zeros与创建一个空数组并通过迭代数组的维度来用零填充它一样快（这是Numba 嵌套循环图中的绿色曲线）。实际上可以通过设置来双重检查NUMBA_DUMP_IR运行脚本之前的环境变量（见下文）。与转储进行比较时numba_loop没有太大区别。

有趣的是，np.zeros通过 32 MB 阈值获得了一点提升。

尽管我远非专家，但我的最佳猜测是 32 MB 限制是操作系统或硬件瓶颈，来自同一进程的缓存中可以容纳的数据量。如果超过这个值，将数据移入和移出缓存来对其进行操作的操作将非常耗时。

相比之下，Numpy 使用calloc https://github.com/numpy/numpy/blob/590facafcbd454fc02f96aa52c1c3412d6cccd14/numpy/core/src/multiarray/alloc.c#L212-L232获取一些内存段，并承诺在访问数据时用零填充数据。

这就是我所取得的进展，我意识到这只是答案的一半，但也许知识渊博的人可以阐明实际发生的情况。

Numba IR 转储：

---------------------------IR DUMP: pure_numba_zeros----------------------------
label 0:
    n = arg(0, name=n)                       ['n']
    $2load_global.0 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['$2load_global.0']
    $4load_attr.1 = getattr(value=$2load_global.0, attr=zeros) ['$2load_global.0', '$4load_attr.1']
    del $2load_global.0                      []
    $10build_tuple.4 = build_tuple(items=[Var(n, script.py:15), Var(n, script.py:15)]) ['$10build_tuple.4', 'n', 'n']
    $12load_global.5 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['$12load_global.5']
    $14load_attr.6 = getattr(value=$12load_global.5, attr=double) ['$12load_global.5', '$14load_attr.6']
    del $12load_global.5                     []
    $18call_function_kw.8 = call $4load_attr.1($10build_tuple.4, func=$4load_attr.1, args=[Var($10build_tuple.4, script.py:15)], kws=[('dtype', Var($14load_attr.6, script.py:15))], vararg=None) ['$10build_tuple.4', '$14load_attr.6', '$18call_function_kw.8', '$4load_attr.1']
    del $4load_attr.1                        []
    del $14load_attr.6                       []
    del $10build_tuple.4                     []
    mat = $18call_function_kw.8              ['$18call_function_kw.8', 'mat']
    del $18call_function_kw.8                []
    $24load_method.10 = getattr(value=mat, attr=reshape) ['$24load_method.10', 'mat']
    del mat                                  []
    $const28.12 = const(int, 2)              ['$const28.12']
    $30binary_power.13 = n ** $const28.12    ['$30binary_power.13', '$const28.12', 'n']
    del n                                    []
    del $const28.12                          []
    $32call_method.14 = call $24load_method.10($30binary_power.13, func=$24load_method.10, args=[Var($30binary_power.13, script.py:16)], kws=(), vararg=None) ['$24load_method.10', '$30binary_power.13', '$32call_method.14']
    del $30binary_power.13                   []
    del $24load_method.10                    []
    $34return_value.15 = cast(value=$32call_method.14) ['$32call_method.14', '$34return_value.15']
    del $32call_method.14                    []
    return $34return_value.15                ['$34return_value.15']

生成图表的脚本：

import numpy as np
from numba import jit
from time import time
import os
import matplotlib.pyplot as plt

os.environ['NUMBA_DUMP_IR'] = '1'

def numpy_zeros(n):
    mat = np.zeros((n,n), dtype=np.double)
    return mat.reshape((n**2))

@jit(nopython=True)
def numba_zeros(n):
    mat = np.zeros((n,n), dtype=np.double)
    return mat.reshape((n**2))

@jit(nopython=True)
def numba_loop(n):
    mat = np.empty((n * 2,n), dtype=np.float32)
    for i in range(mat.shape[0]):
        for j in range(mat.shape[1]):
            mat[i, j] = 0.
    return mat.reshape((2 * n**2))

# To compile
numba_zeros(10)
numba_loop(10)

os.environ['NUMBA_DUMP_IR'] = '0'

max_n = 4100
time_deltas = {
    'numpy_zeros': [],
    'numba_zeros': [],
    'numba_loop': [],
}
call_count = 10
for n in range(0, max_n, 10):
    for f in (numpy_zeros, numba_zeros, numba_loop):
        start = time()
        for i in range(call_count):
              x = f(n)
        delta = time() - start
        time_deltas[f.__name__].append(delta / call_count)
        print(f'{f.__name__:25} n = {n}: {delta}')
    print()

size = np.arange(0, max_n, 10) ** 2 * 8 / 1024 ** 2
fig, ax = plt.subplots()
plt.xticks(np.arange(0, size[-1], 16))
plt.axvline(x=32, color='gray', lw=0.5)
ax.plot(size, time_deltas['numpy_zeros'], label='Numpy zeros (calloc)')
ax.plot(size, time_deltas['numba_zeros'], label='Numba zeros')
ax.plot(size, time_deltas['numba_loop'], label='Numba nested loop')
ax.set_xlabel('Size of array in MB')
ax.set_ylabel(r'Mean $\Delta$t in s')
plt.legend(loc='upper left')
plt.show()