基于Numba优化含NaN的绝对有界求和运算的方法探讨

问题背景与现有代码 #

我基于Numba官方文档修改了如下示例代码，用于对可能包含np.nan的一维NumPy数组执行绝对有界求和运算：

from numba import njit
import numpy as np

@njit
def do_sum(A, lb, ub):
    n = len(A)
    acc = 0.0
    for i in range(n):
        a = 0.0 if np.isnan(A[i]) else A[i]
        acc += abs(max(min(a, ub[i]), lb[i]))
    return acc

其中A、lb、ub为等长一维数组，长度通常在数千至数万之间，且该函数需要被调用数百万次，因此急需性能优化。

核心疑问 #

1. 由于A可能包含np.nan，无法直接使用@njit(fastmath=True)，但无nan场景下该选项能带来显著速度提升，想知道是否有折中方案，既能利用fastmath加速，又能获得比当前实现更优的性能？或者任何性能优化方法都可以推荐。（可先假设lb和ub不含np.nan，若方案能同时处理它们含np.nan的情况则更佳）
2. 该代码与Numba官方文档中演示parallel=True的示例结构类似，但添加parallel=True后性能反而显著下降，对此存在疑问。

优化方案与解释 #

1. 分场景兼容fastmath的折中方案 #

可以通过分支判断+函数重载实现：提前判断数组中是否存在nan，分别调用带fastmath和不带fastmath的版本。因为np.isnan(A).any()的开销相对于数百万次函数调用来说可以忽略，且大部分场景下数组要么全是有效值，要么固定含nan，可以缓存判断结果避免重复计算。

示例代码：

from numba import njit, objmode
import numpy as np

# 无nan场景专用，开启fastmath
@njit(fastmath=True)
def do_sum_no_nan(A, lb, ub):
    n = len(A)
    acc = 0.0
    for i in range(n):
        a = A[i]
        acc += abs(max(min(a, ub[i]), lb[i]))
    return acc

# 兼容nan的版本，不开启fastmath
@njit
def do_sum_with_nan(A, lb, ub):
    n = len(A)
    acc = 0.0
    for i in range(n):
        a = 0.0 if np.isnan(A[i]) else A[i]
        acc += abs(max(min(a, ub[i]), lb[i]))
    return acc

# 对外统一入口
@njit
def do_sum_opt(A, lb, ub):
    with objmode(has_nan='boolean'):
        has_nan = np.isnan(A).any()
    if has_nan:
        return do_sum_with_nan(A, lb, ub)
    else:
        return do_sum_no_nan(A, lb, ub)

如果你的调用场景中，同一批数组的nan状态固定，可以提前在外层判断好，直接调用对应版本，性能会更优。

2. 单函数内的微优化 #

如果不想拆分函数，也可以对原函数做以下调整：

• 把单元素调用的np.isnan替换为更快的判断方式（A[i] != A[i]是Numba中检测nan的原生操作，无额外函数调用开销）
• 手动展开循环，减少循环分支的性能损耗

优化后的单函数版本：

from numba import njit
import numpy as np

@njit
def do_sum_fast(A, lb, ub):
    n = len(A)
    acc = 0.0
    # 处理偶数长度部分，手动展开循环
    for i in range(0, n - n%2, 2):
        # 处理第一个元素
        a1 = A[i]
        if a1 != a1:
            a1 = 0.0
        val1 = abs(max(min(a1, ub[i]), lb[i]))
        # 处理第二个元素
        a2 = A[i+1]
        if a2 != a2:
            a2 = 0.0
        val2 = abs(max(min(a2, ub[i+1]), lb[i+1]))
        acc += val1 + val2
    # 处理剩余的奇数元素
    if n % 2 != 0:
        a = A[-1]
        if a != a:
            a = 0.0
        acc += abs(max(min(a, ub[-1]), lb[-1]))
    return acc

3. 为什么parallel=True会变慢？ #

原因主要有两点：

• 并行开销大于收益：你的数组长度是数千至数万，而每个循环迭代的计算量极小（仅几个比较、取绝对值、累加操作）。并行化需要线程创建、任务分发、结果合并的开销，这些开销远超并行计算节省的时间，导致整体性能下降。
• 累加操作的同步开销：原函数中的acc是全局累加变量，并行时需要原子操作保证线程安全，这会带来额外的同步损耗，进一步拉低性能。即便修改成分块累加再合并的正确并行实现，对于小规模数组来说收益仍然有限。

内容的提问来源于stack exchange，提问作者Keptain