如何使 Numba 访问数组的速度与 Numpy 一样快？
How can I make Numba access arrays as fast as Numpy can?
问题
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)
其他尝试过（但失败了）
- "stay longer in no-python mode";
函数中的虚拟循环
- 从环境中捕获
N（即没有N = len(a)）；
- 强制使用扁平数组：
b.flat[i] = a.flat[i]；
- 使用
fastmath=True and/or 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 模式下实现，因此使用小向量会变得非常慢。
I'm not expecting Numba to be faster but certainly not 70% slower!
根据我的经验，情况几乎总是如此，除非你愿意牺牲准确性（fastmath - 这里不相关，因为我们没有做任何数学运算）或者你可以利用 multi-threading （在这种情况下可能不值得，因为副本基本上是 memory-bandwidth 限制）或 multi-processing （正如另一个答案所示，这可以使更大的数组更快）。毕竟它是将 hand-written（通常是 highly-optimized）代码与 auto-generated 代码进行比较。这也应该回答问题标题：
How can I make Numba access arrays as fast as Numpy can?
使用 numba 不太可能！如果您尝试使用 numba re-implement 一些 native-NumPy 功能并且这已经非常接近了，那么在大阵列上通常会慢 50-200%。
当您需要编写对数组进行操作的代码时，不是已经在 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()
切换到并行版本可能会产生影响
看起来 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()
Windows, 桌面
这里看起来 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]
问题
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)
其他尝试过（但失败了）
- "stay longer in no-python mode"; 函数中的虚拟循环
- 从环境中捕获
N（即没有N = len(a)）； - 强制使用扁平数组：
b.flat[i] = a.flat[i]； - 使用
fastmath=Trueand/ornogil=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 模式下实现，因此使用小向量会变得非常慢。
I'm not expecting Numba to be faster but certainly not 70% slower!
根据我的经验，情况几乎总是如此，除非你愿意牺牲准确性（fastmath - 这里不相关，因为我们没有做任何数学运算）或者你可以利用 multi-threading （在这种情况下可能不值得，因为副本基本上是 memory-bandwidth 限制）或 multi-processing （正如另一个答案所示，这可以使更大的数组更快）。毕竟它是将 hand-written（通常是 highly-optimized）代码与 auto-generated 代码进行比较。这也应该回答问题标题：
How can I make Numba access arrays as fast as Numpy can?
使用 numba 不太可能！如果您尝试使用 numba re-implement 一些 native-NumPy 功能并且这已经非常接近了，那么在大阵列上通常会慢 50-200%。
当您需要编写对数组进行操作的代码时，不是已经在 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()
切换到并行版本可能会产生影响
看起来 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()
Windows, 桌面
这里看起来 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]