使用NumPy对uint16和uint64数组求和时没有加速？

提问者：小点点

使用NumPy对uint16和uint64数组求和时没有加速？

我必须对相对较小的整数进行大量操作（加法），我开始考虑哪种数据类型在 64 位机器上提供最佳性能。

我确信，将4个＜code＞uint16＜/code＞相加所需的时间与一个＜code＜uint64＜/code〕相加所需时间相同，因为ALU只需使用1个＜code＞uint64＞加法器即可进行4个＜code>uint16＞/code＞加法运算。（进位传播意味着这对于单个64位加法器来说不那么容易，但整数SIMD指令就是这样工作的。）

显然情况并非如此：

In [3]: data = np.random.rand(10000)

In [4]: int16 = data.astype(np.uint16)

In [5]: int64 = data.astype(np.uint64)

In [6]: int32 = data.astype(np.uint32)

In [7]: float32 = data.astype(np.float32)

In [8]: float64 = data.astype(np.float64)

In [9]: %timeit int16.sum()
13.4 µs ± 43.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [10]: %timeit int32.sum()
13.9 µs ± 347 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [11]: %timeit int64.sum()
9.33 µs ± 47.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [12]: %timeit float32.sum()
5.79 µs ± 6.51 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [13]: %timeit float64.sum()
6 µs ± 3.54 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

所有 int 操作都花费相同的时间（比浮点操作更多），没有加速。这是由于 numpy 的 C 实现没有完全优化还是存在一些我不知道的硬件限制？

共1个答案

匿名用户

TL；DR:我在Numpy 1.21.1上做了一个实验分析。实验结果表明< code>np.sum并没有(真正)利用SIMD指令:没有SIMD指令用于整数，标量SIMD指令用于浮点数！此外，默认情况下，Numpy将整数转换为较小整数类型的64位值，以避免溢出！

请注意，这可能不会反映所有的Numpy版本，因为目前正在为常用函数提供SIMD支持(尚未发布的版本Numpy 1.22.0rc1继续这项长期工作)。此外，使用的编译器或处理器可能会显著影响结果。下面的实验是使用从带有i5-9600KF处理器的Debian Linux上的pip中检索的Numpy完成的。

对于浮点数，Numpy使用一种成对算法，这种算法在速度相对较快的同时，在数值上非常稳定。这可以在代码中看到，也可以简单地使用一个分析器:< code>TYPE_pairwise_sum是在运行时被调用来计算总和的C函数(其中< code>TYPE是< code>DOUBLE或< code>FLOAT)。

对于整数，Numpy 使用经典的朴素约简。调用的 C 函数在 AVX2 兼容计算机上ULONG_add_avx2。如果类型不是 np.int64，它还会令人惊讶地将项目转换为 64 位项目。

下面是由< code>DOUBLE_pairwise_sum函数执行的汇编代码的热部分

  3,65 │ a0:┌─→add        $0x8,%rcx                  ; Loop iterator
  3,60 │    │  prefetcht0 (%r8,%rax,1)               ; Prefetch data
  9,46 │    │  vaddsd     (%rax,%rbp,1),%xmm1,%xmm1  ; Accumulate an item in xmm1[0]
  4,65 │    │  vaddsd     (%rax),%xmm0,%xmm0         ; Same for xmm0[0]
  6,91 │    │  vaddsd     (%rax,%rdx,1),%xmm4,%xmm4  ; Etc.
  7,77 │    │  vaddsd     (%rax,%rdx,2),%xmm7,%xmm7
  7,41 │    │  vaddsd     (%rax,%r10,1),%xmm2,%xmm2
  7,27 │    │  vaddsd     (%rax,%rdx,4),%xmm6,%xmm6
  6,80 │    │  vaddsd     (%rax,%r11,1),%xmm3,%xmm3
  7,46 │    │  vaddsd     (%rax,%rbx,1),%xmm5,%xmm5
  3,46 │    │  add        %r12,%rax                  ; Switch to the next items (x8)
  0,13 │    ├──cmp        %rcx,%r9                   ; Should the loop continue?
  3,27 │    └──jg         a0                         ; Jump to the beginning if so

Numpy编译的代码显然使用标量SIMD指令vaddsd（仅计算单个双精度项），尽管它成功地展开了8次循环。为FLOAT_pairwise_sum生成相同的代码：vaddss调用8次。

对于 np.uint32，下面是生成的程序集代码的热门部分：

  2,37 │160:┌─→add          $0x1,%rax     ; Loop iterator
 95,95 │    │  add          (%rdi),%rdx   ; Accumulate the values in %rdx
  0,06 │    │  add          %r10,%rdi     ; Switch to the next item
       │    ├──cmp          %rsi,%rax     ; Should the loop continue?
  1,08 │    └──jne          160           ; Jump to the beginning if so

Numpy显然不会对np.uint32类型使用SIMD指令。它甚至不会展开循环。由于对累加器的数据依赖性，执行累加的添加（%rdi），%rdx 指令是这里的瓶颈。在 np.uint64 中可以看到相同的循环（尽管函数的名称是ULONG_add_avx2'）。

然而，np.uint32版本调用C函数_aligned_contig_cast_uint_to_long，以便将整数项转换为更宽的类型。Numpy这样做是为了避免整数溢出。对于np.uint8和np.uint16类型，可以看到相同的情况（尽管函数的名称不同）。希望该函数使用了SIMD指令（SSE），但仍占用了很大一部分执行时间（约占np.sum时间的30%）。

EDIT：正如@user2357112supportsMonica所指出的，可以显式指定np.sum的dtype参数。当它与输入数组的dtype匹配时，则不执行转换。这会导致较小的执行时间，但可能会导致溢出。

以下是我的机器上的结果:

uint16: 7.17 µs ± 80 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)
uint32: 7.11 µs ± 12.3 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)
uint64: 5.05 µs ± 8.57 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)
float32: 2.88 µs ± 9.27 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)
float64: 3.06 µs ± 10.6 ns per loop (mean ± std. dev. of 7 runs, 20000 loops each)

首先，并不是说结果与问题中提供的结果非常相似，这意味着在我的机器上看到的行为可以在另一台机器上成功复制。因此，解释也应该是一致的。

如您所见，64位版本比其他基于整数的版本更快。这是由于转换的开销。前两个同样快，因为标量循环和< code>add指令对于8位、16位和32位整数同样快(对于大多数64位主流平台应该是这样)。整数实现比浮点实现慢，因为缺少(适当的)循环展开。

由于采用了标量 SIMD 指令，浮点实现速度同样快。事实上，指令 vaddss（对于 np.float32）和 vaddsd（对于 np.float64）具有相同的延迟和吞吐量（至少在所有现代英特尔处理器上）。因此，两个实现的吞吐量是相同的，因为两个实现的循环是相似的（相同的展开）。

遗憾的是，np.sum 没有充分利用 SIMD 指令，因为这会大大加快使用它的计算速度（尤其是小整数）。

[更新]查看 Numpy 代码，我发现代码没有矢量化，因为数组步幅是运行时值，编译器不会生成步幅为 1 的专用版本。实际上，这可以在前面的汇编代码中部分看到：编译器使用指令添加 %r10，%rdi，因为 %r10（目标数组的步幅）在编译时是未知的。目前还没有针对 Numpy 代码中这种减少的特定情况进行优化（函数相对通用）。这种情况在不久的将来可能会改变。

除了跨距问题之外，还有一个很大的问题使得编译器很难自动向量化代码:浮点加法不是可交换的，也不是关联的(除非使用了像< code>-ffast-math这样的标志)。