原文: http://numba.pydata.org/numba-doc/latest/extending/interval-example.html
我们将扩展 Numba 前端以支持它目前不支持的类,以便允许:
- 将类的实例传递给 Numba 函数
- 在 Numba 函数中访问类的属性
- 从 Numba 函数构造并返回类的新实例
(以上所有 nopython 模式)
我们将混合来自高级扩展 API 和低级扩展 API 的 API,具体取决于给定任务的可用内容。
我们的示例的起点是以下纯 Python 类:
class Interval(object):
"""
A half-open interval on the real number line.
"""
def __init__(self, lo, hi):
self.lo = lo
self.hi = hi
def __repr__(self):
return 'Interval(%f, %f)' % (self.lo, self.hi)
@property
def width(self):
return self.hi - self.lo
由于 Numba 不知道Interval
类,我们必须创建一个新的 Numba 类型来表示它的实例。 Numba 不直接处理 Python 类型:它有自己的类型系统,允许不同级别的粒度以及常规 Python 类型不具备的各种元信息。
我们首先创建一个类型类IntervalType
,因为我们不需要参数类型,所以我们实例化一个类型实例interval_type
:
from numba import types
class IntervalType(types.Type):
def __init__(self):
super(IntervalType, self).__init__(name='Interval')
interval_type = IntervalType()
就其本身而言,创建一个 Numba 类型并没有做任何事情。我们必须教 Numba 如何推断一些 Python 值作为该类型的实例。在这个例子中,它是微不足道的:Interval
类的任何实例都应被视为属于interval_type
类型:
from numba.extending import typeof_impl
@typeof_impl.register(Interval)
def typeof_index(val, c):
return interval_type
因此,只要它们是Interval
的实例,函数参数和全局值就会被认为属于interval_type
。
我们希望能够从 Numba 函数构造区间对象,因此我们必须教 Numba 识别双参数Interval(lo, hi)
构造函数。参数应该是浮点数:
from numba.extending import type_callable
@type_callable(Interval)
def type_interval(context):
def typer(lo, hi):
if isinstance(lo, types.Float) and isinstance(hi, types.Float):
return interval_type
return typer
type_callable()
装饰器指定在为给定的可调用对象(此处为Interval
类本身)运行类型推断时应调用修饰函数。修饰函数必须简单地返回一个将使用参数类型调用的 typer 函数。这个看似错综复杂的设置的原因是,typer 函数使 _ 与 _ 完全相同的签名与打字的 callable 相同。这允许正确处理关键字参数。
由装饰函数接收的 _ 上下文 _ 参数在更复杂的情况下非常有用,在这种情况下,计算可调用的返回类型需要解析其他类型。
我们已经完成了关于我们的类型推理添加的 Numba 教学。我们现在必须教 Numba 如何为新操作实际生成代码和数据。
作为一般规则, nopython 模式不适用于 Python 对象,因为它们是由 CPython 解释器生成的。解释器使用的表示对于快速本机代码而言效率太低。因此, nopython 模式中支持的每种类型都必须定义定制的本机表示,也称为 _ 数据模型 _。
数据模型的一个常见情况是不可变的类似结构的数据模型,类似于 C struct
。我们的 interval 数据类型很方便地属于该类别,这里有一个可能的数据模型:
from numba.extending import models, register_model
@register_model(IntervalType)
class IntervalModel(models.StructModel):
def __init__(self, dmm, fe_type):
members = [
('lo', types.float64),
('hi', types.float64),
]
models.StructModel.__init__(self, dmm, fe_type, members)
这指示 Numba 将IntervalType
类型的值(或其任何实例)表示为两个字段lo
和hi
的结构,每个字段都是双精度浮点数(types.float64
)。
注意
可变类型需要更复杂的数据模型才能在修改后保留其值。它们通常不能存储和传递到堆栈或寄存器,如不可变类型。
我们希望数据模型属性lo
和hi
以相同的名称公开,以便在 Numba 函数中使用。 Numba 提供了一个方便的功能来做到这一点:
from numba.extending import make_attribute_wrapper
make_attribute_wrapper(IntervalType, 'lo', 'lo')
make_attribute_wrapper(IntervalType, 'hi', 'hi')
这将以只读模式公开属性。如上所述,可写属性不适合此模型。
由于width
属性是计算而不是存储在结构中,我们不能像我们对lo
和hi
那样简单地公开它。我们必须明确地重新实现它:
from numba.extending import overload_attribute
@overload_attribute(IntervalType, "width")
def get_width(interval):
def getter(interval):
return interval.hi - interval.lo
return getter
您可能会问为什么我们不需要为此属性公开类型推断钩子?答案是@overload_attribute
是高级 API 的一部分:它在单个 API 中结合了类型推断和代码生成。
现在我们要实现两个参数Interval
构造函数:
from numba.extending import lower_builtin
from numba import cgutils
@lower_builtin(Interval, types.Float, types.Float)
def impl_interval(context, builder, sig, args):
typ = sig.return_type
lo, hi = args
interval = cgutils.create_struct_proxy(typ)(context, builder)
interval.lo = lo
interval.hi = hi
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 值。
如果此时尝试使用Interval
实例,则肯定会得到错误 “无法将 Interval 转换为本机值”。这是因为 Numba 还不知道如何从 Python Interval
实例创建本机间隔值。让我们教它如何做到:
from numba.extending import unbox, NativeValue
@unbox(IntervalType)
def unbox_interval(typ, obj, c):
"""
Convert a Interval object to a native interval structure.
"""
lo_obj = c.pyapi.object_getattr_string(obj, "lo")
hi_obj = c.pyapi.object_getattr_string(obj, "hi")
interval = cgutils.create_struct_proxy(typ)(c.context, c.builder)
interval.lo = c.pyapi.float_as_double(lo_obj)
interval.hi = c.pyapi.float_as_double(hi_obj)
c.pyapi.decref(lo_obj)
c.pyapi.decref(hi_obj)
is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
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 函数返回间隔值:
from numba.extending import box
@box(IntervalType)
def box_interval(typ, val, c):
"""
Convert a native interval structure to an Interval object.
"""
interval = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val)
lo_obj = c.pyapi.float_from_double(interval.lo)
hi_obj = c.pyapi.float_from_double(interval.hi)
class_obj = c.pyapi.unserialize(c.pyapi.serialize_object(Interval))
res = c.pyapi.call_function_objargs(class_obj, (lo_obj, hi_obj))
c.pyapi.decref(lo_obj)
c.pyapi.decref(hi_obj)
c.pyapi.decref(class_obj)
return res
nopython 模式函数现在可以使用 Interval 对象以及您在其上定义的各种操作。您可以尝试以下功能:
from numba import jit
@jit(nopython=True)
def inside_interval(interval, x):
return interval.lo <= x < interval.hi
@jit(nopython=True)
def interval_width(interval):
return interval.width
@jit(nopython=True)
def sum_intervals(i, j):
return Interval(i.lo + j.lo, i.hi + j.hi)
我们已经展示了如何执行以下任务:
- 通过继承
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
装饰器实现只读属性