使用 `@jitclass` 编译 Python 类

注意

这是 jitclass 支持的早期版本。并非所有编译功能都已公开或实现。

Numba 通过 numba.experimental.jitclass() 装饰器支持类的代码生成。一个类可以使用此装饰器标记进行优化，并指定每个字段的类型。我们将生成的类对象称为 jitclass。jitclass 的所有方法都被编译为 nopython 函数。jitclass 实例的数据在堆上分配为 C 兼容结构，以便任何编译函数可以直接访问底层数据，从而绕过解释器。

基本用法

这是一个 jitclass 的示例

import numpy as np
from numba import int32, float32    # import the types
from numba.experimental import jitclass

spec = [
    ('value', int32),               # a simple scalar field
    ('array', float32[:]),          # an array field
]

@jitclass(spec)
class Bag(object):
    def __init__(self, value):
        self.value = value
        self.array = np.zeros(value, dtype=np.float32)

    @property
    def size(self):
        return self.array.size

    def increment(self, val):
        for i in range(self.size):
            self.array[i] += val
        return self.array

    @staticmethod
    def add(x, y):
        return x + y

n = 21
mybag = Bag(n)

在上面的示例中，spec 以一个包含 2-元组的列表形式提供。这些元组包含字段的名称和字段的 Numba 类型。或者，用户可以使用字典（最好是 OrderedDict 以保持字段顺序稳定），将字段名称映射到类型。

类的定义至少需要一个 __init__ 方法来初始化每个已定义的字段。未初始化的字段包含垃圾数据。可以定义方法和属性（仅限 getter 和 setter）。它们将自动编译。

从类型注解中推断类成员类型，使用 `as_numba_type`

jitclass 的字段也可以从 Python 类型注解中推断。

from typing import List
from numba.experimental import jitclass
from numba.typed import List as NumbaList

@jitclass
class Counter:
    value: int

    def __init__(self):
        self.value = 0

    def get(self) -> int:
        ret = self.value
        self.value += 1
        return ret

@jitclass
class ListLoopIterator:
    counter: Counter
    items: List[float]

    def __init__(self, items: List[float]):
        self.items = items
        self.counter = Counter()

    def get(self) -> float:
        idx = self.counter.get() % len(self.items)
        return self.items[idx]

items = NumbaList([3.14, 2.718, 0.123, -4.])
loop_itr = ListLoopIterator(items)

类上的任何类型注解都将用于扩展 spec，如果该字段尚不存在。Numba 类型与给定的 Python 类型之间的对应关系使用 as_numba_type 进行推断。例如，如果存在以下类

@jitclass([("w", int32), ("y", float64[:])])
class Foo:
    w: int
    x: float
    y: np.ndarray
    z: SomeOtherType

    def __init__(self, w: int, x: float, y: np.ndarray, z: SomeOtherType):
        ...

那么用于 Foo 的完整 spec 将是

"w": int32 （在 spec 中指定）
"x": float64 （从类型注解添加）
"y": array(float64, 1d, A) （在 spec 中指定）
"z": numba.as_numba_type(SomeOtherType) （从类型注解添加）

这里 SomeOtherType 可以是任何支持的 Python 类型（例如 bool, typing.Dict[int, typing.Tuple[float, float]]，或另一个 jitclass）。

请注意，只有类上的类型注解将用于推断 spec 元素。方法类型注解（例如上面 __init__ 的注解）将被忽略。

Numba 需要知道 NumPy 数组的 dtype 和 rank，目前无法通过类型注解来表达。因此，NumPy 数组需要明确地包含在 spec 中。

明确指定 `numba.typed` 容器作为类成员

以下模式展示了如何明确指定 numba.typed.Dict 或 numba.typed.List 作为传递给 jitclass 的 spec 的一部分。

首先，使用显式的 Numba 类型和显式构造。

from numba import types, typed
from numba.experimental import jitclass

# key and value types
kv_ty = (types.int64, types.unicode_type)

# A container class with:
# * member 'd' holding a typed dictionary of int64 -> unicode string (kv_ty)
# * member 'l' holding a typed list of float64
@jitclass([('d', types.DictType(*kv_ty)),
           ('l', types.ListType(types.float64))])
class ContainerHolder(object):
    def __init__(self):
        # initialize the containers
        self.d = typed.Dict.empty(*kv_ty)
        self.l = typed.List.empty_list(types.float64)

container = ContainerHolder()
container.d[1] = "apple"
container.d[2] = "orange"
container.l.append(123.)
container.l.append(456.)
print(container.d) # {1: apple, 2: orange}
print(container.l) # [123.0, 456.0]

另一个有用的模式是使用 numba.typed 容器属性 _numba_type_ 来查找容器的类型，可以直接从 Python 解释器中容器的实例访问。通过在实例上调用 numba.typeof() 也可以获取相同的信息。例如

from numba import typed, typeof
from numba.experimental import jitclass

