如何使用 numba 优化 numpy.packbits？

Question

我正在尝试优化numpy.packbits：

import numpy as np
from numba import njit, prange

@njit(parallel=True)
def _numba_pack(arr, div, su):
    for i in prange(div):
        s = 0
        for j in range(i*8, i*8+8):
            s = 2*s + arr[j]
        su[i] = s
        
def numba_packbits(arr):
    div, mod = np.divmod(arr.size, 8)
    su = np.zeros(div + (mod>0), dtype=np.uint8)
    _numba_pack(arr[:div*8], div, su)
    if mod > 0:
        su[-1] = sum(x*y for x,y in zip(arr[div*8:], (128, 64, 32, 16, 8, 4, 2, 1)))
    return su

>>> X = np.random.randint(2, size=99, dtype=bool)
>>> print(numba_packbits(X))
[ 75  24  79  61 209 189 203 187  47 226 170  61   0]

它看起来比np.packbits(X) 慢 2 - 2.5 倍。这是如何在numpy 内部实现的？这可以在numba 中改进吗？

我在通过conda install 安装的numpy == 1.21.2 和numba == 0.53.1 上工作。我的平台是：

结果：

import benchit
from numpy import packbits
%matplotlib inline
benchit.setparams(rep=5)

sizes = [100000, 300000, 1000000, 3000000, 10000000, 30000000]
N = sizes[-1]
arr = np.random.randint(2, size=N, dtype=bool)
fns = [numba_packbits, packbits]

in_ = {s/1000000: (arr[:s], ) for s in sizes}
t = benchit.timings(fns, in_, multivar=True, input_name='Millions of bits')
t.plot(logx=True, figsize=(12, 6), fontsize=14)

更新

Jérôme 的回应：

@njit('void(bool_[::1], uint8[::1], int_)', inline='never')
def _numba_pack_x64_byJérôme(arr, su, pos):
    for i in range(64):
        j = i * 8
        su[i] = (arr[j]<<7)|(arr[j+1]<<6)|(arr[j+2]<<5)|(arr[j+3]<<4)|(arr[j+4]<<3)|(arr[j+5]<<2)|(arr[j+6]<<1)|arr[j+7]
       
