在 Cython 中循环二维数组的最快方法
Fastest way to do loop over 2D arrays in Cython
我正在尝试在 Cython 中循环遍历 2 个二维数组。阵列具有以下形状：
ranges_1 是 int64 的 6000x3 数组，而 ranges_2 是 int64 的 2000x2 数组。此迭代需要执行大约 10000 次。这意味着嵌套 for 循环内的计算总数将约为 2000x6000x10000 = 1200 亿次。
这是我用来生成“虚拟”数据的代码：
import numpy as np
ranges_1 = np.stack([np.random.randint(0, 10_000, 6_000), np.random.randint(0, 10_000, 6_000), np.arange(0, 6_000)], axis=1)
ranges_2 = np.stack([np.random.randint(0, 10_000, 2_000), np.random.randint(0, 10_000, 2_000)], axis=1)
这给出了 2 个这样的数组：
array([[6131, 1478, 0],
[9317, 7263, 1],
[7938, 6249, 2],
...,
[5153, 426, 5997],
[9164, 9211, 5998],
[1695, 1792, 5999]])
和：
array([[ 433, 558],
[3420, 2494],
[6367, 7916],
...,
[8693, 1692],
[1256, 9013],
[4096, 1860]])
我尝试的第一个实现是一个“朴素”的版本，它如下（里面的函数只是一个测试函数，它使用数组中的所有数据）：
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_func(np.ndarray[DTYPE_t, ndim = 2] ranges_1, np.ndarray[DTYPE_t, ndim=2] ranges_2, int n ):
k = 0
for i in range(n):
for j in range(len(ranges_1)):
r1 = ranges_1[j]
a = r1[0]
b = r1[1]
c = r1[2]
for f in range(len(ranges_2)):
r2 = ranges_2[f]
d = r2[0]
e = r2[1]
k = (a + b + c + d + e)/(d+e)
return k
每个 10_000 外部循环大约需要 5 秒。
所以我然后尝试展平数组，因为我知道另一个轴上的维度，所以访问这样的项目：
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_func_flattened(np.ndarray[DTYPE_t, ndim = 1] ranges_1_, np.ndarray[DTYPE_t, ndim=1] ranges_2_, int n ):
k = 0
for i in range(n):
for j in range(0, len(ranges_1_), 3):
a = ranges_1_[j]
b = ranges_1_[j+1]
c = ranges_1_[j+2]
for f in range(0, len(ranges_2_), 2):
d = ranges_2_[f]
e = ranges_2_[f+1]
k = (a + b + c + d + e)/(d+e)
return k
但这根本没有提供任何加速。执行 10_000 的单次迭代的时间似乎太长，考虑到对于单次迭代它只是循环内的 12_000_000 操作。我还尝试在 Cython 和 python 中实现一个更简单的示例，然后用 numba:
编译
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_1(int n ):
cdef k = 0
cdef a = 0
for i in range(n):
a = i +1
return a
这花了 15 秒到 运行，n = 1_000_000_000。
同时 numba:
def test_1_python(n ):
k = 0
a = 0
for i in range(n):
if i % 2 == 0:
a = a + 1
else:
a = a - 1
return a
test_1_numba= numba.jit(test_1_python)
%%time
test_1_numba(120_000_000_000)
完整的运行 n = 12000亿只用了大约6s，（虽然里面的函数更简单）这意味着它比Cython快500倍，这可能吗？
我是 Cython 的新手，所以我可能遗漏了一些明显的东西，但是由于 numba 版本（没有数组访问）快得多，我认为速度上的差异可能来自与访问项目相关的开销在数组中。
这是一个错误的假设吗？
如果不是，那么在 Cython 中循环遍历二维整数列表的最佳方法是什么？
您在基准测试中衡量的主要是编译人工制品和开销。
首先，Cython 使用 Python 堆栈首选安装在您机器上的默认编译器。在Linux上，应该是GCC。在 Windows 上，如果安装了肯定是 MSVC，否则是 MinGW（如果有的话）。同时，Numba 是基于 LLVM-Lite 的，它像 Clang 一样基于 LLVM 堆栈。因此，在您的情况下，很可能使用 不同的编译器 导致不同的二进制文件具有不同的性能。如果你想做一个公平的基准测试，你需要使用 Clang 来构建你的 Cython 程序。
此外，Cython 的默认优化是 -O2 而 Numba 是 -O3。前者不应启用 auto-vectorization 而后者应启用（这取决于目标编译器——较新版本的 GCC 会更改此行为）。此外，Cython 默认不启用 machine-specific non-portable 优化（因为二进制文件可能会像 pip 一样为其他机器打包）。这意味着 Cython 只能在 x86-64 处理器上默认使用旧的 SSE2 SIMD 指令集。同时，LLVM-JIT 可以使用更快的 AVX2/AVX-512 SIMD 指令集。您需要使用 Cython 手动启用此类优化，以便基准测试公平（即 GCC/Clang 上的 -march=native）。
事实上，在我的 x86-64 主流 Intel 机器上，Numba 确实使用了 AVX2 指令集，而 Cython 在您上次的基准测试中没有使用。例如，这里是 Numba JIT 生成的主循环：
.LBB0_7:
vptestnmq %ymm5, %ymm4, %k1
vpblendmq %ymm5, %ymm6, %ymm18 {%k1}
vpaddq %ymm0, %ymm18, %ymm0
vpaddq %ymm1, %ymm18, %ymm1
vpaddq %ymm2, %ymm18, %ymm2
vpaddq %ymm3, %ymm18, %ymm3
vpaddq %ymm16, %ymm4, %ymm18
vptestnmq %ymm5, %ymm18, %k1
vpblendmq %ymm5, %ymm6, %ymm18 {%k1}
vpaddq %ymm0, %ymm18, %ymm0
vpaddq %ymm1, %ymm18, %ymm1
vpaddq %ymm2, %ymm18, %ymm2
vpaddq %ymm3, %ymm18, %ymm3
vpaddq %ymm17, %ymm4, %ymm4
addq $-2, %rdx
jne .LBB0_7
对于在循环中执行 a = i + 1 的基准测试，它是 有缺陷的 因为一个好的编译器可以优化整个循环（即删除它）并替换它只有一个赋值，因为 只有最后一次迭代很重要 。事实上，同样的事情也适用于 k = (a + b + c + d + e)/(d+e)：只有最后一次迭代很重要。 for i in range(n) 的 i 变量甚至没有使用。 Clang 和 GCC 经常做这样的优化。
最后，如果修改为在 real-world use-case和使用多线程。
请注意，除法非常昂贵，您可以预先计算倒数，以便在主循环中执行乘法。
我正在尝试在 Cython 中循环遍历 2 个二维数组。阵列具有以下形状：
ranges_1 是 int64 的 6000x3 数组，而 ranges_2 是 int64 的 2000x2 数组。此迭代需要执行大约 10000 次。这意味着嵌套 for 循环内的计算总数将约为 2000x6000x10000 = 1200 亿次。
这是我用来生成“虚拟”数据的代码：
import numpy as np
ranges_1 = np.stack([np.random.randint(0, 10_000, 6_000), np.random.randint(0, 10_000, 6_000), np.arange(0, 6_000)], axis=1)
ranges_2 = np.stack([np.random.randint(0, 10_000, 2_000), np.random.randint(0, 10_000, 2_000)], axis=1)
这给出了 2 个这样的数组：
array([[6131, 1478, 0],
[9317, 7263, 1],
[7938, 6249, 2],
...,
[5153, 426, 5997],
[9164, 9211, 5998],
[1695, 1792, 5999]])
和：
array([[ 433, 558],
[3420, 2494],
[6367, 7916],
...,
[8693, 1692],
[1256, 9013],
[4096, 1860]])
我尝试的第一个实现是一个“朴素”的版本，它如下（里面的函数只是一个测试函数，它使用数组中的所有数据）：
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_func(np.ndarray[DTYPE_t, ndim = 2] ranges_1, np.ndarray[DTYPE_t, ndim=2] ranges_2, int n ):
k = 0
for i in range(n):
for j in range(len(ranges_1)):
r1 = ranges_1[j]
a = r1[0]
b = r1[1]
c = r1[2]
for f in range(len(ranges_2)):
r2 = ranges_2[f]
d = r2[0]
e = r2[1]
k = (a + b + c + d + e)/(d+e)
return k
每个 10_000 外部循环大约需要 5 秒。 所以我然后尝试展平数组，因为我知道另一个轴上的维度，所以访问这样的项目：
import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_func_flattened(np.ndarray[DTYPE_t, ndim = 1] ranges_1_, np.ndarray[DTYPE_t, ndim=1] ranges_2_, int n ):
k = 0
for i in range(n):
for j in range(0, len(ranges_1_), 3):
a = ranges_1_[j]
b = ranges_1_[j+1]
c = ranges_1_[j+2]
for f in range(0, len(ranges_2_), 2):
d = ranges_2_[f]
e = ranges_2_[f+1]
k = (a + b + c + d + e)/(d+e)
return k
但这根本没有提供任何加速。执行 10_000 的单次迭代的时间似乎太长，考虑到对于单次迭代它只是循环内的 12_000_000 操作。我还尝试在 Cython 和 python 中实现一个更简单的示例，然后用 numba:
编译import numpy as np
cimport numpy as np
cimport cython
ctypedef np.int_t DTYPE_t
def test_1(int n ):
cdef k = 0
cdef a = 0
for i in range(n):
a = i +1
return a
这花了 15 秒到 运行，n = 1_000_000_000。
同时 numba:
def test_1_python(n ):
k = 0
a = 0
for i in range(n):
if i % 2 == 0:
a = a + 1
else:
a = a - 1
return a
test_1_numba= numba.jit(test_1_python)
%%time
test_1_numba(120_000_000_000)
完整的运行 n = 12000亿只用了大约6s，（虽然里面的函数更简单）这意味着它比Cython快500倍，这可能吗？
我是 Cython 的新手，所以我可能遗漏了一些明显的东西，但是由于 numba 版本（没有数组访问）快得多，我认为速度上的差异可能来自与访问项目相关的开销在数组中。
这是一个错误的假设吗？
如果不是，那么在 Cython 中循环遍历二维整数列表的最佳方法是什么？
您在基准测试中衡量的主要是编译人工制品和开销。
首先，Cython 使用 Python 堆栈首选安装在您机器上的默认编译器。在Linux上，应该是GCC。在 Windows 上，如果安装了肯定是 MSVC，否则是 MinGW（如果有的话）。同时，Numba 是基于 LLVM-Lite 的，它像 Clang 一样基于 LLVM 堆栈。因此，在您的情况下，很可能使用 不同的编译器 导致不同的二进制文件具有不同的性能。如果你想做一个公平的基准测试，你需要使用 Clang 来构建你的 Cython 程序。
此外，Cython 的默认优化是 -O2 而 Numba 是 -O3。前者不应启用 auto-vectorization 而后者应启用（这取决于目标编译器——较新版本的 GCC 会更改此行为）。此外，Cython 默认不启用 machine-specific non-portable 优化（因为二进制文件可能会像 pip 一样为其他机器打包）。这意味着 Cython 只能在 x86-64 处理器上默认使用旧的 SSE2 SIMD 指令集。同时，LLVM-JIT 可以使用更快的 AVX2/AVX-512 SIMD 指令集。您需要使用 Cython 手动启用此类优化，以便基准测试公平（即 GCC/Clang 上的 -march=native）。
事实上，在我的 x86-64 主流 Intel 机器上，Numba 确实使用了 AVX2 指令集，而 Cython 在您上次的基准测试中没有使用。例如，这里是 Numba JIT 生成的主循环：
.LBB0_7:
vptestnmq %ymm5, %ymm4, %k1
vpblendmq %ymm5, %ymm6, %ymm18 {%k1}
vpaddq %ymm0, %ymm18, %ymm0
vpaddq %ymm1, %ymm18, %ymm1
vpaddq %ymm2, %ymm18, %ymm2
vpaddq %ymm3, %ymm18, %ymm3
vpaddq %ymm16, %ymm4, %ymm18
vptestnmq %ymm5, %ymm18, %k1
vpblendmq %ymm5, %ymm6, %ymm18 {%k1}
vpaddq %ymm0, %ymm18, %ymm0
vpaddq %ymm1, %ymm18, %ymm1
vpaddq %ymm2, %ymm18, %ymm2
vpaddq %ymm3, %ymm18, %ymm3
vpaddq %ymm17, %ymm4, %ymm4
addq $-2, %rdx
jne .LBB0_7
对于在循环中执行 a = i + 1 的基准测试，它是 有缺陷的 因为一个好的编译器可以优化整个循环（即删除它）并替换它只有一个赋值，因为 只有最后一次迭代很重要 。事实上，同样的事情也适用于 k = (a + b + c + d + e)/(d+e)：只有最后一次迭代很重要。 for i in range(n) 的 i 变量甚至没有使用。 Clang 和 GCC 经常做这样的优化。
最后，如果修改为在 real-world use-case和使用多线程。
请注意，除法非常昂贵，您可以预先计算倒数，以便在主循环中执行乘法。