6.3。示例：间隔类型

原文： http://numba.pydata.org/numba-doc/latest/extending/interval-example.html

我们将扩展 Numba 前端以支持它目前不支持的类，以便允许：

  • 将类的实例传递给 Numba 函数
  • 在 Numba 函数中访问类的属性
  • 从 Numba 函数构造并返回类的新实例

（以上所有 nopython 模式）

我们将混合来自高级扩展 API 和低级扩展 API 的 API，具体取决于给定任务的可用内容。

我们的示例的起点是以下纯 Python 类：

  1. class Interval(object):
  2.     """
  3.     A half-open interval on the real number line.
  4.     """
  5.     def __init__(self, lo, hi):
  6.         self.lo = lo
  7.         self.hi = hi
  9.     def __repr__(self):
  10.         return 'Interval(%f, %f)' % (self.lo, self.hi)
  12.     @property
  13.     def width(self):
  14.         return self.hi - self.lo

6.3.1。扩展打字层

6.3.1.1。创建一个新的 Numba 类型

由于 Numba 不知道Interval类，我们必须创建一个新的 Numba 类型来表示它的实例。 Numba 不直接处理 Python 类型：它有自己的类型系统，允许不同级别的粒度以及常规 Python 类型不具备的各种元信息。

我们首先创建一个类型类IntervalType，因为我们不需要参数类型，所以我们实例化一个类型实例interval_type：

  1. from numba import types
  3. class IntervalType(types.Type):
  4.     def __init__(self):
  5.         super(IntervalType, self).__init__(name='Interval')
  7. interval_type = IntervalType()

6.3.1.2。 Python 值的类型推断

就其本身而言，创建一个 Numba 类型并没有做任何事情。我们必须教 Numba 如何推断一些 Python 值作为该类型的实例。在这个例子中，它是微不足道的：Interval类的任何实例都应被视为属于interval_type类型：

  1. from numba.extending import typeof_impl
  3. @typeof_impl.register(Interval)
  4. def typeof_index(val, c):
  5.     return interval_type

因此，只要它们是Interval的实例，函数参数和全局值就会被认为属于interval_type。

6.3.1.3。操作类型推断

我们希望能够从 Numba 函数构造区间对象，因此我们必须教 Numba 识别双参数Interval(lo, hi)构造函数。参数应该是浮点数：

  1. from numba.extending import type_callable
  3. @type_callable(Interval)
  4. def type_interval(context):
  5.     def typer(lo, hi):
  6.         if isinstance(lo, types.Float) and isinstance(hi, types.Float):
  7.             return interval_type
  8.     return typer

type_callable() 装饰器指定在为给定的可调用对象（此处为Interval类本身）运行类型推断时应调用修饰函数。修饰函数必须简单地返回一个将使用参数类型调用的 typer 函数。这个看似错综复杂的设置的原因是，typer 函数使 与 完全相同的签名与打字的 callable 相同。这允许正确处理关键字参数。

由装饰函数接收的 上下文 参数在更复杂的情况下非常有用，在这种情况下，计算可调用的返回类型需要解析其他类型。

6.3.2。延伸降低层

我们已经完成了关于我们的类型推理添加的 Numba 教学。我们现在必须教 Numba 如何为新操作实际生成代码和数据。

6.3.2.1。定义原生间隔的数据模型

作为一般规则， nopython 模式不适用于 Python 对象，因为它们是由 CPython 解释器生成的。解释器使用的表示对于快速本机代码而言效率太低。因此， nopython 模式中支持的每种类型都必须定义定制的本机表示，也称为 数据模型 。

数据模型的一个常见情况是不可变的类似结构的数据模型，类似于 C struct。我们的 interval 数据类型很方便地属于该类别，这里有一个可能的数据模型：

  1. from numba.extending import models, register_model
  3. @register_model(IntervalType)
  4. class IntervalModel(models.StructModel):
  5.     def __init__(self, dmm, fe_type):
  6.         members = [
  7.             ('lo', types.float64),
  8.             ('hi', types.float64),
  9.             ]
  10.         models.StructModel.__init__(self, dmm, fe_type, members)

这指示 Numba 将IntervalType类型的值（或其任何实例）表示为两个字段lo和hi的结构，每个字段都是双精度浮点数（types.float64）。

注意

可变类型需要更复杂的数据模型才能在修改后保留其值。它们通常不能存储和传递到堆栈或寄存器，如不可变类型。

6.3.2.2。公开数据模型属性

我们希望数据模型属性lo和hi以相同的名称公开，以便在 Numba 函数中使用。 Numba 提供了一个方便的功能来做到这一点：

  1. from numba.extending import make_attribute_wrapper
  3. make_attribute_wrapper(IntervalType, 'lo', 'lo')
  4. make_attribute_wrapper(IntervalType, 'hi', 'hi')

这将以只读模式公开属性。如上所述，可写属性不适合此模型。

6.3.2.3。公开财产

由于width属性是计算而不是存储在结构中，我们不能像我们对lo和hi那样简单地公开它。我们必须明确地重新实现它：

  1. from numba.extending import overload_attribute
  3. @overload_attribute(IntervalType, "width")
  4. def get_width(interval):
  5.     def getter(interval):
  6.         return interval.hi - interval.lo
  7.     return getter

