对于这个简单的循环,为什么 cython 比 numba 慢得多?

Rap*_*ael 5 python gcc clang cython numba

我有一个简单的循环,仅对 numpy 数组的第二行求和。在 numba 我只需要做:

\n
from numba import njit\n@njit(\'float64(float64[:, ::1])\', fastmath=True)\n    def fast_sum_nb(array_2d):\n        s = 0.0\n        for i in range(array_2d.shape[1]):\n            s += array_2d[1, i]\n        return s\n
Run Code Online (Sandbox Code Playgroud)\n

如果我对代码进行计时,我会得到:

\n
In [3]: import numpy as np\nIn [4]: A = np.random.rand(2, 1000)\nIn [5]: %timeit fast_sum_nb(A)\n305 ns \xc2\xb1 7.81 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 1,000,000 loops each)\n
Run Code Online (Sandbox Code Playgroud)\n

要在 cython 中执行相同的操作,我需要首先 make setup.py ,其中包含:

\n
from setuptools import setup\nfrom Cython.Build import cythonize\nfrom setuptools.extension import Extension\n\next_modules = [\n    Extension(\n        \'test_sum\',\n        language=\'c\',\n        sources=[\'test.pyx\'],  # list of source files\n        extra_compile_args=[\'-Ofast\', \'-march=native\'],  # example extra compiler arguments\n    )\n]\n\nsetup(\n    name = "test module",\n    ext_modules = cythonize(ext_modules, compiler_directives={\'language_level\' : "3"})  \n)\n
Run Code Online (Sandbox Code Playgroud)\n

我有尽可能最高的编译选项。那么 cython 求和代码为:

\n
#cython: language_level=3\nfrom cython cimport boundscheck\nfrom cython cimport wraparound\n@boundscheck(False)\n@wraparound(False)\ndef fast_sum(double[:, ::1] arr):\n    cdef int i=0\n    cdef double s=0.0\n    for i in range(arr.shape[1]):\n        s += arr[1, i]\n    return s\n
Run Code Online (Sandbox Code Playgroud)\n

我用以下方法编译它:

\n
python setup.py build_ext --inplace\n
Run Code Online (Sandbox Code Playgroud)\n

如果我现在计时,我会得到:

\n
In [2]: import numpy as np\nIn [3]: A = np.random.rand(2, 1000)\nIn [4]: %timeit fast_sum(A)\n564 ns \xc2\xb1 1.25 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 1,000,000 loops each)\n
Run Code Online (Sandbox Code Playgroud)\n

为什么 cython 版本慢这么多?

\n
\n

来自 cython 的带注释的 C 代码如下所示:

\n

在此输入图像描述

\n

numba 生成的程序集似乎是:

\n
vaddpd  -96(%rsi,%rdx,8), %ymm0, %ymm0\nvaddpd  -64(%rsi,%rdx,8), %ymm1, %ymm1\nvaddpd  -32(%rsi,%rdx,8), %ymm2, %ymm2\nvaddpd  (%rsi,%rdx,8), %ymm3, %ymm3\naddq    $16, %rdx\ncmpq    %rdx, %rcx\njne .LBB0_5\nvaddpd  %ymm0, %ymm1, %ymm0\nvaddpd  %ymm0, %ymm2, %ymm0\nvaddpd  %ymm0, %ymm3, %ymm0\nvextractf128    $1, %ymm0, %xmm1\nvaddpd  %xmm1, %xmm0, %xmm0\nvpermilpd   $1, %xmm0, %xmm1\nvaddsd  %xmm1, %xmm0, %xmm0\n
Run Code Online (Sandbox Code Playgroud)\n

我不知道如何获取 cython 代码的程序集。它生成的 C 文件很大,并且该.so文件也会分解为一个大文件。

\n

如果我增加二维数组中的列数,这种速度差异仍然存在(实际上它会增加),因此它似乎不是调用开销问题。

\n

我在 Ubuntu 上使用 Cython 版本 0.29.35 和 numba 版本 0.57.0。

\n

Mr *_*ooz 8

看起来像

decorated py code  --Numba-->  LLVM IR  --LLVM-->  machine code 
Run Code Online (Sandbox Code Playgroud)

只是生成比

cython code  --Cython-->  C  --gcc-or-clang-->  machine code
Run Code Online (Sandbox Code Playgroud)

有三件事共同导致 Cython 在此基准测试中表现更差。

  1. Cython 开销(次要)

  2. Cython 生成的循环(主要)

  3. 使用哪个编译器(主要)

开销(原因 1)

Cython 显然有一些额外的开销。您可以通过传入形状 (1,0) 的数组来观察这一点。Numba 功能仍然快得多。这并不奇怪,因为 Cython 是更通用的工具,并且它会尝试对输入、错误处理等格外小心,即使它有些过头了。除非您使用非常小的输入多次调用此函数,否则这应该不重要。

循环展开(原因 2 和 3 一起)

根据更新问题中的反汇编(在此处复制),Numba + LLVM 正在创建一个很好的展开循环。请注意它如何使用 YMM0..YMM3,而不仅仅是一个向量寄存器。

vaddpd  -96(%rsi,%rdx,8), %ymm0, %ymm0
vaddpd  -64(%rsi,%rdx,8), %ymm1, %ymm1
vaddpd  -32(%rsi,%rdx,8), %ymm2, %ymm2
vaddpd  (%rsi,%rdx,8), %ymm3, %ymm3
addq    $16, %rdx
cmpq    %rdx, %rcx
jne .LBB0_5
vaddpd  %ymm0, %ymm1, %ymm0
vaddpd  %ymm0, %ymm2, %ymm0
vaddpd  %ymm0, %ymm3, %ymm0
vextractf128    $1, %ymm0, %xmm1
vaddpd  %xmm1, %xmm0, %xmm0
vpermilpd   $1, %xmm0, %xmm1
vaddsd  %xmm1, %xmm0, %xmm0
Run Code Online (Sandbox Code Playgroud)

相比之下,这是使用 gcc 时来自 Cython 的核心反编译循环。这里没有展开。

    do {
      uVar9 = uVar8 + 1;
      auVar15 = vaddpd_avx(auVar15,*(undefined (*) [32])(local_118 + local_d0 + uVar8 * 0x20));
      uVar8 = uVar9;
    } while (uVar9 < (ulong)local_108 >> 2);
Run Code Online (Sandbox Code Playgroud)

clang 输出的反编译类似,但性能恰好更差(请参阅下面的基准测试结果)。由于某种原因,clang 不想展开 Cython 的循环。

    do {
        auVar1._8_8_ = 0;
        auVar1._0_8_ = *(ulong *)(local_e8.data + lVar4 * 8 + local_e8.strides[0]);
        auVar5 = vaddsd_avx(auVar5,auVar1);
        lVar4 = (long)((int)lVar4 + 1);
    } while (lVar4 < local_e8.shape[1]);
Run Code Online (Sandbox Code Playgroud)

让 Cython 变得更快

有时让 Cython 生成超快代码可能很棘手,但幸运的是还有另一种选择:仅使用 Cython 作为 Python 和 C 之间的粘合剂。

尝试将其放入您的.pyx文件中:

cdef extern from "impl.h":
    double fast_sum_c(double[] x, size_t n) nogil

def fast_sum_cyc(double[:, ::1] arr):
    # The pointer retrieval is only safe 
    # because of the "::1" constraint.
    return fast_sum_c(&arr[1, 0], arr.shape[1])
Run Code Online (Sandbox Code Playgroud)

然后创建一个impl.h包含以下内容的文件:

#pragma once

double fast_sum_c(double const *x, size_t n) {
    double s = 0.0;
    for (size_t i = 0; i < n; ++i) {
        s += x[i];
    }
    return s;
}
Run Code Online (Sandbox Code Playgroud)

在我的机器上,输入数组为 (2,1000),timeit运行时间如下:

Compiler    Numba  OrigCy  CyAndC
----------  -----  ------  ------
LLVM/Clang  240ns   900ns   250ns
gcc         n/a     380ns   380ns
Run Code Online (Sandbox Code Playgroud)

观察结果:

  • Numba 的开销比 Cython 低一些。

  • Cython + Clang 在此基准测试中表现非常差。

  • ...但是一个薄的 Cython 包装器 + 手写 C 代码 + Clang 几乎和 Numba 一样好。

  • GCC 似乎对 Cython 生成的代码不那么敏感。我们使用 Cython 的代码和手写代码获得相同的速度。

clang这是编译版本中最重要的汇编片段fast_sum_c。毫不奇怪,它与 Numba 生成的非常相似(因为它们都使用 LLVM 作为后端):

58b0:       c5 fd 58 04 cf          vaddpd (%rdi,%rcx,8),%ymm0,%ymm0
58b5:       c5 f5 58 4c cf 20       vaddpd 0x20(%rdi,%rcx,8),%ymm1,%ymm1
58bb:       c5 ed 58 54 cf 40       vaddpd 0x40(%rdi,%rcx,8),%ymm2,%ymm2
58c1:       c5 e5 58 5c cf 60       vaddpd 0x60(%rdi,%rcx,8),%ymm3,%ymm3
58c7:       48 83 c1 10             add    $0x10,%rcx
58cb:       48 39 c8                cmp    %rcx,%rax
58ce:       75 e0                   jne    58b0 <fast_sum_c+0x40>
58d0:       c5 f5 58 c0             vaddpd %ymm0,%ymm1,%ymm0
58d4:       c5 ed 58 c0             vaddpd %ymm0,%ymm2,%ymm0
58d8:       c5 e5 58 c0             vaddpd %ymm0,%ymm3,%ymm0
58dc:       c4 e3 7d 19 c1 01       vextractf128 $0x1,%ymm0,%xmm1
58e2:       c5 f9 58 c1             vaddpd %xmm1,%xmm0,%xmm0
58e6:       c4 e3 79 05 c8 01       vpermilpd $0x1,%xmm0,%xmm1
58ec:       c5 fb 58 c1             vaddsd %xmm1,%xmm0,%xmm0
Run Code Online (Sandbox Code Playgroud)

笔记

  • 鼓励更多循环展开的编译器选项没有帮助。在“Making Cython Fast”中,我尝试向 gcc 添加各种编译指示和编译器标志,以鼓励它进行一些展开;没有任何帮助。此外,clang 从来没有遇到过我手写的 C 代码的问题,但也不想展开 Cython 生成的循环。

  • objdump -d test_sum.*.so产生反汇编。查找vaddpd说明有助于找到感兴趣的循环。

  • Ghidra可用于反编译代码。这对理解它有一点帮助。-g使用和编译扩展模块-gdwarf-4使得 Ghidra 的 DWARF 解码可以工作,在反编译中注入更多的元数据。

  • 对于这些测试,我使用了clang14.0.0 和gcc11.3.0。


Jér*_*ard 6

TL;DR:这个答案在 @MrFooz 的好答案的基础上添加了额外的细节,以便理解为什么 Cython 代码在 Clang 和 GCC 中都很慢。性能问题来自于3 个错过的优化的组合:一项来自 Clang,一项来自 GCC,一项来自 Cython ...


在引擎盖下

首先,Cython 生成一个 C 代码,其步幅在编译时未知。这是一个问题,因为编译器的自动向量化无法轻松地使用 SIMD 指令对代码进行向量化,因为理论上数组可能不连续(即使在实践中它始终是连续的)。因此,Clang 自动矢量化器无法优化循环(自动矢量化和展开)。GCC 优化器更聪明一点:它为步幅为 1(即连续数组)生成专门的向量化代码。这是生成的 Cython 代码:

  for (__pyx_t_3 = 0; __pyx_t_3 < __pyx_t_2; __pyx_t_3+=1) {
    __pyx_v_i = __pyx_t_3;
    __pyx_t_4 = 1;
    __pyx_t_5 = __pyx_v_i;
    __pyx_v_s = (__pyx_v_s + (*((double *) ( /* dim=1 */ ((char *) (((double *) ( /* dim=0 */ (__pyx_v_arr.data + __pyx_t_4 * __pyx_v_arr.strides[0]) )) + __pyx_t_5)) ))));
  }
Run Code Online (Sandbox Code Playgroud)

请注意,__pyx_v_arr.strides[0]它在编译时不会被 1 替换,而 Cython 应该知道数组是连续的。对于 Cython 错过的优化有一个解决方法:使用 1D 内存视图。

这是修改后的 Cython 代码:

#cython: language_level=3
from cython cimport boundscheck
from cython cimport wraparound
@boundscheck(False)
@wraparound(False)
def fast_sum(double[:, ::1] arr):
    cdef int i=0
    cdef double s=0.0
    cdef double[::1] line = arr[1]
    for i in range(arr.shape[1]):
        s += line[i]
    return s
Run Code Online (Sandbox Code Playgroud)

不幸的是,由于两个底层编译器问题,该代码并不更快......

默认情况下,GCC 还不会展开这样的循环。这是众所周知的长期被忽视的优化。事实上,这个特定的 C 代码甚至还存在一个未解决的问题-funroll-loops -fvariable-expansion-in-unroller尽管生成的代码仍然不完美,但使用编译标志有助于提高最终性能。

对于 Clang 来说,这是另一个错过的优化,导致代码速度变快。GCC 和 Clang 过去有几个未解决的问题,这些问题与在循环中使用不同大小的类型进行矢量化时错过自动矢量化有关(甚至对于 GCC 有符号/无符号)。要解决此问题,在使用双精度数组时应使用 64 位整数。这是最终修改后的 Cython 代码:

#cython: language_level=3
from cython cimport boundscheck
from cython cimport wraparound
@boundscheck(False)
@wraparound(False)
def fast_sum(double[:, ::1] arr):
    cdef long i=0
    cdef double s=0.0
    cdef double[::1] line = arr[1]
    for i in range(arr.shape[1]):
        s += line[i]
    return s
Run Code Online (Sandbox Code Playgroud)

请注意,Numba 默认使用 64 位整数(例如 for 循环迭代器和索引),而 Numba 使用 LLVM-Lite(基于 LLVM,如 Clang),因此这里不会发生此类问题。


基准

以下是我的机器在配备 i5-9600KF 处理器、GCC 12.2.0、Clang 14.0.6 和 Python 3.11 时的性能结果:

Initial code:
    Cython GCC:      389 ns
    Cython Clang:   1050 ns
    Numba:           232 ns

Modified code:
    Cython GCC:      276 ns
    Cython Clang:    242 ns
Run Code Online (Sandbox Code Playgroud)

Numba 和 Cython+Clang 之间的开销非常小,这是由于启动开销不同。一般来说,这么短的时间应该不是问题,因为不应该长时间从 CPython 调用 Cython/Numba 函数。在这种病态情况下,调用者函数也应该修改为使用 Cython/Numba。如果这是不可能的,Numba 和 Cython 都不会生成快速代码,因此应首选较低级别的 C 扩展。