python - numpy 比 numba 和 cython 快，如何改进 numba 代码

Question

我在这里有一个简单的例子来帮助我理解使用 numba 和 cython。我是 numba 和 cython 的新手。我已经尽力结合所有技巧来使 numba 更快，并且在某种程度上，cython 也是如此，但我的 numpy 代码几乎比 numba 快 2 倍(对于 float64)，如果使用 float32，则快 2 倍以上。不确定我在这里缺少什么。

我在想也许问题不再是编码，而是更多关于编译器之类的，我不太熟悉。

我浏览了很多关于 numpy、numba 和 cython 的 stackoverflow 帖子，但没有找到直接的答案。

numpy 版本:

def py_expsum(x):
    return np.sum( np.exp(x) )

麻麻版本:

@numba.jit( nopython=True)    
def nb_expsum(x):
    nx, ny = x.shape
    val = 0.0
    for ix in range(nx):
        for iy in range(ny):
            val += np.exp(x[ix, iy])
    return val

赛通版本:

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

@cython.boundscheck(False) 
@cython.wraparound(False)
cpdef double cy_expsum2 ( double[:,:] x, int nx, int ny ):
    cdef: 
        double val = 0.0
        int ix, iy    
    for ix in range(nx):
        for iy in range(ny):
            val += exp(x[ix, iy])
    return val

播放大小为 2000 x 1000 的数组并循环超过 100 次。对于 numba，它的第一次激活不计入循环。

使用 python 3(anaconda 发行版)，窗口 10

               float64       /   float32
    1. numpy : 0.56 sec      /   0.23 sec
    2. numba : 0.93 sec      /   0.74 sec      
    3. cython: 0.83 sec

cython 接近 numba。所以对我来说最大的问题是为什么 numba 不能击败 numpy 的运行时？我在这里做错了什么或错过了什么？其他因素如何起作用，我如何找出？

Answer 1

正如我们将看到的，行为取决于使用的 numpy-distribution。

这个答案将集中在 Anacoda-distribution 与 Intel 的 VML(矢量数学库)，考虑到另一个硬件和 numpy 版本，millage 可能会有所不同。

还将展示如何通过 Cython 或 numexpr 使用 VML。 ，以防万一不使用 Anacoda-distribution，它会在引擎盖下插入 VML 以进行一些 numpy 操作。

对于以下维度，我可以重现您的结果

N,M=2*10**4, 10**3
a=np.random.rand(N, M)

我得到:

%timeit py_expsum(a)  #   87ms
%timeit nb_expsum(a)  #  672ms
%timeit nb_expsum2(a)  #  412ms

大部分(约 90%)的计算时间用于评估 exp - 函数，正如我们将看到的，它是一项 CPU 密集型任务。

快速浏览 top -statistics 显示，numpy 的版本是并行执行的，但 numba 的情况并非如此。但是，在我只有两个处理器的 VM 上，仅靠并行化无法解释因子 7 的巨大差异(如 DavidW 的版本 nb_expsum2 所示)。

通过 perf 分析代码两个版本都显示以下内容:
nb_expsum

Overhead  Command  Shared Object                                      Symbol                                                             
  62,56%  python   libm-2.23.so                                       [.] __ieee754_exp_avx
  16,16%  python   libm-2.23.so                                       [.] __GI___exp
   5,25%  python   perf-28936.map                                     [.] 0x00007f1658d53213
   2,21%  python   mtrand.cpython-37m-x86_64-linux-gnu.so             [.] rk_random

py_expsum

  31,84%  python   libmkl_vml_avx.so                                  [.] mkl_vml_kernel_dExp_E9HAynn                                   ▒
   9,47%  python   libiomp5.so                                        [.] _INTERNAL_25_______src_kmp_barrier_cpp_38a91946::__kmp_wait_te▒
   6,21%  python   [unknown]                                          [k] 0xffffffff8140290c                                            ▒
   5,27%  python   mtrand.cpython-37m-x86_64-linux-gnu.so             [.] rk_random  

如您所见:numpy 在底层使用英特尔的并行矢量化 mkl/vml-version，它很容易胜过 numba(或并行版本的 numba 或 cython)使用的 gnu-math-library ( lm.so ) 中的版本那件事)。可以通过使用并行化稍微平整地面，但 mkl 的矢量化版本仍然优于 numba 和 cython。

但是，仅查看一种尺寸的性能并不是很有启发性，如果是 exp (至于其他超越函数)有两个维度需要考虑:

数组中的元素数量 - 缓存效果和不同大小的不同算法(在 numpy 中并非闻所未闻)会导致不同的性能。

取决于 x -value，需要不同的时间来计算exp(x) .通常有三种不同类型的输入导致不同的计算时间:非常小、正常和非常大(具有非有限结果)