d = typed.Dict()
d[1] = "apple"
d[2] = "orange"
l = typed.List()
l.append(123.)
l.append(456.)


@jitclass([('d', typeof(d)), ('l', typeof(l))])
class ContainerInstHolder(object):
    def __init__(self, dict_inst, list_inst):
        self.d = dict_inst
        self.l = list_inst

container = ContainerInstHolder(d, l)
print(container.d) # {1: apple, 2: orange}
print(container.l) # [123.0, 456.0]

值得注意的是，jitclass 中容器的实例在使用前必须初始化，例如，这将导致无效的内存访问，因为在 d 未被初始化为指定类型的 type.Dict 实例的情况下，self.d 被写入。

from numba import types
from numba.experimental import jitclass

dict_ty = types.DictType(types.int64, types.unicode_type)

@jitclass([('d', dict_ty)])
class NotInitialisingContainer(object):
    def __init__(self):
        self.d[10] = "apple" # this is invalid, `d` is not initialized

NotInitialisingContainer() # segmentation fault/memory access violation

支持的操作

jitclass 的以下操作在解释器和 Numba 编译函数中均可工作

调用 jitclass 类对象以构造新实例（例如 mybag = Bag(123)）；
读/写访问属性（例如 mybag.value）；
调用方法（例如 mybag.increment(3)）；
将静态方法作为实例属性调用（例如 mybag.add(1, 1)）；
将静态方法作为类属性调用（例如 Bag.add(1, 2)）；
使用特定的“魔术方法”（例如 __add__ 与 mybag + otherbag）；

在 Numba 编译函数中使用 jitclass 更高效。短方法可以内联（由 LLVM 内联器决定）。属性访问只是从 C 结构中读取。从解释器使用 jitclass 的开销与从解释器调用任何 Numba 编译函数的开销相同。参数和返回值必须在 Python 对象和原生表示之间进行拆箱或装箱。当 jitclass 实例被交给解释器时，由 jitclass 封装的值不会被装箱成 Python 对象。它们是在对字段值进行属性访问时才被装箱的。将静态方法作为类属性调用仅在类定义之外受支持（即，代码不能在 Bag 的另一个方法中调用 Bag.add()）。

支持的魔术方法

jitclass 可以定义以下魔术方法

__abs__
__bool__
__complex__
__contains__
__float__
__getitem__
__hash__
__index__
__int__
__len__
__setitem__
__str__
__eq__
__ne__
__ge__
__gt__
__le__
__lt__
__add__
__floordiv__
__lshift__
__matmul__
__mod__
__mul__
__neg__
__pos__
__pow__
__rshift__
__sub__
__truediv__
__and__
__or__
__xor__
__iadd__
__ifloordiv__
__ilshift__
__imatmul__
__imod__
__imul__
__ipow__
__irshift__
__isub__
__itruediv__
__iand__
__ior__
__ixor__
__radd__
__rfloordiv__
__rlshift__
__rmatmul__
__rmod__
__rmul__
__rpow__
__rrshift__
__rsub__
__rtruediv__
__rand__
__ror__
__rxor__

有关这些方法的描述，请参阅 Python 数据模型文档。

限制

jitclass 类对象在 Numba 编译函数内部被视为一个函数（即构造函数）。
isinstance() 仅在解释器中工作。
目前，在解释器中操作 jitclass 实例尚未优化。
目前仅在 CPU 上支持 jitclass。（注意：GPU 设备的支持计划在未来的版本中提供。）

装饰器：`@jitclass`

numba.experimental.jitclass(cls_or_spec=None, spec=None)

用于创建 jitclass 的函数。可以用作装饰器或函数。

不同的用例会导致设置不同的参数。

如果指定，spec 提供类字段的类型。它必须是字典或序列。如果使用字典，请使用 collections.OrderedDict 以确保稳定的顺序。如果使用序列，它必须包含 (字段名, 字段类型) 的 2-元组。

任何未在 spec 中列出的字段名的类注解都将被添加。对于类注解 x: T，如果 x 尚未是 spec 中的键，我们将把 ("x", as_numba_type(T)) 追加到 spec 中。

返回

如果用作装饰器，返回一个可调用对象，该对象接受一个类对象并
返回一个编译版本。
如果用作函数，返回编译后的类（JitClassType 的一个实例
JitClassType).

示例

cls_or_spec = None, spec = None

>>> @jitclass()
... class Foo:
...     ...

cls_or_spec = None, spec = spec

>>> @jitclass(spec=spec)
... class Foo:
...     ...

cls_or_spec = Foo, spec = None

>>> @jitclass
... class Foo:
...     ...

4) cls_or_spec = spec, spec = None 在这种情况下，我们更新 cls_or_spec, spec = None, cls_or_spec。

>>> @jitclass(spec)
... class Foo:
...     ...

cls_or_spec = Foo, spec = spec

>>> JitFoo = jitclass(Foo, spec)

使用 @jitclass 编译 Python 类