使用 `@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

参见

numba-故障排除

`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 代码