NumPy与Numba在不同复杂度数组运算中的性能差异原因及优化咨询

先贴出测试用的完整代码：

import numpy as np
import timeit as ti
import numba as nb

def f0(a):
    p0 = a[:-2]
    p1 = a[1:-1]
    p2 = a[2:]
    return (p0 < p1) & (p1 > p2)

def f1(a):
    p0 = a[:-4]
    p1 = a[1:-3]
    p2 = a[2:-2]
    p3 = a[3:-1]
    p4 = a[4:]
    return ((p0 < p1) & (p1 == p2) | (p1 < p2)) & ((p2 > p3) | (p2 == p3) & (p3 > p4))

@nb.njit(fastmath=True)
def g0(a):
    r = np.zeros_like(a, dtype=np.bool)
    for i in range(1, a.size-1):
        r[i] = (a[i-1] < a[i]) & (a[i+1] < a[i])
    return r[1:-1]

@nb.njit(fastmath=True)
def g1(a):
    r = np.zeros_like(a, dtype=np.bool)
    for i in range(2, a.size-2):
        r[i] = ((a[i-1] == a[i]) & (a[i-2] < a[i-1]) | (a[i-1] < a[i])) & \
               ((a[i+1] == a[i]) & (a[i+2] < a[i+1]) | (a[i+1] < a[i]))
    return r[2:-2]

a = np.random.randint(0, 256, (500, 500)).astype(np.uint8)
b = a.ravel()

print(f'Minimum, median and maximum execution time in us:')
for fun in ('f0(b)', 'f1(b)', 'g0(b)', 'g1(b)'):
    t = 10**6 * np.array(ti.repeat(stmt=fun, setup=fun, globals=globals(), number=1, repeat=999))
    print(f'{fun:20} {np.amin(t):8,.3f} {np.median(t):8,.3f} {np.amax(t):8,.3f}')

运行在AMD Ryzen 7 3800X、Ubuntu 22.04.5系统上的测试结果：

Minimum, median and maximum execution time in us:
f0(b)                32.261    33.483    95.640
f1(b)               118.974   120.737   129.424
g0(b)                11.081    11.281    19.327
g1(b)               723.319   744.419   794.042

为什么会出现这种性能反转？ #

这个现象其实挺典型的，咱们一步步拆解原因：

• NumPy在复杂向量化运算中的优势
  NumPy的核心是向量化操作，所有数组切片、逻辑运算都是在底层C语言实现的，还会自动利用CPU的SIMD（单指令多数据）指令批量处理数据。像f1这种复杂但纯向量化的表达式，NumPy会把整个运算链优化成连续的内存操作，没有Python层面的循环开销，而且SIMD能一次处理多个uint8元素（比如一次16个），效率拉满。

• Numba复杂循环的性能瓶颈
  再看g1，虽然用Numba编译成了机器码，但几个点拖慢了速度：

  • 内存与SIMD利用不足：循环里每次要访问a[i-2]到a[i+2]的分散元素，加上逻辑判断的分支（|和&的组合），编译器很难自动做SIMD向量化优化，没法批量处理数据。
  • 冗余内存操作：g1先创建了和输入同大小的数组r，最后再切片返回，多了不必要的内存分配、初始化和无效元素赋值操作。
  • 分支预测失效：uint8取值范围大，循环里的==判断结果随机性强，CPU的分支预测器很容易出错，导致流水线停顿，大幅降低执行速度。

可以怎么优化性能？ #

针对g1的问题，试试这些优化方向：

1. 优化Numba版本的内存与并行性
   直接创建目标大小的结果数组，减少冗余操作；用prange开启并行，利用多核心优势：

   @nb.njit(fastmath=True, parallel=True)
   def g1_opt(a):
       res_size = a.size - 4
       r = np.zeros(res_size, dtype=np.bool)
       for i in nb.prange(res_size):
           pos = i + 2
           left = (a[pos-1] == a[pos]) & (a[pos-2] < a[pos-1]) | (a[pos-1] < a[pos])
           right = (a[pos+1] == a[pos]) & (a[pos+2] < a[pos+1]) | (a[pos+1] < a[pos])
           r[i] = left & right
       return r
   

2. 微调NumPy版本的f1
   用sliding_window_view减少重复切片的开销，虽然优化幅度不大，但代码更简洁：

   def f1_opt(a):
       window = np.lib.stride_tricks.sliding_window_view(a, 5)
       p0, p1, p2, p3, p4 = window.T
       return ((p0 < p1) & (p1 == p2) | (p1 < p2)) & ((p2 > p3) | (p2 == p3) & (p3 > p4))
   

3. 调整Numba编译选项
   可以尝试去掉fastmath=True（复杂逻辑下快速数学优化可能适得其反），或者加上cache=True缓存编译后的代码，避免重复编译开销。

内容来源于stack exchange