在 Cython 中循环二维数组的最快方法

Question

我正在尝试在 Cython 中循环 2 个 2d 数组。数组具有以下形状： ranges_1是 的 6000x3 数组int64，而ranges_2是 的 2000x2 数组int64。这个迭代需要执行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

但这根本没有提供任何加速。考虑到对于一次迭代，循环内只有 12_000_000 次操作，执行 10_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 = 120bln 的完整运行大约需要 6s，（虽然里面的函数更简单）这意味着它会比 Cython 快 500 倍，这可能吗？

我是 Cython 的新手，所以我可能错过了一些明显的东西，但由于 numba 版本（没有数组访问）要快得多，我认为速度的差异可能来自与访问数组中的项目相关的开销。

这是一个错误的假设吗？

如果不是，那么在 Cython 中循环整数的二维列表最好的方法是什么？

Answer 1

您在基准测试中测量的主要是编译工件和开销。

首先，Cython 使用 Python 堆栈首选的计算机上安装的默认编译器。在 Linux 上，应该是 GCC。在 Windows 上，如果已安装，则肯定是 MSVC，否则是 MinGW（如果有）。同时，Numba 基于 LLVM-Lite，LLVM-Lite 与 Clang 一样基于 LLVM 堆栈。因此，在您的情况下，很可能使用不同的编译器，导致不同的二进制文件具有不同的性能。如果你想做出一个公平的基准测试，你需要使用 Clang 来构建你的 Cython 程序。

此外，Cython 的默认优化是-O2Numba-O3的默认优化。前者不应启用自动矢量化，而后者则应启用（这取决于目标编译器——较新版本的 GCC 更改了此行为）。此外，Cython 默认情况下不启用特定于机器的不可移植优化（因为二进制文件可能像 pip 一样为其他机器打包）。这意味着 Cython 默认情况下只能在 x86-64 处理器上使用旧的 SSE2 SIMD 指令集。同时，LLVM-JIT 可以使用更快的 AVX2/AVX-512 SIMD 指令集。您需要使用 Cython 手动启用此类优化，以便基准测试公平（即-march=native在 GCC/Clang 上）。

事实上，在我的 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)：只有最后一次迭代才重要。变量i甚至for i in range(n)没有被使用。Clang 和 GCC 经常做这样的优化。

最后，如果修改它以计算现实用例中有意义的内容并且使用多个线程，则初始第一个代码的速度将受到内存限制。

请注意，除法非常昂贵，您可以预先计算倒数，以便在主循环中执行乘法。