📚 课程目录第4章标量数值函数加速实战

4. 标量数值函数加速实战

前几节我们已经理解了 Numba 的 @njit 装饰器如何把 Python 函数编译为机器码。本节把视角落到最典型的“标量循环”场景：向量点积、欧几里得距离、多项式求值。这三类函数在数值计算、机器学习、图形学中反复出现，它们的共同特征是——用 Python 的 for 循环逐元素处理标量，逻辑清晰但解释器开销巨大。我们要做的，是把同一份循环逻辑分别交给纯 Python、NumPy 向量化和 Numba 编译三条路径，观察性能差异从何而来，并掌握写出“对 Numba 友好”的循环的要点。

一、为什么标量循环是 Numba 的主战场

Python 的 for 循环每次迭代都要做类型检查、对象装箱拆箱、字节码分派，处理一个 float64 加法的实际开销远大于加法本身。NumPy 的向量化通过把循环下沉到 C 层规避了这个问题，但它要求运算能表达为数组整体操作，且会为中间结果分配临时数组。Numba 则走第三条路：保留你写的显式循环，但在编译期确定元素类型，生成接近 C 的紧凑机器码，循环体内不再有解释器介入。因此，凡是能用标量循环自然表达、又会被高频调用的数值函数，都是 Numba 的理想目标。

下面这张图对比了三条执行路径在同一份点积逻辑上的数据流差异，注意观察“循环发生在哪一层”以及“是否产生临时数组”：

二、三个函数的 Numba 实现

先看向量点积。核心是显式循环累加，注意用 a.shape[0] 而非 len(a) 获取长度，让编译器更早确定循环边界类型：

import numpy as np
from numba import njit

@njit(cache=True)
def dot_numba(a, b):
    s = 0.0
    for i in range(a.shape[0]):
        s += a[i] * b[i]
    return s

欧几里得距离在点积基础上多一步开方。这里要提醒一个常见坑：math.sqrt 在 njit 中可用，但更稳妥的是用 np.sqrt，Numba 对 NumPy 标量函数支持良好：

@njit(cache=True)
def euclid_numba(a, b):
    s = 0.0
    for i in range(a.shape[0]):
        d = a[i] - b[i]
        s += d * d
    return np.sqrt(s)

多项式求值用 Horner 法（秦九韶算法），从最高次系数开始逐层“乘 x 加下一系数”，只需 n 次乘法和 n 次加法，比逐项求幂更省算力，也更适合循环表达。约定系数数组 coef 按降幂排列，即 coef[0] 是最高次项系数：

@njit(cache=True)
def polyval_numba(coef, x):
    r = 0.0
    for i in range(coef.shape[0]):
        r = r * x + coef[i]
    return r

三个函数都遵循同一原则：循环体内只做标量算术，不调用 Python 对象方法，不动态改变变量类型。累加器 s 和 r 初始化为 0.0（浮点），避免整数与浮点混用导致类型不稳定。

三、对比实验：如何公平地测

性能对比最容易出错的地方是计时方式。要点有三：第一，Numba 首次调用包含编译时间，必须先预热一次再计时；第二，用 time.perf_counter() 而非 time.time()；第三，多次重复取中位数，减少系统抖动。下面是一个可复用的计时骨架：

import time

def bench(fn, *args, repeat=7):
    fn(*args)                      # 预热，触发编译
    ts = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn(*args)
        ts.append(time.perf_counter() - t0)
    ts.sort()
    return ts[len(ts) // 2]        # 中位数

n = 1_000_000
rng = np.random.default_rng(0)
a = rng.random(n)
b = rng.random(n)
coef = rng.random(50)

t_py   = bench(lambda: sum(a[i]*b[i] for i in range(n)))
t_np   = bench(lambda: np.dot(a, b))
t_nb   = bench(dot_numba, a, b)
print(f"Python {t_py*1e3:.2f} ms | NumPy {t_np*1e3:.2f} ms | Numba {t_nb*1e3:.2f} ms")

在一台普通笔记本上，n=100 万时典型量级大致是：纯 Python 生成器求和约 60～100 ms，np.dot 约 0.5～1 ms（走 BLAS），Numba 循环约 0.8～1.5 ms。也就是说，Numba 相对纯 Python 通常有一到两个数量级的提升，与 NumPy 处于同一量级。需要强调，这些数字随 CPU、内存带宽、数组是否连续而变化，属于合理预期而非固定结论，实际应以本机测量为准。

四、结果解读：什么时候 Numba 赢，什么时候 NumPy 赢

对比不能只看一个大数组。把 n 从 10 逐步增大到 10⁶，会看到明显的交叉现象：

小数组（n 约几十到几百）：NumPy 的每次调用都有 Python 层与 C 层的切换开销，甚至不如 Numba 编译后的紧凑循环；而 Numba 若未预热，首次调用会被编译时间拖慢，预热后则非常轻快。
中等数组：两者接近，NumPy 的 SIMD 与 Numba 的标量循环各有优势，取决于运算是否容易被向量化。
大数组：np.dot 调用 BLAS 多线程，往往反超单线程 Numba 循环；但若你的函数是 NumPy 难以整体表达的自定义逻辑（如带条件分支的逐元素处理），Numba 的优势会重新显现。

多项式求值是个好例子：NumPy 有 np.polyval，但它内部同样走循环，且对单个标量 x 求值时调用开销明显；Numba 版本在“对同一组系数反复求值不同 x”的场景（如绘制函数曲线、批量采样）中优势突出。此时更优的写法是把 x 也做成数组，在 Numba 里写双层循环，一次编译调用处理整批数据，避免 Python 层循环调用。

判断是否该用 Numba 的实用准则：函数被高频调用、循环体是纯标量算术、NumPy 难以用整体数组操作表达、且数组规模不是特别巨大到必须依赖 BLAS 多线程。满足前三条时，Numba 通常值得一试。

五、常见陷阱与写法要点

第一，不要在 njit 函数里用 len() 之外的动态容器操作，如向 Python 列表 append，这会迫使 Numba 退回对象模式，性能骤降。第二，保持变量类型稳定：若循环中某变量先被赋值为 int 后被赋值为 float，Numba 需要做类型统一，可能引入额外转换。第三，开启 cache=True，把编译产物缓存到磁盘，避免每次进程启动都重新编译。第四，若函数只接受特定类型，可用 @njit("float64(float64[:], float64[:])") 显式签名，省去类型推断并让错误更早暴露。

最后，验证正确性别省略。对三个函数，用 np.allclose 与 NumPy 参考实现比对，并特意测试空数组、单元素、含负数和极大值等边界情况。性能优化只有在结果正确的前提下才有意义——先让 dot_numba(a,b) 与 np.dot(a,b) 在随机数据上一致到 1e-12，再去谈毫秒级的差距。