Skip to content

1.4。使用@jit 编译 Python 代码

原文： http://numba.pydata.org/numba-doc/latest/user/jit.html

Numba 为代码生成提供了几个实用程序，但它的核心功能是 numba.jit() 装饰器。使用这个装饰器，您可以通过 Numba 的 JIT 编译器标记一个函数进行优化。各种调用模式会触发不同的编译选项和行为。

1.4.1。基本用法

1.4.1.1。懒惰编译

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

1.4.1.2。急切的编译

您还可以告诉 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

1.4.2。调用和内联其他功能

Numba 编译的函数可以调用其他编译函数。函数调用甚至可以在本机代码中内联，具体取决于优化器启发式。例如：

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

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

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

1.4.3。签名规格

显式@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[:]或 8 位整数的二维数组的int8[:,:]。

1.4.4。编译选项

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

1.4.4.1。 nopython

Numba 有两种编译模式： nopython 模式和对象模式。前者产生更快的代码，但有一些限制可以迫使 Numba 回到后者。为防止 Numba 退回，而是引发错误，请传递nopython=True。

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

也可以看看

故障排除和提示

1.4.4.2。 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时，您必须警惕多线程编程的常见缺陷（一致性，同步，竞争条件等）。

1.4.4.3。 cache

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

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

1.4.4.4。 parallel

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

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

也可以看看

使用@jit 自动并行化

回到顶部