我有一个简单的循环,只对numpy数组的第二行求和。在Numba我只需要做:
from numba import njit
@njit('float64(float64[:, ::1])', fastmath=True)
def fast_sum_nb(array_2d):
s = 0.0
for i in range(array_2d.shape[1]):
s += array_2d[1, i]
return s
如果我计时我得到的代码:
In [3]: import numpy as np
In [4]: A = np.random.rand(2, 1000)
In [5]: %timeit fast_sum_nb(A)
305 ns ± 7.81 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
要在cython中做同样的事情,我需要首先创建setup.py,其中有:
from setuptools import setup
from Cython.Build import cythonize
from setuptools.extension import Extension
ext_modules = [
Extension(
'test_sum',
language='c',
sources=['test.pyx'], # list of source files
extra_compile_args=['-Ofast', '-march=native'], # example extra compiler arguments
)
]
setup(
name = "test module",
ext_modules = cythonize(ext_modules, compiler_directives={'language_level' : "3"})
)
我有最高可能的编译选项。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
for i in range(arr.shape[1]):
s += arr[1, i]
return s
我编译它:
python setup.py build_ext --inplace
如果我现在计时,我得到:
In [2]: import numpy as np
In [3]: A = np.random.rand(2, 1000)
In [4]: %timeit fast_sum(A)
564 ns ± 1.25 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
为什么Cython版本这么慢?
cython的注解C代码看起来像这样:
Numba制作的组装似乎是:
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
我不知道如何获取cython代码的程序集。它生成的C文件很大,.so
文件也会反汇编成一个大文件。
如果我增加2d数组中的列数,这种速度差异仍然存在(事实上它会增加),所以它似乎不是一个调用开销问题。
我在Ubuntu上使用Cython 0.29.35和numba 0.57.0。
2条答案
按热度按时间hujrc8aj1#
看起来像
就是产生比
在这个基准测试中,有三个因素共同给予Cython的性能更差。
1.正在使用哪个编译器(主要)
开销(原因1)
Cython显然有一些额外的开销。你可以通过传递一个形状(1,0)的数组来观察这一点。Numba功能仍然快得多。这并不奇怪,因为Cython是更通用的工具,并且它试图在输入、错误处理等方面格外小心,即使是在过度使用时也是如此。除非你用非常小的输入调用了很多这个函数,否则这应该没什么大不了的。
Loop Unrolling(原因2和原因3一起)
基于更新的问题(复制在这里)中的反汇编,Numba + LLVM正在创建一个很好地展开的循环。注意它是如何使用YMM0..YMM3的,而不仅仅是一个向量寄存器。
相比之下,这里是Cython使用gcc时的核心反编译循环。这里不能展开。
clang输出的反编译与此类似,但性能更差(参见下面的基准测试结果)。出于某种原因,clang不想展开Cython的循环。
让Cython快
有时候让Cython生成超快的代码是很棘手的,但幸运的是还有另一种选择:Cython仅用于Python和C之间的粘合。
尝试将其放入
.pyx
文件中:然后创建一个
impl.h
文件,内容如下:在我的机器上,有一个(2,1000)输入数组,下面是
timeit
运行时:观察结果:
下面是
fast_sum_c
的clang
编译版本中最重要的程序集代码段。不出所料,它与Numba产生的非常相似(因为它们都使用LLVM作为后端):注意事项
objdump -d test_sum.*.so
生成反汇编。查找vaddpd
指令有助于定位感兴趣的循环。-g
和-gdwarf-4
编译扩展模块使Ghidra的DWARF解码工作,在反编译中注入更多的元数据。clang
14.0.0和gcc
11.3.0。0x6upsns2#
TL;DR:这个答案在@MrFooz的好答案之上添加了额外的细节,以便理解为什么Cython代码在Clang和GCC上都很慢。性能问题来自以下3个未命中优化的组合:一个来自Clang,一个来自GCC,一个来自Cython...
在引擎盖下
首先,Cython生成了一个C代码,它的步幅在编译时是未知的。这是一个问题,因为编译器的自动向量化不能容易地使用SIMD指令向量化代码,因为阵列在理论上可能不是连续的(即使在实践中总是连续的)。因此,Clang自动向量化器无法优化循环(自动向量化和展开)。GCC优化器更聪明:它为步幅1(即,连续阵列)。下面是生成的Cython代码:
注意
__pyx_v_arr.strides[0]
在编译时没有被1替换,而Cython应该知道数组是连续的。有一个解决办法可以解决这个Cython遗漏的优化:使用1D内存视图。下面是修改后的Cython代码:
不幸的是,由于两个潜在的编译器问题,此代码并没有更快。
GCC默认不会展开这样的循环。这是一个众所周知的长期错过的优化。事实上,甚至还有an open issue for this specific C code。使用编译标志
-funroll-loops -fvariable-expansion-in-unroller
有助于提高结果的性能,尽管生成的代码仍然不完美。当涉及到Clang时,这是另一个错过的优化,防止代码快速。GCC和Clang在过去有几个公开的问题,当在循环中使用不同大小的**类型进行向量化(和even with signed/unsigned for GCC)时,会错过自动向量化。若要解决此问题,在使用双精度数组时应使用64位整数。下面是修改后的Cython代码:
请注意,Numba默认使用64位整数(例如。for loop iterators and indices)和Numba使用LLVM-Lite(基于LLVM,像Clang),所以这样的问题在这里不会发生。
基准测试
以下是我的机器上的性能结果,使用i5- 9600 KF处理器,GCC 12.2.0,Clang 14.0.6和Python 3.11:
Numba和Cython+Clang之间非常小的开销是由于不同的启动开销。一般来说,这么短的时间不应该是问题,因为不应该从CPython调用Cython/Numba函数太多。在这种病态的情况下,调用者函数也应该修改为使用Cython/Numba。当这是不可能的时候,Numba和Cython都不会产生快速的代码,所以应该首选低级的C扩展。