即时编译

JIT 函数

@numba.jit(signature=None, nopython=False, nogil=False, cache=False, forceobj=False, parallel=False, error_model='python', fastmath=False, locals={}, boundscheck=False)

即时编译被装饰的函数以生成高效的机器代码。所有参数都是可选的。

如果存在,signature 是单个签名或签名列表,表示函数参数和返回值的预期类型和签名。每个签名可以有几种形式:

  • 一个类型和签名参数的元组(例如(numba.int32, numba.double)),表示函数参数的类型;Numba 将根据参数推断适当的返回类型。

  • 使用类型和签名的调用签名,同时指定返回类型和参数类型。这可以以直观的形式给出(例如numba.void(numba.int32, numba.double))。

  • 上述之一的字符串表示形式,例如"void(int32, double)"。字符串中使用的所有类型名称都假定在numba.types模块中定义。

nopythonnogil 是布尔标志。locals 是局部变量名到类型和签名的映射。

此装饰器有几种操作模式:

  • 如果在 signature 中给出了一个或多个签名,则会为每个签名编译一个专用版本。调用被装饰的函数时,将尝试选择最匹配的签名,如果函数参数没有适当的转换可用,则会引发TypeError。如果转换成功,则执行编译后的机器代码,并根据签名将返回值转换回来。

  • 如果没有给出 signature,被装饰的函数将实现延迟编译。每次调用被装饰的函数时,如果存在现有专用版本,将尝试重用它(例如,使用两个整数参数的调用可能会重用参数类型为(numba.int64, numba.int64)的专用版本)。如果不存在合适的专用版本,则会即时编译一个新的专用版本,存储以备后用,并使用转换后的参数执行。

如果为真,nopython 强制函数在nopython 模式下编译。如果不可能,编译将引发错误。

如果为真,forceobj 强制函数在对象模式下编译。由于对象模式比 nopython 模式慢,这主要用于测试目的。

如果为真,nogil 尝试在编译函数内部释放全局解释器锁。仅当 Numba 能够在nopython 模式下编译函数时,GIL 才会释放,否则将打印编译警告。

如果为真,cache 启用基于文件的缓存,以缩短函数在先前调用中已编译的编译时间。缓存维护在包含源文件的目录的__pycache__子目录中;但是,如果当前用户不允许写入,则会回退到平台特定的用户范围缓存目录(例如 Unix 平台上的$HOME/.cache/numba)。

如果为真,parallel 启用许多常见 NumPy 构造的自动并行化以及相邻并行操作的融合,以最大化缓存局部性。

error_model 选项控制除零行为。将其设置为“python”会导致除零引发像 CPython 一样的异常。将其设置为“numpy”会导致除零将结果设置为 +/-infnan

并非所有函数都可以缓存,因为某些功能无法总是持久化到磁盘。当函数无法缓存时,会发出警告。

如果为真,fastmath 启用使用 LLVM 文档中描述的否则不安全的浮点转换。此外,如果安装了 Intel SVML,则会使用更快但精度较低的一些数学内部函数版本(答案在 4 ULP 范围内)。

如果为真,boundscheck 为数组索引启用边界检查。越界访问将引发 IndexError。默认情况下不进行边界检查。如果禁用边界检查,越界访问可能会产生垃圾结果或段错误。然而,启用边界检查会使典型函数变慢,因此建议仅将此标志用于调试。您还可以将 NUMBA_BOUNDSCHECK 环境变量设置为 0 或 1 以全局覆盖此标志。

locals 字典可用于强制特定局部变量的类型和签名,例如,如果您想在某个点强制使用单精度浮点数。通常,我们建议您让 Numba 的编译器自行推断局部变量的类型。

以下是两个签名的示例:

@jit(["int32(int32)", "float32(float32)"], nopython=True)
def f(x): ...

装饰器后不加任何括号等同于不带任何参数调用装饰器,即

@jit
def f(x): ...

等同于

@jit()
def f(x): ...

该装饰器返回一个Dispatcher对象。

注意

如果未给出 signature,则在实际编译发生时(即函数首次使用给定参数类型调用时)将引发编译错误。

注意

编译可能会受到一些专门的环境变量的影响。

生成 JIT 函数

