NBEP 2: 扩展点

作者:

安托万·皮特鲁

日期:

2015年7月

状态:

草稿

在 Numba 中实现新类型或函数需要挂钩到编译链中的各种机制（可能在编译链之外）。本文首先旨在考察当前的实现方式，其次提出使扩展更容易的建议。

如果某些提案被实施，我们应该首先在内部使用和实践它们，然后再向公众公开API。

备注

本文档不涉及CUDA或其他任何非CPU后端。

高级 API

目前没有高级API，使得某些用例比它们应有的更加复杂。

提议的更改

专用模块

我们提议增加一个 numba.extending 模块，该模块将公开用于扩展 Numba 的主要 API。

实现一个函数

我们建议增加一个 @overload 装饰器，允许为 nopython 模式 实现给定函数的重载。重载函数具有与实现函数相同的正式签名，并接收实际的参数类型。它应该返回一个为给定类型实现重载函数的 Python 函数。

以下示例使用这种方法实现了 numpy.where()。

import numpy as np

from numba.core import types
from numba.extending import overload

@overload(np.where)
def where(cond, x, y):
    """
    Implement np.where().
    """
    # Choose implementation based on argument types.
    if isinstance(cond, types.Array):
        # Array where() => return an array of the same shape
        if all(ty.layout == 'C' for ty in (cond, x, y)):
            def where_impl(cond, x, y):
                """
                Fast implementation for C-contiguous arrays
                """
                shape = cond.shape
                if x.shape != shape or y.shape != shape:
                    raise ValueError("all inputs should have the same shape")
                res = np.empty_like(x)
                cf = cond.flat
                xf = x.flat
                yf = y.flat
                rf = res.flat
                for i in range(cond.size):
                    rf[i] = xf[i] if cf[i] else yf[i]
                return res
        else:
            def where_impl(cond, x, y):
                """
                Generic implementation for other arrays
                """
                shape = cond.shape
                if x.shape != shape or y.shape != shape:
                    raise ValueError("all inputs should have the same shape")
                res = np.empty_like(x)
                for idx, c in np.ndenumerate(cond):
                    res[idx] = x[idx] if c else y[idx]
                return res

    else:
        def where_impl(cond, x, y):
            """
            Scalar where() => return a 0-dim array
            """
            scal = x if cond else y
            return np.full_like(scal, scal)

    return where_impl

也可以实现Numba已知的函数，以支持额外的类型。以下示例使用这种方法为元组实现了内置函数 len():

@overload(len)
def tuple_len(x):
   if isinstance(x, types.BaseTuple):
      # The tuple length is known at compile-time, so simply reify it
      # as a constant.
      n = len(x)
      def len_impl(x):
         return n
      return len_impl

实现一个属性

我们建议增加一个 @overload_attribute 装饰器，允许在 nopython 模式 中实现属性获取器。

以下示例在 Numpy 数组上实现了 .nbytes 属性:

@overload_attribute(types.Array, 'nbytes')
def array_nbytes(arr):
   def get(arr):
       return arr.size * arr.itemsize
   return get

备注

overload_attribute() 签名允许通过让装饰函数返回一个 getter, setter, deleter 元组来扩展定义设置器和删除器，而不是返回单个 getter。

实现一个方法

我们提议添加一个 @overload_method 装饰器，允许在 nopython 模式 中实现实例方法。

以下示例在 Numpy 数组上实现了 .take() 方法:

@overload_method(types.Array, 'take')
def array_take(arr, indices):
   if isinstance(indices, types.Array):
       def take_impl(arr, indices):
           n = indices.shape[0]
           res = np.empty(n, arr.dtype)
           for i in range(n):
               res[i] = arr[indices[i]]
           return res
       return take_impl

暴露结构成员

我们建议添加一个 make_attribute_wrapper() 函数，将内部字段作为可见的只读属性暴露出来，适用于那些由 StructModel 数据模型支持的类型。

