将uint8
的numpy.ndarray
传递给numpy.logical_and
时,如果我将numpy.view(bool)
应用于其输入,它的运行速度会明显加快。
a = np.random.randint(0, 255, 1000 * 1000 * 100, dtype=np.uint8)
b = np.random.randint(0, 255, 1000 * 1000 * 100, dtype=np.uint8)
%timeit np.logical_and(a, b)
126 ms ± 1.17 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit np.logical_and(a.view(bool), b.view(bool))
20.9 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
有人能解释为什么会发生这种情况吗?
此外,为什么numpy.logical_and
不自动将view(bool)
应用于uint8
的阵列?(有没有什么情况我们不应该使用view(bool)
?(
编辑:
这似乎是Windows环境的问题。我只是在官方的python docker容器(debian(中尝试了同样的东西,但没有发现它们之间的区别。
我的环境:
- 操作系统:Windows 10 Pro 21H2
- CPU:AMD Ryzen 9 5900X
- Python:Python 3.10.2(标签/v3.10.2:a58ebcc,2022年1月17日,14:12:15([MSC v.1929 64位(AMD64(]在win32上
- 数字:1.22.2
这是当前Numpy实现的性能问题。我也可以在Windows上重现这个问题(使用Numpy 1.20.3的Intel Skylake Xeon处理器(。np.logical_and(a, b)
基于慢速条件跳跃执行非常低效的标量汇编代码,而np.logical_and(a.view(bool), b.view(bool))
执行相对快速的SIMD指令。
目前,Numpy对bool
类型使用特定的实现。关于所使用的编译器,如果用于构建Numpy的编译器未能自动向量化代码,则通用实现可能会慢得多,这在Windows上显然是这样(并解释为什么在其他平台上不是这样,因为编译器可能不完全相同(Numpy代码可以针对非bool
类型进行改进。请注意,Numpy的矢量化是一项正在进行的工作,我们计划很快对此进行优化。
更深入的分析
以下是np.logical_and(a, b)
执行的汇编代码:
Block 24:
cmp byte ptr [r8], 0x0 ; Read a[i]
jz <Block 27> ; Jump to block 27 if a[i]!=0
Block 25:
cmp byte ptr [r9], 0x0 ; Read b[i]
jz <Block 27> ; Jump to block 27 if b[i]!=0
Block 26:
mov al, 0x1 ; al = 1
jmp <Block 28> ; Skip the next instruction
Block 27:
xor al, al ; al = 0
Block 28:
mov byte ptr [rdx], al ; result[i] = al
inc r8 ; i += 1
inc rdx
inc r9
sub rcx, 0x1
jnz <Block 24> ; Loop again while i<a.shape[0]
正如您所看到的,循环使用几个数据相关的条件跳转来写入a
和b
读取的每一项。这在这里是非常低效的,因为处理器无法用随机值预测所采取的分支。因此,处理器会停滞几个周期(在现代x86处理器上通常约为10个周期(。
以下是np.logical_and(a.view(bool), b.view(bool))
:执行的汇编代码
Block 15:
movdqu xmm1, xmmword ptr [r10] ; xmm1 = a[i:i+16]
movdqu xmm0, xmmword ptr [rbx+r10*1] ; xmm0 = b[i:i+16]
lea r10, ptr [r10+0x10] ; i += 16
pcmpeqb xmm1, xmm2 ;
pandn xmm1, xmm0 ; | Complex sequence to just do:
pcmpeqb xmm1, xmm2 ; | xmm1 &= xmm0
pandn xmm1, xmm3 ; /
movdqu xmmword ptr [r14+r10*1-0x10], xmm1 ; result[i:i+16] = xmm1
sub rcx, 0x1
jnz <Block 15> ; Loop again while i!=a.shape[0]//16
此代码使用名为SSE的SIMD指令集,该指令集能够在128位宽的寄存器上工作。没有条件跳转。此代码的效率要高得多,因为它每次迭代一次处理16个项目,并且每次迭代都应该快得多。
请注意,最后一段代码也不是最佳的,因为大多数现代x86处理器(如AMD处理器(都支持256位AVX-2指令集(速度是原来的两倍(。此外,编译器生成了一个低效的SIMD指令序列来执行逻辑,并且可以进行优化。编译器似乎认为布尔值可以是不同于0或1的值。也就是说,输入数组太大,无法放入CPU缓存,因此与第一个数组相比,代码受到RAM吞吐量的限制。这就是为什么SIMD友好代码的速度并没有明显加快的原因。两个版本之间的差异肯定要大得多,因为你的处理器上的阵列小于1MIB(就像几乎所有其他现代处理器上一样(。