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

如果类上存在任何类型注解，并且该字段尚未存在，则将使用这些注解来扩展规范。使用 as_numba_type 推断与给定 Python 类型对应的 Numba 类型。例如，如果我们有这个类

@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 的完整规范将是：

• "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）。

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

Numba 需要知道 NumPy 数组的 dtype 和秩，这在当前无法通过类型注解来表达。因此，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 中的容器实例在使用前必须进行初始化，例如，这将导致无效的内存访问，因为 self.d 在被写入时，d 并未被初始化为指定类型的 type.Dict 实例。

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()``）。

支持的双下划线方法

以下是可能为 jitclasses 定义的 dunder 方法：

• __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 数据模型文档。

限制

• 在Numba编译的函数中，jitclass类对象被视为一个函数（构造函数）。

• isinstance() 仅在解释器中有效。

• 在解释器中操作 jitclass 实例尚未优化。

• jitclass 的支持仅适用于 CPU。（注意：对 GPU 设备的支持计划在未来版本中提供。）

装饰器: @jitclass

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

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

不同的使用场景会导致设置不同的参数。

如果指定，spec 给出类字段的类型。它必须是一个字典或序列。使用字典时，使用 collections.OrderedDict 以保持顺序稳定。使用序列时，它必须包含 (字段名, 字段类型) 的 2-元组。

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

返回:

如果用作装饰器，返回一个可调用对象，该对象接受一个类对象并

返回编译后的版本。

如果作为函数使用，返回编译后的类（一个实例）。

JitClassType).

示例

1. cls_or_spec = None, spec = None

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

2. cls_or_spec = None, spec = spec

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

3. cls_or_spec = Foo, spec = None

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

4) cls_or_spec = spec, spec = None In this case we update cls_or_spec, spec = None, cls_or_spec.

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

5. cls_or_spec = Foo，spec = spec

>>> JitFoo = jitclass(Foo, spec)