Numba向量化函数如何更高效利用并行能力提升运算速度

AI 小北

2026-8-29

问题背景

数年前我使用numba编写了一套数学库，用于对大规模元素数组执行向量化操作（例如对两个存储矩阵的数组做乘法运算）。近期我重新梳理该项目，尝试为其加入并行计算支持（我最初开发该库时numba尚未提供并行特性）。实际测试后虽然获得了一定速度提升，但效果远未达到我的预期。

为进一步探究原因，我编写了如下基础测试示例：

import numpy as np
from numba import njit, prange, float64

@njit(float64[:] (float64[:], float64[:]), parallel=True)
def array_sum(array0, array1):
    """
    Sum elements of array0 to array1, similar to doing + on two numpy arrays
    """

    result = np.empty(array0.shape)

    for i in prange(array0.shape[0]):
        result[i] = array0[i] + array1[i]

    return result

随后我使用如下代码开展性能测试：

import time
v0 = np.random.random(10**7) # 10 million floats
v1 = np.random.random(10**7) # 10 million floats
array_sum(v0,v1) # run this once so the code compiles

#--- test start ---#
start_time = time.time()
array_sum(v0,v1) # do the work
end_time = time.time()
delta = end_time-start_time
print (delta)

在我的个人设备上测试得到如下结果：

# direct numpy array0+array1 ---> 0.016865015029907227
# parallel=False             ---> 0.016468048095703125
# parallel=True              ---> 0.010709047317504883

测试结果显示，设置parallel=False时，函数运算速度与numpy原生数组相加速度一致，符合预期；设置parallel=True后仅获得约1.54倍的速度提升。我使用的测试设备配备8核CPU，对于数组逐元素相加这类计算密度极低的简单操作而言，该加速比远低于理论预期。

咨询问题

针对上述类似的数组向量化运算场景，是否存在更高效的并行使用方式，能够充分利用多核CPU的计算能力，获得更高的运算加速比？

回答

首先明确核心前提：逐元素数组加法属于典型的内存带宽绑定任务，而非计算绑定任务，8核CPU拿不到8倍线性加速是正常现象，和numba并行实现本身无关。这类任务每个元素仅需1次浮点运算，对应2次内存读、1次内存写操作，单线程执行时就已经接近吃满消费级平台的内存通道带宽，多线程带来的线程调度、缓存一致性同步开销会抵消大部分理论收益，1.5x左右的加速比是普通DDR4/DDR5平台上的常见结果。

如果要在这类场景下尽可能榨干硬件性能，可以按以下方式调整实现：

手动调整并行粒度，减少线程调度开销
默认prange逐元素调度的粒度过细，线程切换、伪共享的成本极高。可以手动做分块处理，让每个线程处理连续的大块内存，块大小设置为L2缓存可容纳的量级（通常1024~8192个元素之间调试即可），同时将numba线程数设置为物理核心数，不要开启超线程（超线程对内存绑定任务无正向收益，反而会增加内存争抢）。参考实现如下：
```
import numba
# 替换为设备的物理核心数
numba.set_num_threads(8)

@njit(float64[:](float64[:], float64[:]), parallel=True, fastmath=True)
def array_sum_blocked(array0, array1):
    n = array0.shape[0]
    result = np.empty(n, dtype=float64)
    block_size = 4096
    for i in prange(0, n, block_size):
        end = min(i + block_size, n)
        # 块内用普通循环，触发numba自动SIMD向量化
        for j in range(i, end):
            result[j] = array0[j] + array1[j]
    return result
```
装饰器中添加fastmath=True可以允许numba做更激进的指令优化，进一步提升单线程向量化效率。
保证输入数组内存连续
不要将非连续的数组切片、转置后未做连续转换的数组传入并行函数，非连续内存会导致CPU硬件预取失效，内存带宽利用率直接下降40%以上。不确定数组连续性时，可以提前用np.ascontiguousarray做预处理。
优先做算子融合，提升计算密度
不要对单个低计算密度的操作单独开启并行，这是numba并行拿到高加速比的核心原则。如果后续还要对相加结果做乘法、归一化、阈值截断、统计等操作，把所有逻辑合并到同一个prange循环中，让数据从内存加载后一次性完成所有计算再写回，避免数据多次在内存和CPU缓存之间往返。当单个循环内每个元素对应10次以上浮点运算时，并行加速比会快速接近物理核心数。
纯内存操作优先用numpy原生实现
如果确实需要单独执行逐元素拷贝、相加这类极低计算密度的操作，直接使用numpy原生接口即可。numpy后端对接的OpenBLAS/MKL已经对内存拷贝、逐元素运算做了极致的带宽优化，numba的并行调度层在这类场景下很难超过numpy原生性能。

可以用计算绑定任务（比如1024×1024规模的双精度矩阵乘法）做对照测试，这类场景下numba开启parallel=True基本可以拿到接近物理核心数的加速比，即可确认numba并行环境配置没有问题，之前的低加速比完全是任务属性导致的。

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