例如，假设 PdIndexType 是 pandas 索引的 Numba 类型，以下是如何将底层 Numpy 数组作为 ._data 属性暴露出来:

@register_model(PdIndexType)
class PdIndexModel(models.StructModel):
    def __init__(self, dmm, fe_type):
        members = [
            ('values', fe_type.as_array),
            ]
        models.StructModel.__init__(self, dmm, fe_type, members)

make_attribute_wrapper(PdIndexType, 'values', '_data')

打字

Numba 类型

Numba 的标准类型在 numba.types 中声明。要声明一个新类型，可以通过子类化基类 Type 或其现有的抽象子类之一，并实现所需的功能。

提议的更改

无需更改。

值的类型推断

如果新类型的值可以作为函数参数或常量出现，则需要进行类型推断。核心机制在 numba.typing.typeof 中。

在某些情况下，某些 Python 类或类专门映射到新类型，可以扩展通用函数以在这些类上进行调度，例如:

from numba.typing.typeof import typeof_impl

@typeof_impl(MyClass)
def _typeof_myclass(val, c):
   if "some condition":
      return MyType(...)

typeof_impl 特化必须返回一个 Numba 类型实例，如果值类型推导失败则返回 None。

（当控制被类型推断的类时，一种替代 typeof_impl 的方法是在类上定义一个 _numba_type_ 属性）

在更罕见的情况下，新类型可以表示无法枚举的各种Python类，必须在 typeof_impl 泛型函数的回退实现中插入手动检查。

提议的更改

允许人们定义一个通用钩子，而无需对回退实现进行猴子补丁。

函数参数类型推断的快速路径

可选地，可能希望允许新类型参与快速类型解析（用C代码编写），以在用新类型调用JIT编译函数时最小化函数调用开销。然后必须在 _typeof.c 文件中插入所需的检查和实现，大概在 compute_fingerprint() 函数内部。

提议的更改

None. 在C Python扩展中嵌入的C代码中添加通用钩子是一个过于微妙的更改。

操作中的类型推断

各种操作（函数调用、运算符等）产生的结果值使用一组称为“模板”的帮助器进行类型化。可以通过子类化现有的基类并实现所需的推理机制来定义新模板。使用装饰器显式地将模板注册到类型推理机制中。

通过 ConcreteTemplate 基类，可以为给定操作定义一组支持的签名来实现推理。以下示例为取模运算符进行类型标注:

@builtin
class BinOpMod(ConcreteTemplate):
    key = "%"
    cases = [signature(op, op, op)
             for op in sorted(types.signed_domain)]
    cases += [signature(op, op, op)
              for op in sorted(types.unsigned_domain)]
    cases += [signature(op, op, op) for op in sorted(types.real_domain)]

（请注意，签名中使用了类型 实例，这严重限制了可以表达的泛型数量）

通过 AbstractTemplate 基类，可以以编程方式定义推理，从而赋予其完全的灵活性。以下是一个简单的示例，展示了如何表达元组索引（即 __getitem__ 操作符）:

@builtin
class GetItemUniTuple(AbstractTemplate):
    key = "getitem"

    def generic(self, args, kws):
        tup, idx = args
        if isinstance(tup, types.UniTuple) and isinstance(idx, types.Integer):
            return signature(tup.dtype, tup, idx)

The AttributeTemplate 基类允许对给定类型的属性和方法进行类型化。以下是一个示例，对复数的 .real 和 .imag 属性进行类型化:

@builtin_attr
class ComplexAttribute(AttributeTemplate):
    key = types.Complex

    def resolve_real(self, ty):
        return ty.underlying_float

    def resolve_imag(self, ty):
        return ty.underlying_float

备注

AttributeTemplate 仅用于获取属性。设置属性的值在 numba.typeinfer 中是硬编码的。

通过让用户定义一个与正在输入的函数具有相同定义的可调用对象，CallableTemplate 基类提供了一种更简单的方式来解析灵活的函数签名。例如，如果 Numba 支持列表，这里是如何假设性地输入 Python 的 sorted 函数:

