numpy 为什么这两种对二维数组求和的方法会有如此不同的性能?

pwuypxnk  于 2023-03-30  发布在  其他
关注(0)|答案(1)|浏览(159)

考虑以下两种对2d numpy数组中的所有值求和的方法。

import numpy as np
from numba import njit
a = np.random.rand(2, 5000)

@njit(fastmath=True, cache=True)
def sum_array_slow(arr):
    s = 0
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            s += arr[i, j]
    return s
    
@njit(fastmath=True, cache=True)
def sum_array_fast(arr):
    s = 0
    for i in range(arr.shape[1]):
        s += arr[0, i]
    for i in range(arr.shape[1]):
        s += arr[1, i]
    return s

看看sum_array_slow中的嵌套循环,它似乎应该以与sum_array_fast相同的顺序执行完全相同的操作。然而:

In [46]: %timeit sum_array_slow(a)
7.7 µs ± 374 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [47]: %timeit sum_array_fast(a)
951 ns ± 2.63 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

为什么sum_array_fast函数的速度比sum_array_slow函数快8倍,而它似乎是以相同的顺序执行相同的计算?

cnjp1d6j

cnjp1d6j1#

这是因为慢版本是不自动向量化(即编译器无法生成快速SIMD代码),而快版本是。这肯定是因为Numba在第一次循环中没有优化索引 Package ,所以这是Numba的一个错过的优化
这可以通过分析汇编代码看出。下面是慢速版本的热循环:

.LBB0_6:
    addq    %rbx, %rdx
    vaddsd  (%rax,%rdx,8), %xmm0, %xmm0
    leaq    1(%rsi), %rdx
    cmpq    $1, %rbp
    cmovleq %r13, %rdx
    addq    %rbx, %rdx
    vaddsd  (%rax,%rdx,8), %xmm0, %xmm0
    leaq    2(%rsi), %rdx
    cmpq    $2, %rbp
    cmovleq %r13, %rdx
    addq    %rbx, %rdx
    vaddsd  (%rax,%rdx,8), %xmm0, %xmm0
    leaq    3(%rsi), %rdx
    cmpq    $3, %rbp
    cmovleq %r13, %rdx
    addq    $4, %rsi
    leaq    -4(%rbp), %rdi
    addq    %rbx, %rdx
    vaddsd  (%rax,%rdx,8), %xmm0, %xmm0
    cmpq    $4, %rbp
    movl    $0, %edx
    cmovgq  %rsi, %rdx
    movq    %rdi, %rbp
    cmpq    %rsi, %r12
    jne .LBB0_6

我们可以看到,Numba产生了许多无用的索引检查,这使得循环效率非常低。我不知道有任何干净的方法来解决这个问题。这是可悲的,因为这样的问题在实践中远非罕见。使用像C和C++这样的本地语言可以解决这个问题(因为数组中没有索引 Package )。一种不安全/丑陋的方法是在Numba中使用指针,但是提取Numpy数据指针并将其交给Numba似乎相当痛苦(如果可能的话)。
这是最快的一个:

.LBB0_8:
    vaddpd  (%r11,%rsi,8), %ymm0, %ymm0
    vaddpd  32(%r11,%rsi,8), %ymm1, %ymm1
    vaddpd  64(%r11,%rsi,8), %ymm2, %ymm2
    vaddpd  96(%r11,%rsi,8), %ymm3, %ymm3
    vaddpd  128(%r11,%rsi,8), %ymm0, %ymm0
    vaddpd  160(%r11,%rsi,8), %ymm1, %ymm1
    vaddpd  192(%r11,%rsi,8), %ymm2, %ymm2
    vaddpd  224(%r11,%rsi,8), %ymm3, %ymm3
    vaddpd  256(%r11,%rsi,8), %ymm0, %ymm0
    vaddpd  288(%r11,%rsi,8), %ymm1, %ymm1
    vaddpd  320(%r11,%rsi,8), %ymm2, %ymm2
    vaddpd  352(%r11,%rsi,8), %ymm3, %ymm3
    vaddpd  384(%r11,%rsi,8), %ymm0, %ymm0
    vaddpd  416(%r11,%rsi,8), %ymm1, %ymm1
    vaddpd  448(%r11,%rsi,8), %ymm2, %ymm2
    vaddpd  480(%r11,%rsi,8), %ymm3, %ymm3
    addq    $64, %rsi
    addq    $-4, %rdi
    jne .LBB0_8

在这种情况下,循环得到了很好的优化。事实上,它对于大型数组几乎是最优的。对于小型数组,就像你的例子一样,它在像我这样的处理器上不是最优的。事实上,AFAIK,展开的指令没有使用足够的寄存器以隐藏FMA单元的等待时间(这是因为LLVM在内部生成了一个次优代码)。可能需要一个较低级别的本机代码来解决这个问题(至少,在Numba中没有简单的方法来解决这个问题)。

更新

感谢@max9111提供的this link,可以通过使用无符号整数来优化缓慢的代码。这个技巧大大提高了执行时间。下面是修改后的代码:

@njit(fastmath=True, cache=True)
def sum_array_faster(arr):
    s = 0
    for i in range(np.uint64(arr.shape[0])):
        for j in range(np.uint64(arr.shape[1])):
            s += arr[i, j]
    return s

以下是英特尔至强W-2255处理器的性能:

slow:     9.66 µs
fastest:  1.13 µs
fast:     1.14 µs

Theoretical optimal:  0.30-0.35 µs

opt=0替换为opt=2的解决方案(再次感谢@max911)在我的机器上没有给出很好的结果:

slow:     2.12 µs
fastest:  2.17 µs
fast:     2.09 µs

更不用说编译的时间也稍微大一些。
可以实现更快的实现,以便更好地隐藏FMA指令的延迟:

@njit(fastmath=True, cache=True)
def sum_array_fastest(arr):
    s0, s1 = 0, 0
    for i in range(arr.shape[1]):
        s0 += arr[0, i]
        s1 += arr[1, i]
    return s0 + s1

这个需要1.08 µs,更好。
生成的Numba代码仍然有两个限制因素:

  • 与(短)执行时间相比,Numba的开销是显著的:250-300 ns
  • Numba没有使用我的机器上可用的AVX-512(zmm寄存器比AVX寄存器大两倍)。

请注意,可以使用Numba函数的方法inspect_asm提取汇编代码。

相关问题