Skip to content

6.3。示例:间隔类型

原文: 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

6.3.1。扩展打字层

6.3.1.1。创建一个新的 Numba 类型

由于 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()

6.3.1.2。 Python 值的类型推断

就其本身而言,创建一个 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

6.3.1.3。操作类型推断

我们希望能够从 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 相同。这允许正确处理关键字参数。

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

6.3.2。延伸降低层

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

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

作为一般规则, 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类型的值(或其任何实例)表示为两个字段lohi的结构,每个字段都是双精度浮点数(types.float64)。

注意

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

6.3.2.2。公开数据模型属性

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

from numba.extending import make_attribute_wrapper

make_attribute_wrapper(IntervalType, 'lo', 'lo')
make_attribute_wrapper(IntervalType, 'hi', 'hi')

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

6.3.2.3。公开财产

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

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 中结合了类型推断和代码生成。

6.3.2.4。实现构造函数

现在我们要实现两个参数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.float64types.Float的实例)。在类上而不是在特定实例上匹配参数类型通常更具有前瞻性(但是,当 _ 返回 _ 类型时 - 主要是在类型推断阶段 - ,您通常必须返回一个类型实例)。

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

6.3.2.5。拳击和拆箱

如果此时尝试使用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

6.3.3。使用它

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)

6.3.4。结论

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

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


回到顶部