使用 `@jit` 编译 Python 代码

Numba 提供了多种代码生成工具，但其核心功能是 numba.jit() 装饰器。使用此装饰器，您可以将函数标记为由 Numba 的 JIT 编译器进行优化。不同的调用模式会触发不同的编译选项和行为。

基本用法

惰性编译

推荐使用 @jit 装饰器的方式是让 Numba 决定何时以及如何进行优化

from numba import jit

@jit
def f(x, y):
    # A somewhat trivial example
    return x + y

在此模式下，编译将推迟到第一次函数执行时进行。Numba 将在调用时推断参数类型，并根据此信息生成优化代码。Numba 还能够根据输入类型编译不同的专用版本。例如，使用整数或复数调用上面的 f() 函数将生成不同的代码路径

>>> f(1, 2)
3
>>> f(1j, 2)
(2+1j)

即时编译

您还可以告诉 Numba 预期的函数签名。函数 f() 现在看起来会像

from numba import jit, int32

@jit(int32(int32, int32))
def f(x, y):
    # A somewhat trivial example
    return x + y

int32(int32, int32) 是函数的签名。在这种情况下，相应的专用版本将由 @jit 装饰器编译，并且不允许其他专用版本。如果您需要对编译器选择的类型进行细粒度控制（例如，使用单精度浮点数），这会很有用。

如果您省略返回类型，例如通过写入 (int32, int32) 而不是 int32(int32, int32)，Numba 将尝试为您推断。函数签名也可以是字符串，并且您可以将多个签名作为列表传递；有关更多详细信息，请参阅 numba.jit() 文档。

当然，编译后的函数会给出预期的结果

>>> f(1,2)
3

如果我们将 int32 指定为返回类型，则高阶位将被丢弃

>>> f(2**31, 2**31 + 1)
1

调用和内联其他函数

Numba 编译的函数可以调用其他已编译的函数。根据优化器启发式算法，函数调用甚至可以内联到本机代码中。例如

@jit
def square(x):
    return x ** 2

@jit
def hypot(x, y):
    return math.sqrt(square(x) + square(y))

@jit 装饰器必须添加到任何此类库函数中，否则 Numba 可能会生成慢得多的代码。

签名规范

显式的 @jit 签名可以使用多种类型。以下是一些常见的类型

void 是不返回任何内容的函数的返回类型（从 Python 调用时实际上返回 None）
intp 和 uintp 是指针大小的整数（分别为有符号和无符号）
intc 和 uintc 等同于 C 语言的 int 和 unsigned int 整数类型
int8, uint8, int16, uint16, int32, uint32, int64, uint64 是对应位宽的固定宽度整数（有符号和无符号）
float32 和 float64 分别是单精度和双精度浮点数
complex64 和 complex128 分别是单精度和双精度复数
数组类型可以通过索引任何数值类型来指定，例如 float32[:] 表示一维单精度数组，或者 int8[:,:] 表示二维 8 位整数数组。

编译选项

可以将许多仅限关键字的参数传递给 @jit 装饰器。

`nopython`

Numba 有两种编译模式：nopython 模式和对象模式。Nopython 模式是默认模式，它生成更快代码，但有局限性。

@jit  # same as @jit(nopython=True) or @njit since Numba 0.59
def f(x, y):
    return x + y

另请参阅

故障排除和提示

`nogil`

每当 Numba 将 Python 代码优化为仅在本机类型和变量（而非 Python 对象）上工作的本机代码时，就不再需要持有 Python 的全局解释器锁 (GIL)。如果您传递了 nogil=True，Numba 将在进入此类编译函数时释放 GIL。

@jit(nogil=True)
def f(x, y):
    return x + y

释放 GIL 后运行的代码可以与执行 Python 或 Numba 代码的其他线程（无论是相同的编译函数还是其他函数）并发运行，从而允许您利用多核系统。如果函数以对象模式编译，这将无法实现。

使用 nogil=True 时，您必须警惕多线程编程的常见陷阱（一致性、同步、竞态条件等）。

`cache`

为了避免每次调用 Python 程序时的编译时间，您可以指示 Numba 将函数编译结果写入基于文件的缓存中。这通过传递 cache=True 来完成。

@jit(cache=True)
def f(x, y):
    return x + y

注意

编译函数的缓存有几个已知限制

编译函数的缓存不是基于函数进行的。缓存的函数是主要的 jit 函数，所有次要函数（由主函数调用的函数）都包含在主函数的缓存中。
缓存失效无法识别在不同文件中定义的函数的变化。这意味着当一个主要的 jit 函数调用从不同模块导入的函数时，这些其他模块的更改将不会被检测到，并且缓存将不会更新。这带来了在计算中使用“旧”函数代码的风险。
全局变量被视为常量。缓存将记住编译时全局变量的值。在缓存加载时，缓存的函数将不会重新绑定到全局变量的新值。

`parallel`

为函数中已知具有并行语义的操作启用自动并行化（和相关优化）。有关支持的操作列表，请参阅使用 @jit 进行自动并行化。此功能通过传递 parallel=True 启用，并且必须与 nopython=True 结合使用。

@jit(nopython=True, parallel=True)
def f(x, y):
    return x + y

另请参阅

使用 @jit 进行自动并行化

使用 @jit 编译 Python 代码