两种二维数组求和方式性能差异悬殊的原因是什么？

先看这两个用Numba加速的NumPy数组求和实现：

import numpy as np
from numba import njit
a = np.random.rand(2, 5000)

@njit(fastmath=True, cache=True)
def sum_array_slow(arr):
    s = 0
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            s += arr[i, j]
    return s
    
@njit(fastmath=True, cache=True)
def sum_array_fast(arr):
    s = 0
    for i in range(arr.shape[1]):
        s += arr[0, i]
    for i in range(arr.shape[1]):
        s += arr[1, i]
    return s

从逻辑上看，两者的元素访问顺序完全一致：先遍历第一行所有元素，再遍历第二行所有元素。但实际测试结果却差了8倍：

In [46]: %timeit sum_array_slow(a)
7.7 µs ± 374 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [47]: %timeit sum_array_fast(a)
951 ns ± 2.63 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

核心原因：编译器优化能力的差异 #

两者的性能差距本质是Numba对不同循环结构的向量化优化程度不同：

• sum_array_fast的单循环结构：每个循环都是对连续内存块的线性遍历，逻辑简单直接。Numba的LLVM后端能轻松识别这种模式，直接生成SIMD（单指令多数据）指令——一次就能处理多个数组元素，充分利用CPU的并行计算能力。同时，编译器还能做循环展开、寄存器复用等深度优化，进一步降低开销。

• sum_array_slow的嵌套循环结构：虽然逻辑等价，但外层循环的存在会干扰编译器的优化分析。LLVM很难将两层循环合并为可向量化的单一数据流，即使内层循环是连续访问，外层循环的切换也会让优化器无法生成高效的SIMD代码。最终生成的代码只能逐个元素累加，完全没用到CPU的并行计算能力，自然速度慢很多。

额外提一句：嵌套循环的外层条件判断会带来微小的开销，但这不是主导因素，核心还是SIMD向量化的缺失。

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