我正在使用 perfplot 来可视化结果(请参阅附录中的代码)。对于“正常”范围，我们得到以下表现:



虽然 0.0 的性能相似，但我们可以看到，一旦结果变得无限，英特尔的 VML 就会受到相当大的负面影响:



然而，还有其他事情需要注意:

对于矢量大小 <= 8192 = 2^13 numpy 使用 exp 的非并行 glibc 版本(也使用相同的 numba 和 cython)。

我使用的 Anaconda-distribution，overrides numpy's functionality and plugs Intel's VML-library对于大于 8192 的大小，这是矢量化和并行化的 - 这解释了大约 10^4 大小的运行时间下降。

对于较小的尺寸，numba 很容易击败通常的 glibc 版本(对于 numpy 来说开销太大)，但是对于较大的数组，(如果 numpy 不会切换到 VML)没有太大区别。

这似乎是一个受 CPU 限制的任务——我们在任何地方都看不到缓存边界。

并行化的 numba 版本只有在元素超过 500 时才有意义。

那么后果是什么呢？

如果元素不超过 8192 个，则应使用 numba-version。

否则 numpy 版本(即使没有可用的 VML 插件，它也不会损失太多)。

注意:numba 不能自动使用 vdExp来自英特尔的 VML(如评论中部分建议的那样)，因为它计算 exp(x)单独，而 VML 对整个数组进行操作。

在写入和加载数据时，可以减少缓存未命中，这是由 numpy-version 使用以下算法执行的:

执行 VML vdExp 在适合缓存的一部分数据上，但也不太小(开销)。

总结得到的工作数组。

执行 1.+2。对于下一部分数据，直到处理完整个数据。

但是，与 numpy 的版本相比，我不希望获得超过 10%(但也许我错了)，因为无论如何 90% 的计算时间都花在 MVL 上。

尽管如此，这里是 Cython 中可能的快速和肮脏的实现:

%%cython -L=<path_mkl_libs> --link-args=-Wl,-rpath=<path_mkl_libs> --link-args=-Wl,--no-as-needed -l=mkl_intel_ilp64 -l=mkl_core -l=mkl_gnu_thread -l=iomp5
# path to mkl can be found via np.show_config()
# which libraries needed: https://software.intel.com/en-us/articles/intel-mkl-link-line-advisor

# another option would be to wrap mkl.h:
cdef extern from *:
    """
    // MKL_INT is 64bit integer for mkl-ilp64
    // see https://software.intel.com/en-us/mkl-developer-reference-c-c-datatypes-specific-to-intel-mkl
    #define MKL_INT long long int
    void  vdExp(MKL_INT n, const double *x, double *y);
    """
    void vdExp(long long int n, const double *x, double *y)

def cy_expsum(const double[:,:] v):
        cdef:
            double[1024] w;
            int n = v.size
            int current = 0;
            double res = 0.0
            int size = 0
            int i = 0
        while current<n:
            size = n-current
            if size>1024:
                size = 1024
            vdExp(size, &v[0,0]+current, w)
            for i in range(size):
                res+=w[i]
            current+=size
        return res

然而，它究竟是什么 numexpr会做，它也使用英特尔的 vml 作为后端:

 import numexpr as ne
 def ne_expsum(x):
     return ne.evaluate("sum(exp(x))")

至于时间，我们可以看到以下内容:



有以下值得注意的细节:

numpy、numexpr 和 cython 版本对于更大的数组具有几乎相同的性能 - 这并不奇怪，因为它们使用相同的 vml 功能。

在这三个中，cython-version 开销最少，numexpr 最多

numexpr-version 可能是最容易编写的(鉴于并非每个 numpy 分发插件 mvl 功能)。

房源:

情节:

import numpy as np
def py_expsum(x):
    return np.sum(np.exp(x))

import numba as nb
@nb.jit( nopython=True)    
def nb_expsum(x):
    nx, ny = x.shape
    val = 0.0
    for ix in range(nx):
        for iy in range(ny):
            val += np.exp( x[ix, iy] )
    return val

@nb.jit( nopython=True, parallel=True)    
def nb_expsum2(x):
    nx, ny = x.shape
    val = 0.0
    for ix in range(nx):
        for iy in nb.prange(ny):
            val += np.exp( x[ix, iy]   )
    return val

import perfplot
factor = 1.0 # 0.0 or 1e4
perfplot.show(
    setup=lambda n: factor*np.random.rand(1,n),
    n_range=[2**k for k in range(0,27)],
    kernels=[
        py_expsum, 
        nb_expsum,
        nb_expsum2, 
        ],
    logx=True,
    logy=True,
    xlabel='len(x)'
    )