为什么 numba 这么快？
Why is numba so fast?
我想编写一个函数，它将采用形状为 (N_ROWS,) 的索引 lefts 我想编写一个函数，它将创建一个矩阵 out = (N_ROWS, N_COLS) 矩阵，使得 out[i, j] = 1 当且仅当 j >= lefts[i]。这里有一个在循环中执行此操作的简单示例：
class Looped(Strategy):
def copy(self, lefts):
out = np.zeros([N_ROWS, N_COLS])
for k, l in enumerate(lefts):
out[k, l:] = 1
return out
现在我希望它尽可能快，所以我有这个函数的不同实现：
- 普通python循环
- numba 与
@njit
- 我用
ctypes 调用的纯 c++ 实现
以下是 100 次运行的平均值结果：
Looped took 0.0011599776260009093
Numba took 8.886413300206186e-05
CPP took 0.00013200821400096175
所以 numba 大约是下一个最快的实现（c++ 实现）的 1.5 倍。我的问题是为什么？
- 我在类似的问题中听说过 cython 速度较慢，因为它没有在设置所有优化标志的情况下进行编译，但是 cpp 实现是使用
-O3 编译的，这足以让我拥有所有可能的编译器会给我优化吗？
- 我不完全明白如何将 numpy 数组交给 c++，我是否无意中复制了这里的数据？
# numba implementation
@njit
def numba_copy(lefts):
out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)
for k, l in enumerate(lefts):
out[k, l:] = 1.
return out
class Numba(Strategy):
def __init__(self) -> None:
# avoid compilation time when timing
numba_copy(np.array([1]))
def copy(self, lefts):
return numba_copy(lefts)
// array copy cpp
extern "C" void copy(const long *lefts, float *outdatav, int n_rows, int n_cols)
{
for (int i = 0; i < n_rows; i++) {
for (int j = lefts[i]; j < n_cols; j++){
outdatav[i*n_cols + j] = 1.;
}
}
}
// compiled to a .so using g++ -O3 -shared -o array_copy.so array_copy.cpp
# using cpp implementation
class CPP(Strategy):
def __init__(self) -> None:
lib = ctypes.cdll.LoadLibrary("./array_copy.so")
fun = lib.copy
fun.restype = None
fun.argtypes = [
ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ctypes.c_long,
ctypes.c_long,
]
self.fun = fun
def copy(self, lefts):
outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )
self.fun(lefts, outdata, N_ROWS, N_COLS)
return outdata
带有时序等的完整代码：
import time
import ctypes
from itertools import combinations
import numpy as np
from numpy.ctypeslib import ndpointer
from numba import njit
N_ROWS = 1000
N_COLS = 1000
class Strategy:
def copy(self, lefts):
raise NotImplementedError
def __call__(self, lefts):
s = time.perf_counter()
n = 1000
for _ in range(n):
out = self.copy(lefts)
print(f"{type(self).__name__} took {(time.perf_counter() - s)/n}")
return out
class Looped(Strategy):
def copy(self, lefts):
out = np.zeros([N_ROWS, N_COLS])
for k, l in enumerate(lefts):
out[k, l:] = 1
return out
@njit
def numba_copy(lefts):
out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)
for k, l in enumerate(lefts):
out[k, l:] = 1.
return out
class Numba(Strategy):
def __init__(self) -> None:
numba_copy(np.array([1]))
def copy(self, lefts):
return numba_copy(lefts)
class CPP(Strategy):
def __init__(self) -> None:
lib = ctypes.cdll.LoadLibrary("./array_copy.so")
fun = lib.copy
fun.restype = None
fun.argtypes = [
ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ctypes.c_long,
ctypes.c_long,
]
self.fun = fun
def copy(self, lefts):
outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )
self.fun(lefts, outdata, N_ROWS, N_COLS)
return outdata
def copy_over(lefts):
strategies = [Looped(), Numba(), CPP()]
outs = []
for strategy in strategies:
o = strategy(lefts)
outs.append(o)
for s_0, s_1 in combinations(outs, 2):
for a, b in zip(s_0, s_1):
np.testing.assert_allclose(a, b)
if __name__ == "__main__":
copy_over(np.random.randint(0, N_COLS, size=N_ROWS))
Numba 目前使用 LLVM-Lite 将代码有效地编译为二进制文件（在 Python 代码被转换为 LLVM 中间表示之后）。代码经过优化，就像 C++ 代码使用带有标志 -O3 和 -march=native 的 Clang 一样。最后一个参数非常重要，因为它使 LLVM 能够在相对较新的 x86-64 处理器上使用 更宽的 SIMD 指令 ：AVX 和 AVX2（对于最新的 Intel 处理器可能是 AVX512）。否则，默认情况下 Clang 和 GCC 仅使用 SSE/SSE2 指令（因为 向后兼容性 ）。
另一个区别来自于 GCC 和 Numba 的 LLVM 代码之间的比较。 Clang/LLVM 倾向于主动展开循环，而 GCC 通常不会。这对生成的程序有显着的性能影响。事实上，你可以看到 generated assembly code from Clang:
使用 Clang（每个循环 128 个项目）：
.LBB0_7:
vmovups ymmword ptr [r9 + 4*r8 - 480], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 448], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 416], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 384], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 352], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 320], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 288], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 256], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 224], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 192], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 160], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 128], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 96], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 64], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 32], ymm0
vmovups ymmword ptr [r9 + 4*r8], ymm0
sub r8, -128
add rbp, 4
jne .LBB0_7
使用 GCC（每个循环 8 个项目）：
.L5:
mov rdx, rax
vmovups YMMWORD PTR [rax], ymm0
add rax, 32
cmp rdx, rcx
jne .L5
因此，为了公平起见，您需要将 Numba 代码与使用 Clang 和上述优化标志编译的 C++ 代码进行比较。
请注意，根据您的需求和最后一级处理器缓存的大小，您可以使用非临时存储（NT 存储）编写更快的特定于平台的 C++ 代码). NT-stores 告诉处理器不要将数组存储在它的缓存中。使用 NT-stores 写入数据在 RAM 中写入大数组的速度更快，但是如果数组适合缓存（因为必须从 RAM 重新加载数组），在复制后读取存储的数组时速度会变慢。在您的情况下（4 MiB 阵列），目前还不清楚这会更快。
结合其他 answers/comments 的所有建议，我在以下方面比 Numba 做得更好：
- 使用 cython + memoryiews（使用 ctypes 在某处有一些开销）
- 优化 cpp 实现
- 将 cython 编译器更改为 clang 并设置 -march=skylake
现在我有
CPP took 9.407872100018721e-05
Numba took 9.336918499957392e-05
Cythonised took 9.22323310005595e-05
// array_copy.cpp
#include "array.h"
const int n_rows = 1000;
const int n_cols = 1000;
void copy(const long *lefts, float *outdatav)
{
const float one = 1.;
for (int i = 0; i < n_rows; i++) {
const int l = lefts[i];
float* start = outdatav + i*n_cols + l;
std::fill(start, start + n_cols - l, one);
}
}
# setup.py
import os
os.environ["CXX"] = "clang"
os.environ["CC"] = "clang"
from setuptools import setup, Extension
from Cython.Build import cythonize
ex = Extension(
"array_copy_cython",
["./cpp_ext/array_copy_ext.pyx", "./cpp_ext/array.cpp" ],
include_dirs=["./cpp_ext"],
extra_compile_args=["-march=skylake"],
language="c++")
setup(
name='copy',
ext_modules=cythonize(ex),
zip_safe=False,
)
# array_copy_ext.pyx
cdef extern from "array.h":
void copy(const long* lefts, float* outdatav)
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
def copy_array(const long[:] lefts, float[:,:] outdatav):
copy(&lefts[0], &outdatav[0][0])
return outdatav
我想编写一个函数，它将采用形状为 (N_ROWS,) 的索引 lefts 我想编写一个函数，它将创建一个矩阵 out = (N_ROWS, N_COLS) 矩阵，使得 out[i, j] = 1 当且仅当 j >= lefts[i]。这里有一个在循环中执行此操作的简单示例：
class Looped(Strategy):
def copy(self, lefts):
out = np.zeros([N_ROWS, N_COLS])
for k, l in enumerate(lefts):
out[k, l:] = 1
return out
现在我希望它尽可能快，所以我有这个函数的不同实现：
- 普通python循环
- numba 与
@njit - 我用
ctypes 调用的纯 c++ 实现
以下是 100 次运行的平均值结果：
Looped took 0.0011599776260009093
Numba took 8.886413300206186e-05
CPP took 0.00013200821400096175
所以 numba 大约是下一个最快的实现（c++ 实现）的 1.5 倍。我的问题是为什么？
- 我在类似的问题中听说过 cython 速度较慢，因为它没有在设置所有优化标志的情况下进行编译，但是 cpp 实现是使用
-O3编译的，这足以让我拥有所有可能的编译器会给我优化吗？ - 我不完全明白如何将 numpy 数组交给 c++，我是否无意中复制了这里的数据？
# numba implementation
@njit
def numba_copy(lefts):
out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)
for k, l in enumerate(lefts):
out[k, l:] = 1.
return out
class Numba(Strategy):
def __init__(self) -> None:
# avoid compilation time when timing
numba_copy(np.array([1]))
def copy(self, lefts):
return numba_copy(lefts)
// array copy cpp
extern "C" void copy(const long *lefts, float *outdatav, int n_rows, int n_cols)
{
for (int i = 0; i < n_rows; i++) {
for (int j = lefts[i]; j < n_cols; j++){
outdatav[i*n_cols + j] = 1.;
}
}
}
// compiled to a .so using g++ -O3 -shared -o array_copy.so array_copy.cpp
# using cpp implementation
class CPP(Strategy):
def __init__(self) -> None:
lib = ctypes.cdll.LoadLibrary("./array_copy.so")
fun = lib.copy
fun.restype = None
fun.argtypes = [
ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ctypes.c_long,
ctypes.c_long,
]
self.fun = fun
def copy(self, lefts):
outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )
self.fun(lefts, outdata, N_ROWS, N_COLS)
return outdata
带有时序等的完整代码：
import time
import ctypes
from itertools import combinations
import numpy as np
from numpy.ctypeslib import ndpointer
from numba import njit
N_ROWS = 1000
N_COLS = 1000
class Strategy:
def copy(self, lefts):
raise NotImplementedError
def __call__(self, lefts):
s = time.perf_counter()
n = 1000
for _ in range(n):
out = self.copy(lefts)
print(f"{type(self).__name__} took {(time.perf_counter() - s)/n}")
return out
class Looped(Strategy):
def copy(self, lefts):
out = np.zeros([N_ROWS, N_COLS])
for k, l in enumerate(lefts):
out[k, l:] = 1
return out
@njit
def numba_copy(lefts):
out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)
for k, l in enumerate(lefts):
out[k, l:] = 1.
return out
class Numba(Strategy):
def __init__(self) -> None:
numba_copy(np.array([1]))
def copy(self, lefts):
return numba_copy(lefts)
class CPP(Strategy):
def __init__(self) -> None:
lib = ctypes.cdll.LoadLibrary("./array_copy.so")
fun = lib.copy
fun.restype = None
fun.argtypes = [
ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),
ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
ctypes.c_long,
ctypes.c_long,
]
self.fun = fun
def copy(self, lefts):
outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )
self.fun(lefts, outdata, N_ROWS, N_COLS)
return outdata
def copy_over(lefts):
strategies = [Looped(), Numba(), CPP()]
outs = []
for strategy in strategies:
o = strategy(lefts)
outs.append(o)
for s_0, s_1 in combinations(outs, 2):
for a, b in zip(s_0, s_1):
np.testing.assert_allclose(a, b)
if __name__ == "__main__":
copy_over(np.random.randint(0, N_COLS, size=N_ROWS))
Numba 目前使用 LLVM-Lite 将代码有效地编译为二进制文件（在 Python 代码被转换为 LLVM 中间表示之后）。代码经过优化，就像 C++ 代码使用带有标志 -O3 和 -march=native 的 Clang 一样。最后一个参数非常重要，因为它使 LLVM 能够在相对较新的 x86-64 处理器上使用 更宽的 SIMD 指令 ：AVX 和 AVX2（对于最新的 Intel 处理器可能是 AVX512）。否则，默认情况下 Clang 和 GCC 仅使用 SSE/SSE2 指令（因为 向后兼容性 ）。
另一个区别来自于 GCC 和 Numba 的 LLVM 代码之间的比较。 Clang/LLVM 倾向于主动展开循环，而 GCC 通常不会。这对生成的程序有显着的性能影响。事实上，你可以看到 generated assembly code from Clang:
使用 Clang（每个循环 128 个项目）：
.LBB0_7:
vmovups ymmword ptr [r9 + 4*r8 - 480], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 448], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 416], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 384], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 352], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 320], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 288], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 256], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 224], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 192], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 160], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 128], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 96], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 64], ymm0
vmovups ymmword ptr [r9 + 4*r8 - 32], ymm0
vmovups ymmword ptr [r9 + 4*r8], ymm0
sub r8, -128
add rbp, 4
jne .LBB0_7
使用 GCC（每个循环 8 个项目）：
.L5:
mov rdx, rax
vmovups YMMWORD PTR [rax], ymm0
add rax, 32
cmp rdx, rcx
jne .L5
因此，为了公平起见，您需要将 Numba 代码与使用 Clang 和上述优化标志编译的 C++ 代码进行比较。
请注意，根据您的需求和最后一级处理器缓存的大小，您可以使用非临时存储（NT 存储）编写更快的特定于平台的 C++ 代码). NT-stores 告诉处理器不要将数组存储在它的缓存中。使用 NT-stores 写入数据在 RAM 中写入大数组的速度更快，但是如果数组适合缓存（因为必须从 RAM 重新加载数组），在复制后读取存储的数组时速度会变慢。在您的情况下（4 MiB 阵列），目前还不清楚这会更快。
结合其他 answers/comments 的所有建议，我在以下方面比 Numba 做得更好：
- 使用 cython + memoryiews（使用 ctypes 在某处有一些开销）
- 优化 cpp 实现
- 将 cython 编译器更改为 clang 并设置 -march=skylake
现在我有
CPP took 9.407872100018721e-05
Numba took 9.336918499957392e-05
Cythonised took 9.22323310005595e-05
// array_copy.cpp
#include "array.h"
const int n_rows = 1000;
const int n_cols = 1000;
void copy(const long *lefts, float *outdatav)
{
const float one = 1.;
for (int i = 0; i < n_rows; i++) {
const int l = lefts[i];
float* start = outdatav + i*n_cols + l;
std::fill(start, start + n_cols - l, one);
}
}
# setup.py
import os
os.environ["CXX"] = "clang"
os.environ["CC"] = "clang"
from setuptools import setup, Extension
from Cython.Build import cythonize
ex = Extension(
"array_copy_cython",
["./cpp_ext/array_copy_ext.pyx", "./cpp_ext/array.cpp" ],
include_dirs=["./cpp_ext"],
extra_compile_args=["-march=skylake"],
language="c++")
setup(
name='copy',
ext_modules=cythonize(ex),
zip_safe=False,
)
# array_copy_ext.pyx
cdef extern from "array.h":
void copy(const long* lefts, float* outdatav)
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
def copy_array(const long[:] lefts, float[:,:] outdatav):
copy(&lefts[0], &outdatav[0][0])
return outdatav