Neumaier求和可以加快吗？

Question

Neumaier 求和是 Kahan 求和的改进，用于对浮点数数组进行精确求和。

import numba as nb

@nb.njit
def neumaier_sum(arr):
    s = arr[0]
    c = 0.0
    for i in range(1, len(arr)):
        t = s + arr[i]
        if abs(s) >= abs(arr[i]):
            c += (s - t) + arr[i]
        else:
            c += (arr[i] - t) + s
        s = t
    return s + c

这效果很好，但比添加 fastmath=True 至少慢四倍。不幸的是，fastmath 允许重新括号总和（关联性），这会破坏其准确性，所以我们不能这样做。

这是一个测试，显示不同求和方式的结果。首先我们创建一个长度为 1001 的数组。

import numpy as np
n = 10 ** 3 + 1
a = np.full(n, 0.01, dtype=np.float64)
a[0] = 10**10
a[-1] = -10**10

我们

import math

，这样我们就可以使用

fsum

来查看正确答案。我们还使用 fastmath 定义了一个 Neumaier 版本。

# Bad version using fastmath
@nb.njit(fastmath=True)
def neumaier_sum_fm(arr):
    s = arr[0]
    c = 0.0
    for i in range(1, len(arr)):
        t = s + arr[i]
        if abs(s) >= abs(arr[i]):
            c += (s - t) + arr[i]
        else:
            c += (arr[i] - t) + s
        s = t
    return s + c

结果是：

math.fsum       : 9.99
nb_neumaier_sum : 9.99
nb_neumaier_sum_fm: 9.99001693725586

这些是时间安排：

 %timeit nb_neumaier_sum_fm(a)
350 ns ± 0.983 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
 %timeit nb_neumaier_sum(a)
1.5 µs ± 18.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

有什么方法可以加快这段代码的速度，但让它输出的结果与现在完全一样？

---

组装如下。这些评论是由chatgpt添加的，所以如果有任何错误请告诉我。

    movq    8(%rsp), %rcx          # Load the address of the array into rcx
    movq    16(%rsp), %rax         # Load the length of the array into rax
    vmovsd  (%rcx), %xmm1          # Load the first element of the array into xmm1 (s = arr[0])
    leaq    -1(%rax), %rdx         # Calculate (length of array - 1) into rdx
    testq   %rdx, %rdx             # Test if rdx is zero
    jle .LBB0_1               # If rdx <= 0, jump to .LBB0_1
    movabsq $.LCPI0_0, %rsi       # Load address of constant value (mask for abs)
    movl    $1, %edx               # Initialize loop counter (i = 1)
    vxorpd  %xmm0, %xmm0, %xmm0    # Clear xmm0 (c = 0.0)
    vmovapd %xmm1, %xmm3           # Copy xmm1 (s) to xmm3
    vmovapd (%rsi), %xmm2          # Load abs mask into xmm2

    .p2align    4, 0x90
.LBB0_3:
    vmovsd  (%rcx,%rdx,8), %xmm4   # Load arr[i] into xmm4
    vandpd  %xmm2, %xmm3, %xmm5    # Compute abs(s) into xmm5
    incq    %rdx                   # Increment rdx (i++)
    vaddsd  %xmm4, %xmm3, %xmm1    # Compute t = s + arr[i]
    vandpd  %xmm2, %xmm4, %xmm6    # Compute abs(arr[i]) into xmm6
    vcmpnlesd   %xmm5, %xmm6, %xmm5  # Compare abs(s) and abs(arr[i])
    vsubsd  %xmm1, %xmm3, %xmm6    # Compute (s - t) into xmm6
    vaddsd  %xmm6, %xmm4, %xmm6    # Compute (s - t) + arr[i] into xmm6
    vsubsd  %xmm1, %xmm4, %xmm4    # Compute (arr[i] - t) into xmm4
    vaddsd  %xmm4, %xmm3, %xmm3    # Compute (arr[i] - t) + s into xmm3
    vblendvpd   %xmm5, %xmm3, %xmm6, %xmm3  # Conditional blend based on the comparison
    vaddsd  %xmm3, %xmm0, %xmm0    # Update c
    vmovapd %xmm1, %xmm3           # Update s (s = t)
    cmpq    %rdx, %rax             # Compare loop counter with array length
    jne .LBB0_3                # If not equal, loop again

    vaddsd  %xmm1, %xmm0, %xmm0    # Final addition (s + c)
    xorl    %eax, %eax             # Clear eax
    vmovsd  %xmm0, (%rdi)          # Store the result
    retq                        # Return from function

Answer 1

分析

Numba 代码的速度受 3 个因素限制：

• （1000）个依赖操作的非常长的依赖链*（90%的时间在我的机器上）
• 从 CPython 调用 Numba 函数、检查参数类型并将其转换为本机 Numba 类型的开销（在我的机器上为 10%）

虽然我们可以尝试展开循环并打破依赖关系，但在大多数主流 x86 CPU 上，好处会很小（如果有的话），因为结果循环中要执行的指令数量较多（这是瓶颈的另一个来源）。

关键点是使用压缩SIMD指令而不是标量指令。从技术上讲，它意味着诸如

vsubsd

/

vaddsd

