3. Numba类型系统与延迟编译

Numba 的加速能力来自“先确定类型、再生成机器码”这条主线。同一个 Python 函数，如果传入的参数类型不同，Numba 会为每种类型组合分别编译一份专用版本。理解类型系统，就等于理解了 Numba 什么时候编译、编译出什么、以及为什么有时快有时慢。

类型推断：从 Python 对象到 Numba 类型

当你用 @njit 装饰函数并直接调用时，Numba 进入延迟编译（lazy compilation）：它先观察本次调用实参的 Python 类型，推断出对应的 Numba 类型，再据此编译。例如传入 np.float64 标量和 float64[:] 一维数组，Numba 会推断出 (float64, array(float64, 1d, C)) 这样的签名。推断是逐层进行的：函数体内的局部变量、循环变量、中间结果都要能落到具体类型上，否则推断失败。

推断的难点在于 Python 的动态性。空列表、None、混合类型容器、返回类型随分支变化的函数，都会让 Numba 无法确定唯一类型。典型报错是 Cannot determine Numba type of <class 'NoneType'> 或 cannot infer the type of variable。处理思路是消除歧义：给空容器显式初始化（如 np.zeros(n) 而非 []），避免同一变量在不同分支被赋成不同类型，必要时用 numba.typed.List 替代普通列表。

显式签名：eager compilation 的写法

与延迟编译相对的是急切编译（eager compilation）：在装饰器里直接写明签名，函数定义时即完成编译。签名用字符串或类型对象表达，返回类型写在参数列表之后：

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

# 字符串签名：参数为 float64 一维数组与 float64 标量，返回 float64
@njit("float64(float64[:], float64)")
def scale_sum(arr, k):
    s = 0.0
    for i in range(arr.shape[0]):
        s += arr[i] * k
    return s

# 类型对象写法，等价
@njit(float64(float64[:], float64))
def scale_sum2(arr, k):
    return scale_sum(arr, k)

显式签名带来三个好处：编译错误在定义阶段就暴露，而不是等到第一次调用；可以预先为多种类型组合注册多个版本；便于在部署时固定行为。代价是灵活性下降——若实际传入类型与签名不符，Numba 会尝试类型转换，转换失败则抛 TypeError。实践中常见做法是：开发期用延迟编译快速迭代，稳定后用显式签名锁定接口。

类型映射与常见错误诊断

Python 与 NumPy 类型到 Numba 类型的映射大体直观：int 对应 int64，float 对应 float64，bool 对应 bool_，NumPy 数组按元素类型与维度映射为 array(dtype, ndim, layout)，其中 layout 为 C（行优先连续）、F（列优先）或 A（任意）。布局信息很关键：C 连续数组的遍历效率最高，切片产生的非连续视图可能触发额外检查。

诊断时最有效的工具是 inspect_types() 与 signatures。前者打印函数体内每个变量的推断类型，后者列出已编译的签名集合。若发现某变量被推断成 pyobject，说明它退回了 Python 对象层，通常就是性能瓶颈所在。

一次可复现的诊断示例

假设有一个求加权和的函数，初版写成：

from numba import njit
import numpy as np

@njit
def wsum(x, w):
    total = 0          # 整数 0，后续与 float 相加
    for i in range(x.shape[0]):
        total += x[i] * w[i]
    return total

x = np.array([1.0, 2.0, 3.0])
w = np.array([0.5, 0.5, 0.5])
print(wsum(x, w))

这段代码能跑，但 total 初值为 int64，与 float64 相乘相加时 Numba 会做类型提升，通常不报错但可能引入不必要的转换。更稳妥的写法是 total = 0.0。若把 w 换成 Python 列表 [0.5, 0.5, 0.5]，Numba 会推断为 reflected list(float64)，循环中每次索引都要走反射列表路径，性能明显下降——这正是“类型推断结果决定性能”的直接体现。

诊断时执行 wsum.inspect_types()，输出中若看到 total = const(int, 0) 后紧跟 total = total + x[i] * w[i] 且变量类型标注为 float64，说明提升发生在循环内；若看到 pyobject，则说明该变量已退回对象层，必须改写。另一个常用检查是 wsum.signatures，它返回类似 [(array(float64, 1d, C), array(float64, 1d, C))] 的列表，可确认实际编译了哪些版本。

选择策略小结

延迟编译适合探索阶段：改代码即改行为，无需维护签名列表，但要注意首次调用的编译延迟和签名膨胀。急切编译适合接口固定的库函数：编译期校验、行为可预期、便于分发给不关心内部类型的调用方。两者并非互斥——可以先以延迟编译开发，稳定后把实际出现的签名抄进装饰器，完成从“推断”到“声明”的收敛。真正决定性能的，始终是推断出的类型是否落在 Numba 原生支持的数值类型与连续数组布局上。