jit() 装饰器类似,但在编译时调用被装饰的函数,并传入函数参数的类型。被装饰的函数必须返回一个可调用对象,该对象将编译为这些类型的函数实现,从而允许灵活的专用化。

如果您正在寻找此功能,请参阅高级扩展 API 中的 @overload 系列装饰器。

调度器对象

class Dispatcher

通过调用jit()创建的对象类。您不应该尝试以任何其他方式创建此类对象。调用 Dispatcher 对象将调用针对所调用参数的编译专用版本,使其作为已编译 Python 函数的加速替代品。

此外,Dispatcher 对象具有以下方法和属性

py_func

被编译的纯 Python 函数。

inspect_types(file=None, pretty=False)

打印函数源代码列表,逐行注释对应的 Numba IR,以及推断的各种变量类型。如果指定了 file,则打印到该文件对象,否则打印到 sys.stdout。如果 pretty 设置为 True,则在终端中生成彩色 ANSI,在 notebook 中生成 HTML。

另请参阅

Numba 架构

inspect_llvm(signature=None)

返回一个字典,将编译后的函数签名映射到为该函数生成的易读 LLVM IR。如果指定了 signature 关键字,则返回对应单个签名的字符串。

inspect_asm(signature=None)

返回一个字典,将编译后的函数签名映射到该函数的人类可读本机汇编代码。如果指定了 signature 关键字,则返回对应单个签名的字符串。

inspect_cfg(signature=None, show_wrapped)

返回一个字典,将编译后的函数签名映射到该函数的控制流图对象。如果指定了 signature 关键字,则返回对应单个签名的字符串。

控制流图对象可以字符串化(strrepr)以获取 DOT 格式的图的文本表示。或者,使用其 .display(filename=None, view=False) 方法绘制图形。filename 选项可以设置为渲染输出写入的特定路径。如果 view 选项为 True,则由系统默认应用程序打开图像格式(PDF)的图形。在 IPython notebook 中,返回的对象可以内联绘制。

用法

@jit
def foo():
  ...

# opens the CFG in system default application
foo.inspect_cfg(foo.signatures[0]).display(view=True)
inspect_disasm_cfg(signature=None)

返回一个字典,将编译后的函数签名映射到基础编译ELF对象的反汇编控制流图。如果指定了 signature 关键字,则返回对应单个签名的控制流图。此函数具有执行环境感知能力,将在 Jupyter notebook 中生成 SVG 输出,在终端中生成 ASCII。

示例

@njit
def foo(x):
    if x < 3:
        return x + 1
    return x + 2

foo(10)

print(foo.inspect_disasm_cfg(signature=foo.signatures[0]))

给出

[0x08000040]>  # method.__main__.foo_241_long_long (int64_t arg1, int64_t arg3);
 ─────────────────────────────────────────────────────────────────────┐
│  0x8000040                                                          │
│ ; arg3 ; [02] -r-x section size 279 named .text                     │
│   ;-- section..text:                                                │
│   ;-- .text:                                                        │
│   ;-- __main__::foo$241(long long):                                 │
│   ;-- rip:                                                          │
│ 25: method.__main__.foo_241_long_long (int64_t arg1, int64_t arg3); │
│ ; arg int64_t arg1 @ rdi                                            │
│ ; arg int64_t arg3 @ rdx                                            │
│ ; 2                                                                 │
│ cmp rdx, 2                                                          │
│ jg 0x800004f                                                        │
└─────────────────────────────────────────────────────────────────────┘
        f t
        │ │
        │ └──────────────────────────────┐
        └──┐                             │
           │                             │
    ┌─────────────────────────┐   ┌─────────────────────────┐
    │  0x8000046              │   │  0x800004f              │
    │ ; arg3                  │   │ ; arg3                  │
    │ inc rdx                 │   │ add rdx, 2              │
    │ ; arg3                  │   │ ; arg3                  │
    │ mov qword [rdi], rdx    │   │ mov qword [rdi], rdx    │
    │ xor eax, eax            │   │ xor eax, eax            │
    │ ret                     │   │ ret                     │
    └─────────────────────────┘   └─────────────────────────┘
recompile()

重新编译所有现有签名。例如,如果全局或闭包变量被您的函数冻结并且其在 Python 中的值已更改,则这可能很有用。由于编译不便宜,这主要用于测试和交互式使用。

