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
  • intpuintp是指针大小的整数(分别是有符号和无符号)
  • intcuintc相当于 C intunsigned int整数类型
  • int8uint8int16uint16int32uint32int64uint64是相应位宽的有限宽度整数(有符号和无符号)
  • float32float64分别是单精度和双精度浮点数
  • complex64complex128分别是单精度和双精度复数
  • 数组类型可以通过索引任何数字类型来指定,例如一维单精度数组的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 自动并行化


我们一直在努力

apachecn/AiLearning

【布客】中文翻译组