在迭代NumPy数组时,为什么Cython比Numba慢得多？

Question

迭代NumPy数组时,Numba似乎比Cython快得多.
我可能错过了什么Cython优化？

这是一个简单的例子:

纯Python代码:

import numpy as np

def f(arr):
  res=np.zeros(len(arr))

  for i in range(len(arr)):
     res[i]=(arr[i])**2

  return res

arr=np.random.rand(10000)
%timeit f(arr)

out:每循环4.81 ms±72.2μs(平均值±标准偏差,7次运行,每次100次循环)

---

Cython代码(在Jupyter中):

%load_ext cython
%%cython

import numpy as np
cimport numpy as np
cimport cython
from libc.math cimport pow

#@cython.boundscheck(False)
#@cython.wraparound(False)

cpdef f(double[:] arr):
   cdef np.ndarray[dtype=np.double_t, ndim=1] res
   res=np.zeros(len(arr),dtype=np.double)
   cdef double[:] res_view=res
   cdef int i

   for i in range(len(arr)):
      res_view[i]=pow(arr[i],2)

   return res

arr=np.random.rand(10000)
%timeit f(arr)

输出:每回路445μs±5.49μs(平均值±标准偏差,7次运行,每次1000次循环)

---

Numba代码:

import numpy as np
import numba as nb

@nb.jit(nb.float64[:](nb.float64[:]))
def   f(arr):
   res=np.zeros(len(arr))

   for i in range(len(arr)):
       res[i]=(arr[i])**2

   return res

arr=np.random.rand(10000)
%timeit f(arr)

输出:每个环路9.59μs±98.8 ns(平均值±标准偏差,7次运行,每次100000次循环)

---

在这个例子中,Numba几乎比Cython快50倍.
作为一名Cython初学者,我想我错过了一些东西.

当然,在这个简单的情况下,使用NumPy square矢量化函数会更合适:

%timeit np.square(arr)

输出:每循环5.75μs±78.9 ns(平均值±标准偏差,7次运行,每次100000次循环)

Answer 1

正如@Antonio所指出的那样,使用pow简单的乘法并不是很明智,并且导致相当大的开销:

因此,替换pow(arr[i], 2)通过arr[i]*arr[i]导致相当大的加速:

cython-pow-version        356 µs
numba-version              11 µs
cython-mult-version        14 µs

剩下的差异可能是由于编译器和优化级别之间的差异(在我的情况下为llvm vs MSVC).你可能想用clang来匹配numba性能(参见例如这个SO-answer)

为了使编译器的优化更容易,您应该将输入声明为连续数组,即double[::1] arr(请参阅此问题为什么它对于矢量化很重要),使用@cython.boundscheck(False)(使用选项-a查看黄色较少)并添加编译器标志(即-O3,-march=native或类似,取决于您的编译器启用矢量化,请注意默认使用的构建标志,这可能会禁止某些优化,例如-fwrapv).最后,您可能希望在C中编写working-horse-loop,使用flags/compiler的正确组合进行编译,并使用Cython来包装它.

顺便说一句,通过nb.float64[:](nb.float64[:])在降低numba的性能时输入函数的参数- 不再允许假设输入数组是连续的,从而统治了矢量化.让numba检测类型(或将其定义为连续类型,即nb.float64[::1](nb.float64[::1]),您将获得更好的性能:

@nb.jit(nopython=True)
def nb_vec_f(arr):
   res=np.zeros(len(arr))

   for i in range(len(arr)):
       res[i]=(arr[i])**2

   return res

导致以下改进:

%timeit f(arr)  # numba version
# 11.4 µs ± 137 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit nb_vec_f(arr)
# 7.03 µs ± 48.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

正如@ max9111所指出的那样,我们不必用零来初始化生成的数组,但可以使用np.empty(...)而不是np.zeros(...)- 这个版本甚至可以胜过numpy的np.square()

我机器上不同方法的表现如下:

numba+vectorization+empty     3µs
np.square                     4µs
numba+vectorization           7µs
numba missed vectorization   11µs
cython+mult                  14µs
cython+pow                  356µs