@builtin
class Sorted(CallableTemplate):
    key = sorted

    def generic(self):
        def typer(iterable, key=None, reverse=None):
            if reverse is not None and not isinstance(reverse, types.Boolean):
                return
            if key is not None and not isinstance(key, types.Callable):
                return
            if not isinstance(iterable, types.Iterable):
                return
            return types.List(iterable.iterator_type.yield_type)

        return typer

（注意你可以只返回函数的返回类型，而不是完整的签名）

提议的更改

各种装饰器的命名相当模糊且令人困惑。我们建议将 @builtin 重命名为 @infer，将 @builtin_attr 重命名为 @infer_getattr，并将 builtin_global 重命名为 infer_global。

全局值的两步声明有点冗长，我们建议通过允许使用 infer_global 作为装饰器来简化它:

@infer_global(len)
class Len(AbstractTemplate):
    key = len

    def generic(self, args, kws):
        assert not kws
        (val,) = args
        if isinstance(val, (types.Buffer, types.BaseTuple)):
            return signature(types.intp, val)

基于类的API可能会显得笨拙，我们可以为某些模板类型添加一个功能性API：

@type_callable(sorted)
def type_sorted(context):
    def typer(iterable, key=None, reverse=None):
        # [same function as above]

    return typer

代码生成

Numba 类型的值的具体表示

任何具体的 Numba 类型都必须能够以 LLVM 形式表示（用于变量存储、参数传递等）。通过实现一个数据模型类并使用装饰器注册它来定义该表示。标准类型的数据模型类在 numba.datamodel.models 中定义。

提议的更改

无需更改。

类型转换

Numba 类型之间的隐式转换目前是在 BaseContext.cast() 方法中作为一系列选择和类型检查实现的。要添加新的隐式转换，可以在该方法中追加一个特定类型的检查。

布尔值评估是隐式转换的一个特例（目标类型为 types.Boolean）。

备注

显式转换被视为常规操作，例如构造函数调用。

提议的更改

添加一个用于隐式转换的通用函数，基于源类型和目标类型的多重分派。以下是一个展示如何编写浮点数到整数转换的示例:

@lower_cast(types.Float, types.Integer)
def float_to_integer(context, builder, fromty, toty, val):
    lty = context.get_value_type(toty)
    if toty.signed:
        return builder.fptosi(val, lty)
    else:
        return builder.fptoui(val, lty)

操作的实现

其他操作是通过一组通用函数和装饰器来实现和注册的。例如，以下是如何实现对 Numpy 数组上 .ndim 属性的查找:

@builtin_attr
@impl_attribute(types.Kind(types.Array), "ndim", types.intp)
def array_ndim(context, builder, typ, value):
    return context.get_constant(types.intp, typ.ndim)

以下是如何在元组值上调用 len() 的实现方式:

@builtin
@implement(types.len_type, types.Kind(types.BaseTuple))
def tuple_len(context, builder, sig, args):
    tupty, = sig.args
    retty = sig.return_type
    return context.get_constant(retty, len(tupty.types))

提议的更改

审查并优化API。去掉显式编写 types.Kind(...) 的要求。移除单独的 @implement 装饰器，并将 @builtin 重命名为 @lower_builtin，@builtin_attr 重命名为 @lower_getattr 等。

添加装饰器以实现 setattr() 操作，命名为 @lower_setattr 和 @lower_setattr_generic。

从/到 Python 对象的转换

某些类型需要从或转换为 Python 对象，如果它们可以作为函数参数传递或从函数返回。相应的装箱和拆箱操作使用泛型函数实现。标准 Numba 类型的实现位于 numba.targets.boxing 中。例如，以下是布尔值的装箱实现:

@box(types.Boolean)
def box_bool(c, typ, val):
    longval = c.builder.zext(val, c.pyapi.long)
    return c.pyapi.bool_from_long(longval)

提议的更改

将实现签名从 (c, typ, val) 改为 (typ, val, c)，以匹配为 typeof_impl 通用函数选择的签名。