使用NumPy对uint16和uint64阵列求和时没有加速？ [英] No speedup when summing uint16 vs uint64 arrays with NumPy?
问题描述
我必须对相对较小的整数进行大量的运算(加法)，我开始考虑哪种数据类型在64位计算机上性能最好。
我确信，将4uint16加在一起将花费与uint64相同的时间，因为ALU只使用1uint64加法器就可以进行4uint16加法。(进位传播意味着这对于单个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操作都需要相同的时间(比Float操作多)，并且没有加速。这是因为NumPy的C实现没有完全优化，还是有一些我不知道的硬件限制？
推荐答案
np.sum并没有(真正)使用SIMD指令：没有SIMD指令用于整数，而标量SIMD指令用于浮点数！此外，Numpy默认情况下会将整数转换为较小整数类型的64位值，以避免溢出！
请注意，这可能不会反映所有的Numpy版本，因为正在进行为常用功能提供SIMD支持的工作(尚未发布的版本Numpy 1.22.0rc1继续这项长期工作)。此外，使用的编译器或处理器可能会显著影响结果。以下实验是在i5-9600KF处理器的Debian Linux上使用从PIP中检索到的Numpy进行的。
np.sum
对于浮点数，Numpy使用成对算法，这是众所周知的非常数值稳定的，同时相对较快。这可以在代码中看到，但也可以简单地使用分析器：TYPE_pairwise_sum是在运行时调用C函数来计算总和(其中TYPE是DOUBLE或FLOAT)。
对于整数，Numpy使用经典的朴素归约。在兼容AVX2的机器上调用的C函数是ULONG_add_avx2。如果类型不是np.int64，它还会令人惊讶地将项目转换为64位项目。
以下是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指令。它甚至不会展开循环。由于对累加器的数据依赖性，执行累加的add (%rdi),%rdx指令是这里的瓶颈。对于np.uint64(despite the name of the function isulong_add_avx2`)也可以看到相同的循环。
但是，np.uint32版本调用C函数_aligned_contig_cast_uint_to_ulong以便将整型项转换为更宽的类型。Numpy这样做是为了避免整数溢出。np.uint8和np.uint16类型也是如此(尽管函数名不同)。希望此函数使用SIMD指令(SSE)，但仍占用很大一部分执行时间(约占np.sum时间的30%)。
编辑：正如@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位版本比其他基于整数的版本速度更快。这是由于转换的开销造成的。由于标量循环和add指令对于8位、16位和32位整数都是同样快的(对于大多数64位主流平台来说应该是这样)，所以前两个指令是同样快的。整数实现比浮点实现慢，因为缺少(正确的)循环展开。
vaddss(对于np.float32)和vaddsd(对于np.float64)具有相同的延迟和吞吐量(至少在所有现代英特尔处理器上)。因此，两个实现的吞吐量是相同的，因为两个实现的循环相似(相同的展开)。
其他说明
这很遗憾np.sum没有充分利用SIMD指令，因为这会大大加快使用它的计算速度(特别是小整数)。
[更新]查看Numpy代码时，我发现代码没有矢量化，因为数组跨度是运行期，并且编译器不会生成跨度为1的专用版本。事实上，这可以在前面的汇编代码中部分看出：编译器使用指令add %r10, %rdi，因为%r10(目标数组的跨度)在编译时是未知的。目前还没有针对Numpy代码缩减这一特定情况的优化(这些函数相对通用)。这种情况在不久的将来可能会改变。
除了跨度问题，还有一个大问题使编译器很难自动向量化代码：浮点加法是not commutative, nor associative(除非使用了-ffast-math这样的标志)。
这篇关于使用NumPy对uint16和uint64阵列求和时没有加速？的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！