您可能会问为什么我们不需要为此属性公开类型推断钩子？答案是@overload_attribute是高级 API 的一部分：它在单个 API 中结合了类型推断和代码生成。

6.3.2.4。实现构造函数

现在我们要实现两个参数Interval构造函数：

  1. from numba.extending import lower_builtin
  2. from numba import cgutils
  4. @lower_builtin(Interval, types.Float, types.Float)
  5. def impl_interval(context, builder, sig, args):
  6.     typ = sig.return_type
  7.     lo, hi = args
  8.     interval = cgutils.create_struct_proxy(typ)(context, builder)
  9.     interval.lo = lo
  10.     interval.hi = hi
  11.     return interval._getvalue()

还有一点在这里发生。 @lower_builtin为某些特定的参数类型修饰给定的可调用或操作（这里是Interval构造函数）的实现。这允许定义给定操作的特定于类型的实现，这对于重载过多的函数很重要，例如 len() 。

types.Float是所有浮点类型的类（types.float64是types.Float的实例）。在类上而不是在特定实例上匹配参数类型通常更具有前瞻性（但是，当 返回 类型时 - 主要是在类型推断阶段 - ，您通常必须返回一个类型实例）。

由于 Numba 如何传递价值，cgutils.create_struct_proxy()和interval._getvalue()是一个样板。值作为 llvmlite.ir.Value 的实例传递，这可能太有限：LLVM 结构值尤其是非常低级别。结构代理是 LLVM 结构值的临时包装，允许轻松获取或设置结构的成员。 _getvalue()调用只是从包装器中获取 LLVM 值。

6.3.2.5。拳击和拆箱

如果此时尝试使用Interval实例，则肯定会得到错误 “无法将 Interval 转换为本机值”。这是因为 Numba 还不知道如何从 Python Interval实例创建本机间隔值。让我们教它如何做到：

  1. from numba.extending import unbox, NativeValue
  3. @unbox(IntervalType)
  4. def unbox_interval(typ, obj, c):
  5.     """
  6.     Convert a Interval object to a native interval structure.
  7.     """
  8.     lo_obj = c.pyapi.object_getattr_string(obj, "lo")
  9.     hi_obj = c.pyapi.object_getattr_string(obj, "hi")
  10.     interval = cgutils.create_struct_proxy(typ)(c.context, c.builder)
  11.     interval.lo = c.pyapi.float_as_double(lo_obj)
  12.     interval.hi = c.pyapi.float_as_double(hi_obj)
  13.     c.pyapi.decref(lo_obj)
  14.     c.pyapi.decref(hi_obj)
  15.     is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
  16.     return NativeValue(interval._getvalue(), is_error=is_error)

Unbox 是“将 Python 对象转换为本机值”的另一个名称（它将 Python 对象的想法视为包含简单本机值的复杂框）。该函数返回一个NativeValue对象，该对象使其调用者可以访问计算的本机值，错误位以及可能的其他信息。

上面的代码片段充分利用了c.pyapi对象，可以访问 Python 解释器的 C API 的子集。注意使用c.pyapi.err_occurred()来检测拆箱时可能发生的任何错误（例如尝试传递Interval('a', 'b')）。

我们还想做反向操作，称为 拳击 ，以便从 Numba 函数返回间隔值：

  1. from numba.extending import box
  3. @box(IntervalType)
  4. def box_interval(typ, val, c):
  5.     """
  6.     Convert a native interval structure to an Interval object.
  7.     """
  8.     interval = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
  9.     lo_obj = c.pyapi.float_from_double(interval.lo)
  10.     hi_obj = c.pyapi.float_from_double(interval.hi)
  11.     class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(Interval))
  12.     res = c.pyapi.call_function_objargs(class_obj, (lo_obj, hi_obj))
  13.     c.pyapi.decref(lo_obj)
  14.     c.pyapi.decref(hi_obj)
  15.     c.pyapi.decref(class_obj)
  16.     return res

6.3.3。使用它

nopython 模式函数现在可以使用 Interval 对象以及您在其上定义的各种操作。您可以尝试以下功能：

  1. from numba import jit
  3. @jit(nopython=True)
  4. def inside_interval(interval, x):
  5.     return interval.lo <= x < interval.hi
  7. @jit(nopython=True)
  8. def interval_width(interval):
  9.     return interval.width
  11. @jit(nopython=True)
  12. def sum_intervals(i, j):
  13.     return Interval(i.lo + j.lo, i.hi + j.hi)

6.3.4。结论

我们已经展示了如何执行以下任务：

  • 通过继承Type类来定义新的 Numba 类型类
  • 为非参数类型定义单例 Numba 类型实例
  • 教 Numba 如何使用typeof_impl.register推断某类的 Numba 类型的 Python 值
  • 使用StructModel和register_model定义 Numba 类型的数据模型
  • 使用@box装饰器为 Numba 类型实现装箱功能
  • 使用@unbox装饰器和NativeValue类为 Numba 类型实现拆箱功能
  • 使用@type_callable和@lower_builtin装饰器键入并实现可调用
  • 使用make_attribute_wrapper便捷功能公开只读结构属性
  • 使用@overload_attribute装饰器实现只读属性