为什么类型列表的元素访问比使用 Numba 的数组慢得多？答案

【问题标题】：Why is element access for typed lists so much slower than for arrays with Numba?为什么类型列表的元素访问比使用 Numba 的数组慢得多？
【发布时间】：2021-12-12 17:52:18
【问题描述】：

我想知道为什么在使用 Numba 时，类型化列表的元素访问速度比 NumPy 数组慢得多。我有一个如下所示的最小示例。我随机生成索引以防止在后台进行任何编译器优化。似乎在纠正了生成所有这些随机数所需的时间之后，在 NumPy 数组的情况下，每个元素访问几乎都是即时发生的（


import numpy as np
import numba as nb

@nb.njit
def only_rand(N):
    
    for _ in range(10000):
        
        i = np.random.randint(N)
        j = np.random.randint(N)

@nb.njit
def foo(pos, N):
    
    for _ in range(10000):
        
        i = np.random.randint(N)
        j = np.random.randint(N)
        
        dx = pos[i][0] - pos[j][0]
        dy = pos[i][1] - pos[j][1]
        dz = pos[i][2] - pos[j][2]
        

N = 100

Array = np.random.rand(N,3)
List = nb.typed.List(Array)

print('Random number generation:')
%timeit only_rand(N)
print('Numpy Array:')
%timeit foo(Array, N)
print('Typed List:')
%timeit foo(List, N)

输出：

Random number generation:
133 µs ± 4.92 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Numpy Array:
132 µs ± 881 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Typed List:
947 µs ± 74 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

【问题讨论】：

标签： python arrays list performance numba

【解决方案1】：

一个问题是基准有缺陷。实际上，Numba JIT 编译器可以（部分地）看到您的计算大部分是无用的，因为它大部分没有计算可见影响：dx、dy 和dz 不会被读取，因此可以简单地忽略它们的计算（例如pos[i][0] - pos[j][0]）。乍一看似乎同样适用于i 和j，但事实并非如此：np.random.randint 修改内部种子会导致副作用。这种副作用迫使编译器仍然计算部分循环。

但是，除了上述几点之外，一旦基准修复，基于列表的实现确实会变慢。它来自临时列表的引用计数。以及汇编代码被 JIT 优化得不太好这一事实（列表往往会生成更复杂的代码，更难优化）。

深入分析：

要看到 JIT 优化了代码，你可以增加很多 N 的值。以下是我机器上的时间：

With N=100:
  Random number generation:
  109 µs ± 13.1 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
  Numpy Array:
  113 µs ± 22.7 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
  Typed List:
  806 µs ± 197 µs per loop (mean ± std. dev. of 7 runs, 250 loops each)

With N=1_000_000:
  Random number generation:
  64.7 µs ± 13.3 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
  Numpy Array:
  68.6 µs ± 23.1 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
  Typed List:
  804 µs ± 215 µs per loop (mean ± std. dev. of 7 runs, 250 loops each)

With N=10_000_000:
  Random number generation:
  185 µs ± 13.7 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
  Numpy Array:
  190 µs ± 23.2 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
  Typed List:
  839 µs ± 200 µs per loop (mean ± std. dev. of 7 runs, 250 loops each)

注意时间不是很依赖N。

这两种实现的汇编代码非常庞大，但可以看到主循环相似，并且在两种情况下都包含对numba_rnd_shuffle 的调用，这些调用未优化（由于np.random.randint 的副作用）。这是一个例子：

.LBB0_20: <----------\
        cmpl    $624, %eax
        jae     .LBB0_21
.LBB0_22: <----------\
        movl    %eax, %ecx
        movl    4(%rsi,%rcx,4), %ebp
        leal    1(%rax), %ecx
        movl    %ecx, (%rsi)
        movl    %ebp, %edx
        shrl    $11, %edx
        xorl    %ebp, %edx
        movl    %edx, %ebp
        shll    $7, %ebp
        andl    $-1658038656, %ebp
        xorl    %edx, %ebp
        movl    %ebp, %edx
        shll    $15, %edx
        andl    $-272236544, %edx
        xorl    %ebp, %edx
        movl    %edx, %ebp
        shrl    $18, %ebp
        xorl    %edx, %ebp
        andl    %edi, %ebp
        cmpl    $623, %eax
        jae     .LBB0_23
.LBB0_24: <----------\
        movl    %ecx, %eax
        movl    4(%rsi,%rax,4), %eax
        incl    %ecx
        movl    %ecx, (%rsi)
        movl    %eax, %edx
        shrl    $11, %edx
        xorl    %eax, %edx
        movl    %edx, %eax
        shll    $7, %eax
        andl    $-1658038656, %eax
        xorl    %edx, %eax
        movl    %eax, %edx
        shll    $15, %edx
        andl    $-272236544, %edx
        xorl    %eax, %edx
        movl    %edx, %eax
        shrl    $18, %eax
        xorl    %edx, %eax
        shlq    $32, %rbp
        orq     %rax, %rbp
        movl    %ecx, %eax
        cmpq    %r14, %rbp
        jge     .LBB0_20 ---------->
        jmp     .LBB0_12
.LBB0_21:
        movq    %rsi, %rcx
        movabsq $numba_rnd_shuffle, %rax
        callq   *%rax
        movl    $0, (%rsi)
        xorl    %eax, %eax
        jmp     .LBB0_22 ---------->
.LBB0_23:
        movq    %rsi, %rcx
        movabsq $numba_rnd_shuffle, %rax
        callq   *%rax
        movl    $0, (%rsi)
        xorl    %ecx, %ecx
        jmp     .LBB0_24 ---------->
        .p2align        4, 0x90

问题是在每次迭代结束时，下面的汇编代码会重复 6 次：

        movabsq $numba_list_size_address, %rdi

        movq    %r13, %rcx
        movabsq $NRT_incref, %rax
        callq   *%rax                  # NRT_incref(ptrVar);

        movq    %r15, %rcx
        callq   *%rdi                  # tmp1 = numba_list_size_address(listVar);

        movq    %rbp, %r12
        sarq    $63, %r12
        movq    (%rax), %rbx           # tmp2 = fancy_operation(*tmp1)
        andq    %r12, %rbx
        addq    %rbp, %rbx
        js      .LBB0_34               # Conditional goto to the end (overflow check?)

        movq    %r15, %rcx
        callq   *%rdi                  # tmp3 = numba_list_size_address(listVar);

        movq    %r15, %rdi
        movq    (%rax), %r15
        movq    %r13, %rcx
        movabsq $NRT_decref, %rax
        callq   *%rax                  # NRT_decref(ptrVar);

        cmpq    %r15, %rbx             # if(*tmp3 >= tmp2)
        jge     .LBB0_33               #     goto end;

        movq    %rdi, %rcx
        movabsq $numba_list_base_ptr, %rax
        callq   *%rax                  # numba_list_base_ptr(listVar);

可以看到引用计数调用以及与列表相关的函数。这部分汇编代码来自表达式pos[i] 和pos[j]。 JIT 未优化列表对象的引用计数。相关检查似乎也是如此。

我猜这是因为在这种情况下无法优化 Numba 函数调用，或者 JIT 认为它不够昂贵。列表相关函数的代码可以在here找到。我发现 JIT 不优化与列表相关的函数调用很奇怪，因为它们被 Numba 标记为 alwaysinline 和 readonly... 无论如何，我认为这是一个错过的优化和可以改进。

我向 Numba 开发者here提交了一个问题。

【讨论】：