parallel_diagnostics(signature=None, level=1)

打印给定签名的并行诊断信息。如果不存在签名,则为所有已知签名打印。 level 用于调整详细程度,level=1(默认)是最小详细程度,级别 2、3 和 4 提供递增的详细程度。

get_metadata(signature=None)

获取给定签名的编译元数据。这对于 Numba 和 Numba 扩展的开发人员很有用。

向量化函数(ufuncs 和 DUFuncs)

@numba.vectorize(*, signatures=[], identity=None, nopython=True, target='cpu', forceobj=False, cache=False, locals={})

编译被装饰的函数,并将其包装为 NumPy ufunc 或 Numba DUFunc。可选的 nopythonforceobjlocals 参数与 numba.jit() 中的含义相同。

signatures 是签名的可选列表,其形式与 numba.jit()signature 参数相同。如果 signatures 非空,则装饰器将用户 Python 函数编译为 NumPy ufunc。如果没有给出 signatures,则装饰器将用户 Python 函数包装在 DUFunc 实例中,该实例将在调用时编译用户函数,只要 NumPy 找不到与输入参数匹配的循环。如果 target"parallel",则 signatures 是必需的。

identity 是所实现函数的恒等(或单位)值。可能的值为 0、1、None 和字符串 "reorderable"。默认值为 None。None 和 "reorderable" 都表示函数没有恒等值;"reorderable" 另外指定可以重新排列沿多个轴的归约。

如果有多个 signatures,它们必须按从更具体到最不具体的顺序排列。否则,NumPy 的基于类型的调度可能无法按预期工作。例如,以下是错误的:

@vectorize(["float64(float64)", "float32(float32)"])
def f(x): ...

因为对单精度数组运行它将选择编译函数的float64版本,导致效率大大降低。正确的调用是:

@vectorize(["float32(float32)", "float64(float64)"])
def f(x): ...

target 是后端目标的字符串;可用值有“cpu”、“parallel”和“cuda”。要使用多线程版本,将目标更改为“parallel”(需要指定签名)。

@vectorize(["float64(float64)", "float32(float32)"], target='parallel')
def f(x): ...

对于 CUDA 目标,使用“cuda”

@vectorize(["float64(float64)", "float32(float32)"], target='cuda')
def f(x): ...

编译后的函数可以缓存,以减少未来的编译时间。通过将 cache 设置为 True 来启用。只有“cpu”和“parallel”目标支持缓存。

此函数创建的 ufuncs 遵循 NEP-13,即 NumPy 覆盖 ufuncs 的机制。如果 ufunc 的 __call__ 的任何参数具有 __array_ufunc__ 方法,则将调用该方法(在 Python 中,而不是编译上下文中),该方法可能会预处理和/或后处理编译后的 ufunc 的参数和返回值(或者甚至可能不调用它)。

@numba.guvectorize(signatures, layout, *, identity=None, nopython=True, target='cpu', forceobj=False, cache=False, locals={})

numba.vectorize() 的通用版本。虽然 numba.vectorize() 将生成一个简单的 ufunc,其核心功能(您正在装饰的函数)操作标量操作数并返回标量值,但 numba.guvectorize() 允许您创建一个 NumPy ufunc,其核心功能接受各种维度的数组参数。

额外的参数 layout 是一个字符串,以符号形式指定参数类型和返回类型的维度和大小关系。例如,矩阵乘法将具有 "(m,n),(n,p)->(m,p)" 的布局字符串。它的定义可能是(函数体省略):

@guvectorize(["void(float64[:,:], float64[:,:], float64[:,:])"],
             "(m,n),(n,p)->(m,p)")
def f(a, b, result):
    """Fill-in *result* matrix such as result := a * b"""
    ...

如果其中一个参数应该是标量,则相应的布局规范是 (),并且该参数将真正以零维数组的形式提供给您(您必须解引用它才能获取标量值)。例如,一个带可参数化窗口宽度的一维移动平均可能具有 "(n),()->(n)" 的布局字符串。

请注意,任何输出都将预先分配为附加函数参数:您的代码必须用适当的值填充它以实现您正在实现的函数。

