如何让 Numba 访问数组的速度与 Numpy 一样快？

Question

问题

Numpy 在将数组内容复制到另一个数组时基本上更快（或者看起来，对于足够大的数组）。

我不希望 Numba 更快，但几乎一样快似乎是一个合理的目标。

最小工作示例

import numpy as np
import numba as nb

def copyto_numpy(a, b):
    np.copyto(a, b, 'no')

@nb.jit(nopython=True)
def copyto_numba(a, b):
    N = len(a)
    for i in range(N):
        b[i] = a[i]

a = np.random.rand(2**20)
b = np.empty_like(a)
copyto_numpy(a, b)
copyto_numba(a, b)

%timeit copyto_numpy(a, b)
%timeit copyto_numba(a, b)

时间安排

1.28 ms ± 5.58 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
2.19 ms ± 222 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

尝试过的其他事情（但失败了）

• 函数中的虚拟循环“在无python模式下停留更长时间”；
• 从环境中捕获N （即 no N = len(a) ）；
• 强制使用扁平数组： b.flat[i] = a.flat[i] ;
• 使用fastmath=True和/或nogil=True以防它触发矢量化；
• 将循环计数器的类型改为无符号整数（32位和64位）；
• 删除循环并立即复制完整数组： b[:] = a （慢两倍！）；

• 带有 Cython 的新版本（与 Numba 一样快，比 Numpy 慢，与 @MSeifert 的答案略有不同）：

   %load_ext cython %%cython -c=-march=native -c=-O3 -c=-ftree-vectorize -c=-flto -c=-fuse-linker-plugin cimport cython @cython.boundscheck(False) # Deactivate bounds checking @cython.wraparound(False) # Deactivate negative indexing. cpdef copyto_cython_flags(double[::1] in_array, double[::1] out_array): cdef Py_ssize_t idx, N = len(in_array) for idx in range(N): out_array[idx] = in_array[idx]

• 检查 CPU 使用率以确认 Numpy 仅使用一个内核。

同样，我不希望 Numba 更快，但肯定不会慢 70%！ 我能做些什么来加快速度吗？ 请注意， np.copyto未在 nopython 模式下实现，因此使用小向量时它会变得非常慢。

---

Answer 1

我不希望 Numba 更快，但肯定不会慢 70%！

根据我的经验，情况几乎总是如此，除非您愿意牺牲准确性（ fastmath - 这里不相关，因为我们没有做任何数学运算）或者您可以利用多线程（在这种情况下可能不值得，因为复制基本上是内存带宽限制）或多处理（如另一个答案所示，这可以使更大的阵列更快）。 毕竟它是将手写（通常是高度优化的）代码与自动生成的代码进行比较。 这也应该回答问题标题：

如何让 Numba 访问数组的速度与 Numpy 一样快？

你不太可能用 numba ！ 如果您尝试使用 numba 重新实现一些本机 NumPy 功能，那么在大数组上通常会慢 50-200%，这已经非常接近了。

当你需要Numba擅长写上是不是已经在与NumPy，SciPy的，或者任何其他优化的库中实现阵列代码的运行。

---

然而 numba 已经比 Cython 的类似代码更快（我觉得很神奇）：

%load_ext cython

%%cython
cimport cython

@cython.boundscheck(False)  # Deactivate bounds checking
@cython.wraparound(False)   # Deactivate negative indexing.
cpdef copyto_cython(double[::1] in_array, double[::1] out_array):
    cdef Py_ssize_t idx
    for idx in range(len(in_array)):
        out_array[idx] = in_array[idx]

import numpy as np
import numba as nb

def copyto_numpy(a, b):
    np.copyto(a, b, 'no')

@nb.jit(nopython=True)
def copyto_numba(a, b):
    N = len(a)
    for i in range(N):
        b[i] = a[i]

我在这里使用我自己的库simple_benchmark进行性能测量：

from simple_benchmark import BenchmarkBuilder, MultiArgument

b = BenchmarkBuilder()

b.add_functions([copyto_cython, copyto_numpy, copyto_numba])

@b.add_arguments('array size')
def argument_provider():
    for exp in range(4, 21):
        size = 2**exp
        arr = np.random.rand(size)
        arr2 = np.empty(size)
        yield size, MultiArgument([arr, arr2])

r = b.run()
r.plot()

Answer 2

切换到并行化版本可能会有影响

看起来 Numba 在小型阵列上要快一点，而在中型阵列上则相当。 对于更大的数组，看起来 Numpy 正在切换到并行化副本，这必须在 Numba 中手动实现。 我不知道为什么 cython 性能明显变慢，但这可能是由于一些编译器标志。

代码

%load_ext cython

%%%%cython -c=-march=native -c=-O3 -c=-ftree-vectorize -c=-flto -c=-fuse-linker-plugin
cimport cython

@cython.boundscheck(False)  # Deactivate bounds checking
@cython.wraparound(False)   # Deactivate negative indexing.
cpdef copyto(double[::1] in_array, double[::1] out_array):
    cdef Py_ssize_t idx
    for idx in range(len(in_array)):
        out_array[idx] = in_array[idx]

import numpy as np
import numba as nb

def copyto_numpy(a, b):
    np.copyto(a, b, 'no')

@nb.jit(nopython=True,parallel=True)
def copyto_numba_p(a, b):
    N = len(a)
    for i in nb.prange(N):
        b[i] = a[i]

@nb.jit(nopython=True)
def copyto_numba_s(a, b):
    N = len(a)
    for i in nb.prange(N):
        b[i] = a[i]

@nb.jit(nopython=True)
def copyto_numba_combined(a, b):
    if a.shape[0]>4*10**4:
        copyto_numba_p(a, b)
    else:
        copyto_numba_s(a, b)

from simple_benchmark import BenchmarkBuilder, MultiArgument

b = BenchmarkBuilder()

b.add_functions([copyto, copyto_numpy, copyto_numba_combined,copyto_numba_s,copyto_numba_p])

@b.add_arguments('array size')
def argument_provider():
    for exp in range(4, 23):
        size = 2**exp
        arr = np.random.rand(size)
        arr2 = np.empty(size)
        yield size, MultiArgument([arr, arr2])

r = b.run()
r.plot()

窗户，桌面

这里看起来 numpy 没有在某个阈值切换到并行复制（这也可能是配置问题）。 使用 Numba，我手动实现了切换。

Linux，工作站

更新：并行 cython 版本

%%cython -c=-march=native -c=-O3 -c=-ftree-vectorize -c=-flto -c=-fuse-linker-plugin -c=-fopenmp

cimport cython
from cython.parallel import prange
@cython.boundscheck(False)  # Deactivate bounds checking
@cython.wraparound(False)   # Deactivate negative indexing.
cpdef copyto_p(double[::1] in_array, double[::1] out_array):
    cdef Py_ssize_t idx
    cdef Py_ssize_t size=len(in_array)
    for idx in prange(size,nogil=True):
        out_array[idx] = in_array[idx]