python - 为什么在遍历 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)

输出:每个循环 4.81 毫秒 ± 72.2 微秒(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 与 MSVC)。您可能希望使用 clang 来匹配 numba 性能(例如，参见 SO-answer )

为了使编译器更容易优化，您应该将输入声明为连续数组，即 double[::1] arr(参见 this question 为什么它对矢量化很重要) ，使用 @cython.boundscheck(False)(使用选项 -a 可以看到黄色变少了)并且还添加了编译器标志(即 -O3、-march=native 或类似代码，具体取决于您的编译器以启用矢量化，请注意默认使用的构建标志，这可能会抑制某些优化，例如 -fwrapv )。最后，您可能想用 C 语言编写 working-horse-loop，使用正确的标志/编译器组合进行编译，并使用 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