/等指令。 （即标量）应替换为

vsubpd

/

vaddpd

/等。 （即“包装”的）。然而，显然没有简单且可移植的方法来让 Numba 生成它。

Cython 中也是如此（可以使用依赖于编译器的 GCC 向量扩展 结合一些 Cython 技巧，但无论如何都缺少一些像混合这样的操作......

由于浮点数学不具有关联性并且

fastmath

被禁用，编译器无法更改指令的顺序并破坏依赖链。唯一的解决方案似乎是手动。

---

更快的解决方案

一个简单的解决方案是编写 C/C++ 代码。一种方法是使用低级 x86 SIMD 内在函数。这种方式会产生非常好的性能。然而，它并不是在所有 CPU 上都可移植的（例如，仅在支持 AVX-2 的 x86 CPU 上），高度容易出现 bug，并且非常难以维护（我尝试过，但我得到了一个快速代码，读起来很疯狂，而且肯定不准确）由于一个偷偷摸摸的错误）。另一种方法是使用 C++ 库来做到这一点，但支持混合（或 SIMD 预测）等操作、生成高效可移植代码的库实际上并不常见（XSIMD 似乎是一个不错的候选者）。最重要的是，生成的代码通常非常复杂，并且对于非熟练用户来说并不友好（代码充满了带有 2 个手动缩减的模板：一个用于剩余项目，另一个用于 SIMD 项目本身）。 更好的解决方案是使用 SIMD 友好的本机语言，如

ISPC（类似于着色器/CUDA/OpenCL，尽管 ISPC 与其他语言相比对 CPU 的支持非常好）。这是一个实现：

// sum.ispc

export uniform double neumaier(uniform double arr[], uniform size_t size)
{
    double s = 0.0d;
    double c = 0.0d;

    foreach (i = 0...size)
    {
        const double item = arr[i];
        const double t = s + item;

        if(abs(s) >= abs(item))
            c += (s - t) + item;
        else
            c += (item - t) + s;

        s = t;
    }

    uniform double uni_s = 0.0d;
    uniform double uni_c = 0.0d;

    // Also applies neumaier to perform the final reduction between SIMD lanes
    foreach_active(laneId)
    {
        const uniform double item = extract(s, laneId);
        const uniform double t = uni_s + item;

        if(abs(uni_s) >= abs(item))
            uni_c += (uni_s - t) + item;
        else
            uni_c += (item - t) + uni_s;

        uni_s = t;
    }

    return uni_s + (reduce_add(c) + uni_c);
}

请注意，

uniform

可以视为标量变量，

varying

可以视为 SIMD 变量。默认情况下，变量为

varying

。如需了解更多信息，请阅读 ISPC 文档。与其他本机实现（例如 C/C++）相比，此实现实际上相当简单，并且它是“可移植的”，因为 ISPC 支持几乎所有主流 x86/ARM CPU。它比 Numba 代码更复杂（主要是由于最终的 SIMD 减少，这只是为了提供非常准确的结果），但速度也明显更快。

要编译它，您需要 ISPC 编译器和链接器。以下是我在 Linux 上构建二进制文件所用的命令：

# Compilation (for CPUs supporting AVX-2)
# On windows, --dllexport is maybe required
ispc sum.ispc -o sum_ispc.o --target avx2-i64x4 --pic --addressing=64

# Link: build the library
gcc --shared sum_ispc.o -o sum_ispc.so

要从 CPython 调用此函数，可以使用 Numba、ctypes 或 Cython。单独使用，ctypes 的开销相当大。它可以与 Numba 混合使用，以减少开销（这要归功于 Numba 相当快的 JIT 生成的包装函数）。最快的解决方案显然是使用 CPython 的 C 扩展，但这写起来相当痛苦（尽管在我的机器上开销要小 2 倍）。

这是使用 ctypes + Numba 的示例：

import numba as nb
import ctypes

ispcLib = ctypes.CDLL('./sum_ispc.so')
ispcLib.neumaier.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
ispcLib.neumaier.restype = ctypes.c_double
neumaier_ctypes_ispc = ispcLib.neumaier

# Faster Numba wrapper (than calling neumaier_ctypes_ispc directly)
@nb.njit
def neumaier_numba_ispc(arr):
    return neumaier_ctypes_ispc(arr.ctypes.data, arr.size)

准确性和性能结果在

---

np.random.rand(1000)

和

np.random.normal(0,1,1000)

上，我使用 ISPC

 获得与 Numba 和 Cython 完全相同的输出。

以下是我的机器上的性能结果（Linux 上使用 i5-9600KF CPU）：

Numba:              1.41 µs ± 0.234 ns/loop
Cython:             1.21 µs ± 0.334 ns/loop
Pythran:            1.16 µs ± 0.128 ns/loop

ISPC + ctypes:      1.60 µs ± 15.9 ns/loop
ISPC + Numba:       0.48 µs ± 0.123 ns/loop     <-----
ISPC + C-ext:       0.38 µs ± 0.064 ns/loop     <-----

Estimated optimal:  0.30 µs

* mean ± std. dev. of 7 runs, 50,000 loops each

最快的 ISPC 解决方案比 Numba/Cython/Pythran 快约

3 倍

。