2. 第一个@jit函数

上一节我们理解了Numba“在运行时把Python函数翻译成机器码”的基本定位。这一节要真正动手：写一个纯Python数值函数，给它加上@jit装饰器，观察它从“解释执行”变成“编译执行”的过程，并学会用nopython模式判断编译是否成功，用计时手段区分编译时间与运行时间。这是后续所有Numba优化的起点。

从纯Python函数开始

选一个足够简单、又能体现数值计算特征的函数作为起点：计算一维数组的平方和。纯Python版本如下：

import numpy as np

def sum_sq_py(a):
    total = 0.0
    for i in range(a.shape[0]):
        total += a[i] * a[i]
    return total

a = np.arange(1_000_000, dtype=np.float64)
print(sum_sq_py(a))   # 3.333328333335e+17

这个函数用显式循环逐元素累加，正是Numba最擅长加速的形态：循环次数多、每轮只做少量标量运算、没有复杂对象操作。如果改用np.sum(a*a)，向量化本身已经很快，Numba的收益反而不明显。因此教学示例刻意保留循环写法。

加上@jit装饰器

Numba的核心入口是numba.jit。最简用法是把它作为装饰器套在函数上：

from numba import jit

@jit
def sum_sq_numba(a):
    total = 0.0
    for i in range(a.shape[0]):
        total += a[i] * a[i]
    return total

此时函数体一个字都没改，只是多了一行装饰器。第一次调用sum_sq_numba(a)时，Numba会检查传入参数的类型（这里是float64[:]一维数组），据此生成一份针对该类型的机器码并缓存，然后执行。第二次及以后用相同类型调用，就直接复用已编译版本，不再编译。

装饰器还可以带参数，最常用的是nopython=True和cache=True：

@jit(nopython=True, cache=True)
def sum_sq_numba(a):
    total = 0.0
    for i in range(a.shape[0]):
        total += a[i] * a[i]
    return total

cache=True会把编译产物写入__pycache__，下次启动进程时若函数与类型未变，可跳过编译直接加载，这对反复调试的脚本很实用。

nopython模式与object模式

这是本节最需要建立的概念区分。@jit默认采用“惰性回退”策略：Numba先尝试把整个函数编译成不依赖Python解释器的原生代码，这就是nopython模式；如果函数里出现了Numba不支持的Python特性（例如操作字典的某些方法、调用未支持的第三方库、使用try/except的复杂形式等），编译会失败，Numba就退回到object模式——把函数中能编译的部分编译，不能编译的对象仍以Python对象方式处理，运行时频繁进出解释器。

object模式虽然“能跑”，但性能往往接近甚至不如纯Python，因为它没有真正摆脱Python对象模型。因此工程上应尽量让函数停留在nopython模式。判断方法有两种：一是显式写@jit(nopython=True)，此时若无法编译会直接抛出TypingError或NumbaError，问题立刻暴露；二是用.nopython_signatures或.signatures属性检查已编译签名，或调用numba.typeof辅助排查。推荐在开发阶段一律使用nopython=True，把回退变成显式报错，避免“看起来加速了、其实还在解释执行”的假象。

测量编译时间与运行时间

Numba的加速效果必须用数据说话，而计时最容易踩的坑，是把首次调用的编译开销算进了“运行时间”。正确做法是：先做一次“预热调用”触发编译，再对后续调用计时。下面用time.perf_counter演示：

import time

# 预热：触发编译，不计入运行时间
sum_sq_numba(a)

# 测量纯运行时间（多次取平均）
t0 = time.perf_counter()
for _ in range(20):
    sum_sq_numba(a)
t_run = (time.perf_counter() - t0) / 20

# 单独测量编译时间：换一个新类型触发重新编译
b = np.arange(1_000_000, dtype=np.float32)
t0 = time.perf_counter()
sum_sq_numba(b)          # 首次以 float32 调用，触发新编译
t_compile = time.perf_counter() - t0

print(f"运行时间: {t_run*1e3:.3f} ms")
print(f"编译时间: {t_compile*1e3:.1f} ms")

在一台普通笔记本上，这类示例的典型量级是：纯Python循环约几十毫秒，nopython模式约零点几毫秒，编译时间约几十到上百毫秒。也就是说，编译开销通常远大于单次运行时间，只有当函数被反复调用（例如在循环、仿真步进或大批量数据处理中调用成千上万次）时，摊销后的收益才显著。这正是Numba适合“热函数”而不适合“一次性脚本”的原因。

一个可复现的对比示例

把上面的片段组织成一个完整脚本，就能得到一张清晰的对比表（以下数值为示例量级，实际结果随机器与数组规模变化）：

读表时要抓住两点：第一，比较“后续单次运行”才能反映真实加速比；第二，object模式那一行说明“加了装饰器不等于加速”，必须确认落在nopython模式。若你的函数在nopython=True下报错，应优先检查是否用了Numba不支持的容器操作或第三方调用，而不是简单退回默认@jit了事。

常见问题与排查方向

• 首次调用特别慢：这是编译开销，属正常现象；用预热调用把它排除在计时之外。
• 加了@jit却没变快：很可能回退到了object模式。改用nopython=True让错误显式抛出，再逐项修复不支持的写法。
• 换一种数组类型后又变慢：Numba按类型分别编译，float32与float64是两份机器码，各自都有首次编译成本。
• 函数里用了全局变量或闭包：nopython模式对全局常量的支持有限，尽量把依赖作为参数传入。