计算复杂numpy ndarray的abs()**2的最有效方法是什么?

b09cbbtk  于 2023-03-23  发布在  其他
关注(0)|答案(1)|浏览(127)

我正在寻找最省时的方法来计算python中复杂ndarray的绝对平方值。

arr = np.empty((8, 4000), dtype="complex128")  # typical size

我尝试了这些选项:
x一个一个一个一个x一个一个二个x
事实证明,numba实现比numpy快大约4倍,我想知道是否存在更快的方法。
我读过这篇question文章,其中提到了几个方法,但这篇文章是面向内存效率的,这在我的情况下不是一个约束。

m3eecexj

m3eecexj1#

所提供的函数的执行时间是非常具有挑战性的。在这种情况下,找到是否有更好的可能实现的最佳解决方案是 * 查看生成的汇编代码 *,因为如果生成的汇编代码接近最佳,那么盲目地尝试编写许多替代函数是浪费时间的。事实上,这里就是这样的情况:生成的汇编代码在输入布局方面非常好
Numba在内部生成了许多方法来计算数组,并在运行时根据输入选择最佳方法。当输入数组在内存中是连续的时,最佳实现使用以下汇编热循环:

.LBB0_55:
    vmovupd -192(%r9,%rbp,2), %ymm0
    vmovupd -160(%r9,%rbp,2), %ymm1
    vmovupd -128(%r9,%rbp,2), %ymm2
    vmovupd -96(%r9,%rbp,2), %ymm3
    vmovupd -64(%r9,%rbp,2), %ymm4
    vmovupd -32(%r9,%rbp,2), %ymm5
    vmovupd (%r9,%rbp,2), %ymm6
    vmulpd  %ymm1, %ymm1, %ymm1
    vmulpd  %ymm0, %ymm0, %ymm0
    vhaddpd %ymm1, %ymm0, %ymm0
    vmovupd 32(%r9,%rbp,2), %ymm1
    vmulpd  %ymm3, %ymm3, %ymm3
    vmulpd  %ymm2, %ymm2, %ymm2
    vhaddpd %ymm3, %ymm2, %ymm2
    vmulpd  %ymm5, %ymm5, %ymm3
    vmulpd  %ymm4, %ymm4, %ymm4
    vhaddpd %ymm3, %ymm4, %ymm3
    vmulpd  %ymm1, %ymm1, %ymm1
    vmulpd  %ymm6, %ymm6, %ymm4
    vhaddpd %ymm1, %ymm4, %ymm1
    vpermpd $216, %ymm0, %ymm0
    vpermpd $216, %ymm2, %ymm2
    vpermpd $216, %ymm3, %ymm3
    vpermpd $216, %ymm1, %ymm1
    vmovupd %ymm0, (%r11,%rbp)
    vmovupd %ymm2, 32(%r11,%rbp)
    vmovupd %ymm3, 64(%r11,%rbp)
    vmovupd %ymm1, 96(%r11,%rbp)
    subq    $-128, %rbp
    addq    $-16, %rdx
    jne .LBB0_55

这个循环非常高效。实际上,它是使用AVX-2 SIMD instruction set进行向量化的(在我的机器上这种情况下是最好的)。这可以通过查看操作码前缀/后缀和操作数来看出:v前缀用于AVX,pd后缀用于压缩双精度运算,ymm是256位AVX寄存器。循环也是展开的,因此循环迭代的开销可以忽略不计。我几乎不相信任何替代的Numba函数可以生成更好的热循环。事实上,我也不期望任何自动向量化的本机代码会明显更快。
虽然这个循环非常好,但我不认为它是最佳的。(如vshufpdvunpckhpdvunpcklpd以及可能的vpermpd)。也就是说,速度提升可能很小。实际上,输入需要64 KiB,输出需要32 KiB。这意味着操作肯定会在L2缓存中完成,L2缓存通常较慢,可能是更优化代码的瓶颈。事实上,在我的机器上,一个只返回x.real的函数需要同样的时间!更不用说用汇编语言编写这样的代码或用本机语言使用intrinsic了(如C/C++)导致代码不可移植,并且需要高度的技能(在C/C++,汇编,并有一个非常好的知识如何x86-64处理器操作)。简单地说,我认为这个循环是你能得到的最好的顺序Python代码
在目标平台上,使用多线程会更快。在这种粒度下,创建线程、共享工作和等待线程会花费大量时间,因此实际上可能比某些计算机上的计算速度要慢(尤其是有很多核心的计算服务器)。在我的机器上开销很低(几微秒)。因此,使用多线程稍微快一点。这远远不是很好,因为使用更多线程就像浪费了大部分内核,因为速度相当令人失望。最重要的是,Numba无法使用SIMD指令对代码进行自动向量化,这里是普通循环(vectorize似乎还不支持parallel=True标志)。我们可以通过预分配一次数组来加快函数的速度,因为临时数组的创建(和第一次填充)有点慢。下面是生成的代码:

@numba.njit('(complex128[:,::1],float64[:,::1])', parallel=True)
def abs2_numba(x, out):
    for i in numba.prange(x.shape[0]):
        tin = x[i, :]
        tout = out[i, :]
        for j in range(x.shape[1]):
            v = tin[j]
            real = v.real
            imag = v.imag
            tout[j] = real * real + imag * imag

out = np.empty(arr.shape)
%timeit -n 10_000 abs2_numba(x, out)

在我的6核i5-9600机器上,这需要9.8 µs,而初始代码需要12.9 µs。大部分时间都花在了并行开销上,而且代码没有像其他顺序代码那样矢量化。
我不建议你在你的情况下使用这种并行实现,除非每微秒都很重要,而且你不打算在一台有很多内核的机器上运行它。相反,优化调用这个函数的代码肯定是你最好的选择。

相关问题