@njit(parallel=True)
def _numba_pack_byJérôme(arr, div, su):
    for i in prange(div//64):
        _numba_pack_x64_byJérôme(arr[i*8:(i+64)*8], su[i:i+64], i)
    for i in range(div//64*64, div):
        j = i * 8
        su[i] = (arr[j]<<7)|(arr[j+1]<<6)|(arr[j+2]<<5)|(arr[j+3]<<4)|(arr[j+4]<<3)|(arr[j+5]<<2)|(arr[j+6]<<1)|arr[j+7]
        
def numba_packbits_byJérôme(arr):
    div, mod = np.divmod(arr.size, 8)
    su = np.zeros(div + (mod>0), dtype=np.uint8)
    _numba_pack_byJérôme(arr[:div*8], div, su)
    if mod > 0:
        su[-1] = sum(x*y for x,y in zip(arr[div*8:], (128, 64, 32, 16, 8, 4, 2, 1)))
    return su

用法：

>>> print(numba_packbits_byJérôme(X))
[ 75  24  79  61 209 189 203 187  47 226 170  61   0]

结果：

Answer 1

Numba 实现存在几个问题。其中之一是并行循环破坏了 LLVM-Lite 中的常量传播优化（Numba 使用的 JIT 编译器）。这会导致诸如数组步长之类的关键信息无法传播，从而导致标量实现速度较慢，而不是 SIMD 实现，以及额外的非指令指令以计算偏移量。在 C 代码中也可以看到这样的问题。 Numpy 添加了特定的宏，以便在工作维度的步幅实际上为 1 时帮助编译器自动矢量化代码（即使用 SIMD 指令）。

克服持续传播问题的解决方案是调用另一个 Numba 函数。此函数必须不内联。应该手动提供签名，以便编译器可以在编译时知道数组的步长为 1 并生成更快的代码。最后，该函数应该在固定大小的块上工作，因为函数调用很昂贵并且编译器可以向量化代码。 使用移位展开循环也可以生成更快的代码（尽管它更丑陋）。这是一个例子：

@njit('void(bool_[::1], uint8[::1], int_)', inline='never')
def _numba_pack_x64(arr, su, pos):
    for i in range(64):
        j = i * 8
        su[i] = (arr[j]<<7)|(arr[j+1]<<6)|(arr[j+2]<<5)|(arr[j+3]<<4)|(arr[j+4]<<3)|(arr[j+5]<<2)|(arr[j+6]<<1)|arr[j+7]

@njit('void(bool_[::1], int_, uint8[::1])', parallel=True)
def _numba_pack(arr, div, su):
    for i in prange(div//64):
        _numba_pack_x64(arr[i*8:(i+64)*8], su[i:i+64], i)
    for i in range(div//64*64, div):
        j = i * 8
        su[i] = (arr[j]<<7)|(arr[j+1]<<6)|(arr[j+2]<<5)|(arr[j+3]<<4)|(arr[j+4]<<3)|(arr[j+5]<<2)|(arr[j+6]<<1)|arr[j+7]

---

基准测试

以下是在我的 6 核机器 (i5-9600KF) 上以十亿个随机项作为输入的性能结果：

Initial Numba (seq):    189 ms  (x0.7)
Initial Numba (par):    141 ms  (x1.0)
Numpy (seq):             98 ms  (x1.4)
Optimized Numba (par):   35 ms  (x4.0)
Theoretical optimal:     27 ms  (x5.2)  [fully memory-bound case]

这个新实现比最初的并行实现快 4 倍，比 Numpy 快大约 3 倍。

---

深入研究生成的汇编代码

当设置parallel=False 并将prange 替换为range 时，在我支持AVX-2 的Intel 处理器上生成以下汇编代码：

.LBB0_7:
    vmovdqu 112(%rdx,%rax,8), %xmm1
    vmovdqa 384(%rsp), %xmm3
    vpshufb %xmm3, %xmm1, %xmm0
    vmovdqu 96(%rdx,%rax,8), %xmm2
    vpshufb %xmm3, %xmm2, %xmm3
    vpunpcklwd  %xmm0, %xmm3, %xmm3
    vmovdqu 80(%rdx,%rax,8), %xmm15
    vmovdqa 368(%rsp), %xmm5
    vpshufb %xmm5, %xmm15, %xmm4
    vmovdqu 64(%rdx,%rax,8), %xmm0
    [...] <------------------------------  ~180 other instructions discarded
    vpcmpeqb    %xmm3, %xmm11, %xmm2
    vpandn  %xmm8, %xmm2, %xmm2
    vpor    %xmm2, %xmm1, %xmm1
    vpcmpeqb    %xmm3, %xmm0, %xmm0
    vpaddb  %xmm0, %xmm1, %xmm0
    vpsubb  %xmm4, %xmm0, %xmm0
    vmovdqu %xmm0, (%r11,%rax)
    addq    $16, %rax
    cmpq    %rax, %rsi
    jne .LBB0_7

代码不是很好，因为它使用了许多不需要的指令（例如 SIMD 比较指令，可能是由于布尔类型的隐式转换），很多寄存器是临时存储的（寄存器溢出），而且它使用 128 位 AVX 向量而不是我的机器上支持的 256 位 AVX。话虽如此，代码是矢量化的，每次循环迭代一次写入 16 字节，没有任何条件分支（循环之一除外），因此结果性能还不错。

事实上，Numpy 代码更小更高效。这就是为什么它比我的机器上输入大的顺序 Numba 代码快 2 倍左右。这是热组装循环：

4e8:
    mov      (%rdx,%rax,8),%rcx
    bswap    %rcx
    mov      %rcx,0x20(%rsp)
    mov      0x8(%rdx,%rax,8),%rcx
    add      $0x2,%rax
    movq     0x20(%rsp),%xmm0
    bswap    %rcx
    mov      %rcx,0x20(%rsp)
    movhps   0x20(%rsp),%xmm0
    pcmpeqb  %xmm1,%xmm0
    pcmpeqb  %xmm1,%xmm0
    pmovmskb %xmm0,%ecx
    mov      %cl,(%rsi)
    movzbl   %ch,%ecx
    mov      %cl,(%rsi,%r13,1)
    add      %r9,%rsi
    cmp      %rax,%r8
    jg       4e8

它按 8 字节块读取值，并使用 128 位 SSE 指令部分计算它们。每次迭代写入 2 个字节。话虽如此，它也不是最优的，因为没有使用 256 位 SIMD 指令，我认为代码可以进一步优化。

使用初始并行代码时，这里是热循环的汇编代码：

.LBB3_4:
     movq %r9, %rax
     leaq (%r10,%r14), %r9
     movq %r15, %rsi
     sarq $63, %rsi
     andq %rdx, %rsi
     addq %r11, %rsi
     cmpb $0, (%r14,%rsi)
     setne     %cl
     addb %cl, %cl
     [...] <---------------  56 instructions (with few 256-bit AVX ones)
     orb  %bl, %cl
     orb  %al, %cl
     orb  %dl, %cl
     movq %rbp, %rdx
     movb %cl, (%r8,%r15)
     incq %r15
     decq %rdi
     addq $8, %r14
     cmpq $1, %rdi
     jg   .LBB3_4

上面的代码主要是没有向量化，效率很低。它使用大量指令（包括相当慢的指令，如 setne/cmovlq/cmpb 来执行许多条件存储）​​，每次迭代仅写入 1 个字节。对于相同数量的写入字节，Numpy 执行的指令减少了大约 8 倍。使用多线程可以减轻此代码的低效率。最后，并行版本在具有许多内核（例如 >= 6）的机器上会更快一些。

此答案开头提供的改进实现会生成类似于上述顺序实现的代码，但使用多线程（仍然远未达到最佳状态，但还差很多）。