如果您的函数不接受输出数组,您应该省略布局字符串中的“箭头”(例如 "(n),(n)")。这样做时,重要的是要注意对输入数组的更改不能总是依赖于在 ufunc 执行之外可见,因为 NumPy 可能会传入临时数组作为输入(例如,如果需要类型转换)。

另请参阅

NumPy 支持的布局字符串规范。请注意,NumPy 使用“signature”一词,我们不幸地将其用于其他用途。

编译后的函数可以缓存,以减少未来的编译时间。通过将 cache 设置为 True 来启用。只有“cpu”和“parallel”目标支持缓存。

class numba.DUFunc

通过不带签名的 numba.vectorize() 调用创建的对象类。

DUFunc 实例的行为应与 NumPy ufunc 对象类似,但有一个重要区别:调用时循环生成。调用 ufunc 时,NumPy 会查看为该 ufunc 注册的现有循环,如果找不到它可以安全地将输入转换为适合的循环,则会引发 TypeError。调用 DUFunc 时,Numba 将调用委托给 NumPy。如果 NumPy ufunc 调用失败,则 Numba 尝试为给定的输入类型构建一个新循环,并再次调用 ufunc。如果第二次调用尝试失败或发生编译错误,DUFunc 会将异常传递给调用者。

另请参阅

用户指南中的“动态通用函数”部分演示了 DUFunc 的调用时行为,并讨论了调用顺序对 Numba 生成底层 ufunc 的影响。

ufunc

DUFunc 实例构建的实际 NumPy ufunc 对象。请注意,DUFunc 对象维护了几个重要的数据结构,这些结构是 ufunc 正常功能所必需的(特别是动态编译的循环)。用户不应在不确保底层 DUFunc 不会被垃圾回收的情况下传递 ufunc 值。

nin

DUFunc (ufunc) 输入的数量。请参阅 ufunc.nin

nout

DUFunc 输出的数量。请参阅 ufunc.nout

nargs

DUFunc 参数的总数(应为 nin + nout)。请参阅 ufunc.nargs

ntypes

DUFunc 支持的输入类型数量。请参阅 ufunc.ntypes

types

支持的类型列表,以字符串形式给出。请参阅 ufunc.types

identity

将 ufunc 作为归约使用时的恒等值。请参阅 ufunc.identity

reduce(A, *, axis, dtype, out, keepdims)

通过沿一个轴应用 DUFunc 将 A 的维度减少一。请参阅 ufunc.reduce

accumulate(A, *, axis, dtype, out)

累积将运算符应用于所有元素的结果。请参阅 ufunc.accumulate

reduceat(A, indices, *, axis, dtype, out)

在单个轴上执行指定切片的(局部)归约。请参阅 ufunc.reduceat

outer(A, B)

A 中的所有 aB 中的所有 b 应用 ufunc 到所有对 (a, b)。请参阅 ufunc.outer

at(A, indices, *, B)

A 中由 indices 指定的元素执行非缓冲原地操作。如果您使用的是 NumPy 1.7 或更早版本,则此方法将不存在。请参阅 ufunc.at

注意

在极少数情况下,向量化函数可能会显示意外的警告或错误

C 回调

@numba.cfunc(signature, nopython=False, cache=False, locals={})

即时编译被装饰的函数以生成高效的机器代码。编译后的代码被封装在一个薄的 C 回调中,使其可以使用自然的 C ABI 调用。

signature 是单个签名,表示 C 回调的签名。其形式必须与 jit() 中的形式相同。装饰器不检查签名中的类型在 C 中是否具有明确定义的表示。

nopythoncache 是布尔标志。locals 是局部变量名到类型和签名的映射。它们的含义与 jit() 中的含义相同。

该装饰器返回一个CFunc对象。

注意

C 回调目前不支持对象模式

class CFunc

cfunc() 创建的对象类。CFunc 对象公开以下属性和方法:

address

编译后的 C 回调的地址,以整数形式。

cffi

一个 cffi 函数指针实例,用于作为参数传递给 cffi 包装的函数。该指针的类型是 void *,因此在将其传递给 cffi 时只会进行最少的类型检查。

ctypes

一个 ctypes 回调实例,就像它使用 ctypes.CFUNCTYPE() 创建一样。

native_name

编译后的 C 回调的名称。

inspect_llvm()

返回为 C 回调生成的易读 LLVM IR。native_name 是此回调在 IR 中定义的名称。