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

Question

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

\n

我确信uint16将 4 加在一起所需的时间与 1 相同uint64，因为 ALU 可以uint16仅使用 1 个uint64加法器进行 4 次加法。（进位传播意味着这对于单个 64 位加法器来说并不容易，但这就是整数 SIMD 指令的工作原理。）

\n

显然情况并非如此：

\n

In [3]: data = np.random.rand(10000)\n\nIn [4]: int16 = data.astype(np.uint16)\n\nIn [5]: int64 = data.astype(np.uint64)\n\nIn [6]: int32 = data.astype(np.uint32)\n\nIn [7]: float32 = data.astype(np.float32)\n\nIn [8]: float64 = data.astype(np.float64)\n\nIn [9]: %timeit int16.sum()\n13.4 \xc2\xb5s \xc2\xb1 43.3 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n\nIn [10]: %timeit int32.sum()\n13.9 \xc2\xb5s \xc2\xb1 347 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n\nIn [11]: %timeit int64.sum()\n9.33 \xc2\xb5s \xc2\xb1 47.8 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n\nIn [12]: %timeit float32.sum()\n5.79 \xc2\xb5s \xc2\xb1 6.51 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n\nIn [13]: %timeit float64.sum()\n6 \xc2\xb5s \xc2\xb1 3.54 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100000 loops each)\n

\n

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

\n

Answer 1

TL;DR：我对 Numpy 1.21.1 进行了实验分析。实验结果表明，它np.sum（确实）没有使用 SIMD 指令：没有 SIMD 指令用于整数，标量 SIMD 指令用于浮点数！此外，Numpy默认将较小整数类型的整数转换为 64 位值，以避免溢出！

\n\n

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

\n\n

在引擎盖下np.sum

\n

对于浮点数，Numpy 使用成对算法，该算法在数值上相当稳定，同时速度相对较快。这可以在代码中看到，但也可以简单地使用探查器：TYPE_pairwise_sum调用 C 函数来在运行时计算总和（其中TYPE是DOUBLE或FLOAT）。

\n

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

\n

DOUBLE_pairwise_sum这是函数执行的汇编代码的热门部分

\n

  3,65 \xe2\x94\x82 a0:\xe2\x94\x8c\xe2\x94\x80\xe2\x86\x92add        $0x8,%rcx                  ; Loop iterator\n  3,60 \xe2\x94\x82    \xe2\x94\x82  prefetcht0 (%r8,%rax,1)               ; Prefetch data\n  9,46 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%rbp,1),%xmm1,%xmm1  ; Accumulate an item in xmm1[0]\n  4,65 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax),%xmm0,%xmm0         ; Same for xmm0[0]\n  6,91 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%rdx,1),%xmm4,%xmm4  ; Etc.\n  7,77 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%rdx,2),%xmm7,%xmm7\n  7,41 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%r10,1),%xmm2,%xmm2\n  7,27 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%rdx,4),%xmm6,%xmm6\n  6,80 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%r11,1),%xmm3,%xmm3\n  7,46 \xe2\x94\x82    \xe2\x94\x82  vaddsd     (%rax,%rbx,1),%xmm5,%xmm5\n  3,46 \xe2\x94\x82    \xe2\x94\x82  add        %r12,%rax                  ; Switch to the next items (x8)\n  0,13 \xe2\x94\x82    \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80cmp        %rcx,%r9                   ; Should the loop continue?\n  3,27 \xe2\x94\x82    \xe2\x94\x94\xe2\x94\x80\xe2\x94\x80jg         a0                         ; Jump to the beginning if so\n

\n

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

\n

对于np.uint32，这是生成的汇编代码的热门部分：

\n

  2,37 \xe2\x94\x82160:\xe2\x94\x8c\xe2\x94\x80\xe2\x86\x92add          $0x1,%rax     ; Loop iterator\n 95,95 \xe2\x94\x82    \xe2\x94\x82  add          (%rdi),%rdx   ; Accumulate the values in %rdx\n  0,06 \xe2\x94\x82    \xe2\x94\x82  add          %r10,%rdi     ; Switch to the next item\n       \xe2\x94\x82    \xe2\x94\x9c\xe2\x94\x80\xe2\x94\x80cmp          %rsi,%rax     ; Should the loop continue?\n  1,08 \xe2\x94\x82    \xe2\x94\x94\xe2\x94\x80\xe2\x94\x80jne          160           ; Jump to the beginning if so\n

\n

Numpy 显然不使用 SIMD 指令进行np.uint32类型处理。它甚至不展开循环。add (%rdi),%rdx由于数据对累加器的依赖性，执行累加的指令是这里的瓶颈。对于 np.uint64 ULONG_add_avx2`）可以看到相同的循环(despite the name of the function is。

\n

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

\n

编辑：正如 @user2357112supportsMonica 所指出的，可以显式指定dtype的参数。np.sum当它与dtype输入数组的 匹配时，不执行转换。这会导致执行时间更短，但代价是可能发生溢出。

\n\n

基准测试结果

\n

这是我机器上的结果：

\n

uint16: 7.17 \xc2\xb5s \xc2\xb1 80 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 20000 loops each)\nuint32: 7.11 \xc2\xb5s \xc2\xb1 12.3 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 20000 loops each)\nuint64: 5.05 \xc2\xb5s \xc2\xb1 8.57 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 20000 loops each)\nfloat32: 2.88 \xc2\xb5s \xc2\xb1 9.27 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 20000 loops each)\nfloat64: 3.06 \xc2\xb5s \xc2\xb1 10.6 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 20000 loops each)\n

\n

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

\n

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

\n

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

\n\n

补充笔记

\n

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

\n

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

\n

除了跨步问题之外，还有一个重要问题使得编译器很难自动向量化代码：浮点加法不是可交换的，也不是关联的（除非-ffast-math使用类似标志）。

\n