ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

NBEP 3:Numba JIT Classes(jitclass)设计提案与实现深度解析

NBEP 3:Numba JIT Classes(jitclass)设计提案与实现深度解析 NBEP 3Numba JIT Classesjitclass设计提案与实现深度解析【免费下载链接】numbaNumPy aware dynamic Python compiler using LLVM项目地址: https://gitcode.com/gh_mirrors/nu/numbaNumba 的 JIT Classes下文统一称jitclass是让用户自定义的 Python 类获得 nopython 模式编译加速的机制开发者通过jitclass(spec)声明每个字段的 Numba 类型类中的方法即被编译为原生机器码实例数据以 C 兼容结构存放在堆上可供编译函数直接访问。本文以仓库中的设计提案 docs/source/proposals/jit-classes.rstNBEP 3为主体脉络结合当前仓库中 jitclass 的实际实现 numba/experimental/jitclass、用户指南 docs/source/user/jitclass.rst 与测试用例 numba/tests/test_jitclasses.py逐层拆解其存储模型、类型推断、引用计数、装箱box/unbox等核心机制并给出可直接运行的实战示例。读完本文你将掌握 jitclass 的设计边界、spec 的完整写法、从解释器到编译代码的调用链以及它的全部已知限制。提案背景为什么 Numba 需要用户自定义类NBEP 3作者 Siu Kwan Lam2015 年 12 月状态 Draft开篇指出彼时 Numba 尚不支持用户自定义类。类提供数据属性与操作方法的抽象封装合理使用能提升模块化程度一个类的实例是类的实例化产物。该提案聚焦最朴素的使用场景——仅包含属性和方法的类类方法classmethod、静态方法staticmethod与继承被推迟到后续提案但提案认为基于文中描述的底层基础这些特性可以很容易地补上。从当前仓库看这一设想已基本落地jitclass现在位于 numba/experimental/init.py 并对外导出同时支持 property仅 getter/setter、staticmethod但仍不支持继承详见下文支持边界。提案是理解设计动机的第一手资料而 docs/source/user/jitclass.rst 则是其当代使用形态的权威说明。核心设计决策类对象不是一等公民提案为 jitclass 划定了比 Python 类更严格的操作集合围绕类及其实例仅支持四类操作实例化以类对象为构造函数调用cls(*args, **kwargs)析构释放实例化期间分配的资源并解除对所有其他对象的引用属性访问通过instance.attr语法加载和存储属性方法访问通过instance.method语法加载方法。关键设计决策在于类对象而非实例不需要被物化materialize。把类对象用作构造函数时具体运行时实现会在编译器的类型推断typing阶段被完全解析出来这意味着类对象在 Numba 中不是一等公民first class——若实现一等公民的类对象将需要一种接口类型或类的类型。同理方法从不存储于实例中而是挂在类上既然类对象只存在于类型域方法也会在 typing 阶段被完全解析。这个决策在源码中留下了清晰的证据在 numba/experimental/jitclass/base.py 的ctor_impl中构造函数直接在代码生成阶段完成分配内存 → 置零 → 调用 JIT 编译的__init__整条流程整个编译期根本不需要存在一个运行时类对象。对编译内层的影响类对象被当作构造函数函数由于类对象不是一等公民在 Numba 编译函数内部jitclass 的类对象被当作一个函数即构造函数来使用。这一点在 docs/source/user/jitclass.rst 的 Limitations 中被明确列为第一条A jitclass class object is treated as a function (the constructor) inside a Numba compiled function.也就是说你可以在njit函数里写Vec(1, 2)来构造实例但无法把Vec本身当作值传来传去。存储模型C 兼容的 plain-old-data 结构提案规定为了与 C 兼容属性被存放在一个简单的 plain-old-dataPOD结构中。每个属性按用户定义的顺序存放在按需填充padding 以保证正确对齐、连续的内存区域中。提案给出了一个包含int32、float32、complex64三个字段的实例可与如下 C 结构完全兼容struct { int32 field0; float32 field1; complex64 field2; };该结构同时也与对齐后的 NumPy 结构化 dtypestructured dtype兼容。在当代实现中这个模型由 numba/experimental/jitclass/base.py 中的两个数据模型类体现InstanceModel描述实例本身成员为meminfo指向 NRT MemInfo 的指针用于引用计数与data指向实际属性数据块的CPointerInstanceDataModel描述属性数据块成员由clsty.struct即 spec 定义的字段名与类型映射而来字段名经_mangle_attr处理为m_ name形式。属性访问在 get_attr_impl 与 set_attr_impl 中实现读取时从inst.data拿到数据指针用make_data_helper构造出结构体再取对应字段写入时先对旧值decref、对新值incref详见下一节。整个数据块就是一块连续内存因此编译函数内的属性访问只是对 C 结构的一次读取开销极低短方法还可以由 LLVM 内联器决定内联这是 jitclass 在编译函数内使用更高效的根本原因。方法与特殊方法__init__的定位与__del__的缺席提案指出方法就是普通的可绑定到实例上的函数可以像普通函数一样被 Numba 编译运行时通过getattr(instance, name)将实例绑定到对应方法。特殊的__init__方法也按普通函数处理而__del__在提案阶段不受支持参考 docs/source/proposals/jit-classes.rst 中 Methods 一节且时至今日依旧不支持用户自定义__del__——实例的清理统一由 NRT 的析构器完成见下一节。在源码层面构造流程是__new__分配 __init__初始化的编译期组合register_class_type 把类中所有FunctionType成员用njit编译为jit_methodsproperty 的fget/fset分别编译为jit_propsstaticmethod编译为jit_static_methodsConstructorTemplatebase.py把构造调用转发给__init__的签名解析并要求__init__()必须返回None否则抛出NumbaTypeErrorctor_impl真正执行meminfo_alloc_dtor分配带析构器的内存 → 将数据块置零 → 以(instance, *args)调用 JIT 编译的__init__。测试 numba/tests/test_jitclasses.pytest_jitclass_usage_from_python验证了从解释器构造、访问属性/方法、设置属性、MemInfo 引用计数变化等完整行为是理解上述调用链最直接的运行证据。引用计数与析构器NRT 生命周期管理实例可能包含其他 NRT 跟踪对象作为属性。提案规定jitclass 实例由 NRT 进行引用计数当引用计数降为零时必须调用析构器析构器会将所有属性的引用计数各减一。提案同时明确两个已知限制不支持用户自定义__del__循环引用的正确清理当时未处理环会导致内存泄漏这一限制至今仍然成立需要开发者自己避免属性间的循环引用。实现中析构器由 imp_dtor 生成它为一个命名如_Dtor.type的 LLVM 函数接收(voidptr, size, voidptr)三个参数把数据指针 bitcast 回属性结构体类型后对整个结构体执行一次nrt.decref结构体模型的 decref 会递归到每个 NRT 字段。而在set_attr_impl中写入新值时先incref新值、再decref旧值从而保证替换属性时引用计数不泄漏。类型推断spec 的两种写法提案指出要物化实例例如分配存储必须有属性的类型信息最简单的方式是由用户提供每个属性的类型及其顺序例如用OrderedDictdct OrderedDict() dct[x] int32 dct[y] float32不过这种静态类型语义不如 Python 那种通用类灵活。提案还回顾了一个早期的实现尝试通过特化__init__来捕获存入属性的类型——但方法可以包含任意逻辑若类型按值条件性赋值问题会退化为依赖类型dependent typing问题很少有语言实现依赖类型且多限于定理证明器因此被否决最终采用用户显式提供 spec 的方案。写法一OrderedDictfrom collections import OrderedDict import numba spec OrderedDict() spec[x] numba.int32 spec[y] numba.float32 jitclass(spec) class Vec(object): def __init__(self, x, y): self.x x self.y y def add(self, dx, dy): self.x dx self.y dy使用普通字典也可以但为了字段顺序稳定提案与 decorators.py 的文档字符串都推荐使用OrderedDict。写法二2-tuple 列表spec [(x, numba.int32), (y, numba.float32)] jitclass(spec) class Vec(object): ...两种写法在当前实现中统一处理register_class_type 会把Sequence形式的 spec 转成OrderedDict再经_validate_spec校验——键必须是字符串、值必须是 Numba 类型实例随后通过_fix_up_private_attr对__x这类私有名做 CPython 式的名字重整mangling并检查方法/属性名与字段名之间是否存在名字遮蔽name shadowing有则抛出NameError。当代扩展用类型注解推断字段类型提案之后实现演进出了一个更便捷的玩法字段类型也可以从类的类型注解中推断。在 numba/tests/doc_examples/test_jitclass.py 中有完整示例其要点也记录在 docs/source/user/jitclass.rst 中如下from numba.experimental import jitclass jitclass class Counter: value: int def __init__(self): self.value 0 def get(self) - int: ret self.value self.value 1 return ret规则是register_class_type通过pt.get_type_hints(cls)收集类级注解凡是不在 spec 中的字段就用as_numba_type(py_type)推断出对应的 Numba 类型补入 spec。注意两点边界只有类级注解参与推断方法含__init__的参数注解会被忽略NumPy 数组的 dtype 与维度无法用注解表达因此数组字段必须显式写进 spec例如(y, float64[:])。另外numba.typed.Dict、numba.typed.List等容器作为成员时可以显式以 Numba 类型如types.DictType(types.int64, types.unicode_type)写进 spec也可以用numba.typeof(实例)或容器的_numba_type_属性取得类型后写入 spec但容器成员必须先初始化再使用否则会触发非法内存访问segfaultdocs/source/user/jitclass.rst 中专门给出了反面示例。同源多型同一个类对象可创建多个 jitclass 类型提案特别强调即使作用于同一个类对象、同一份 specjitclass(spec)每次调用都会创建一个新的 jitclass 类型class Vec(object): ... Vec1 jitclass(spec)(Vec) Vec2 jitclass(spec)(Vec) # Vec1 和 Vec2 是两个不同的 jitclass 类型这一点在 decorators.py 中有对应的全部五种调用形态jitclass()—— 无参装饰jitclass(specspec)—— 关键字传 specjitclass—— 裸装饰类上无注解时等价于空 specjitclass(spec)—— 位置参数传 spec此时内部会做cls_or_spec, spec None, cls_or_spec的参数规整JitFoo jitclass(Foo, spec)—— 函数式调用返回编译后的类。其中第 3、4 种形态在装饰器源码注释中被逐一列出并标号是理解jitclass参数规整逻辑decorators.py的直接依据。从解释器使用box/unbox 与装箱开销提案描述了 jitclass 在解释器中的行为构造实例时会创建一个box包装底层的 jitclass 实例属性和方法都可从解释器访问实际实现则运行在 Numba 编译代码中。任何 Python 对象在传入 Numba 侧时都会被转换为原生表示返回值同理被转回 Python 表示因此在解释器中操作 jitclass 实例存在一定的装箱开销——但提案认为这个开销很小且很容易被编译方法内更高效的计算所摊薄。当前实现中这套机制位于 numba/experimental/jitclass/boxing.py_specialize_boxboxing.py为每个 jitclass 类型动态生成一个_box.Box的子类字段被注入为propertygetter/setter 都是njit编译的小函数并缓存结果避免代码膨胀方法被注入为_generate_method生成的 JIT 包装器_box_class_instance/_unbox_class_instanceboxing.py实现解释器与原生表示之间的转换通过_boxC 扩展中的字节偏移读写meminfo与data指针字段值只有在解释器中真正访问属性时才会被逐个 box 成 Python 对象——jitclass 实例本身被传给解释器时并不会把内部值整体装箱这也是 docs/source/user/jitclass.rst 中性能小节描述的优化点。typeof对 box 类型也有专门实现boxing.py用于短路类型判定避免实现了自定义__hash__的 jitclass 在类型推断时触发无限递归源码注释详细记录了这一坑。支持边界property、静态方法、继承与目标平台提案与当前实现共同划定了如下边界特性支持情况依据property仅支持 getter 与 setterdeleter 不支持register_class_type中对含fdel的 property 抛出TypeErrorbase.pystaticmethod提案时代不支持当前已支持可从实例属性mybag.add(1, 1)或类属性Bag.add(1, 2)调用但类内方法中不可调用Bag.add()docs/source/user/jitclass.rstclassmethod不支持提案 Support for property, staticmethod and classmethod 一节继承仅允许基类为object对 jitclass 再子类化会抛出TypeErrorcannot subclass from a jitclassbase.py目标平台仅 CPU 目标含 parallel 目标提案称 GPUCUDA、HSA通过实例的不可变版本支持详见另一份 NBEP提案 Supported targets 一节注意上表中的staticmethod 当前已支持是与提案原文的差异点提案写作于 2015 年声明 staticmethod 不支持当前仓库实现base.py、boxing.py已经落地支持且 test_jitclass.py 中的Bag.add静态方法示例对此有测试覆盖。因此引用该提案时应注意其历史草稿属性Status: Draft以 docs/source/user/jitclass.rst 为现行行为依据。此外JitClassTypebase.py作为所有 jitclass 的元类要求有且仅有一个基类类中不能出现除方法、property、staticmethod 之外的其他成员如普通类变量否则抛出TypeError(class members are not yet supported: ...)。其他属性与运行语义给定spec [(x, numba.int32), (y, numba.float32)] jitclass(spec) class Vec(object): ...提案明确了两个容易踩坑的语义isinstance(Vec(1, 2), Vec)为True由JitClassType.__instancecheck__实现见 base.pytype(Vec(1, 2))可能不等于Vec——因为解释器拿到的是 box 包装对象其动态类型是_specialize_box生成的特化子类。同时要注意isinstance()仅在解释器中可用在编译函数内isinstance不受支持。这也解释了为什么解释器中对实例的修改是可行的_box.Box子类上的字段 property 会触发 JIT getter/setter间接读写原生数据块与njit函数中看到的实例共享同一块 NRT 内存test_jitclasses.py 通过_get_meminfo断言了obj与传入njit函数后的副本共享同一个 MemInfo 与数据指针。支持的操作与可用的 dunder 方法当前 jitclass 在解释器与 Numba 编译函数中均支持的操作docs/source/user/jitclass.rst调用类对象构造实例mybag Bag(123)读写属性与 propertymybag.value调用方法mybag.increment(3)以实例属性形式调用静态方法mybag.add(1, 1)以类属性形式调用静态方法Bag.add(1, 2)仅限类外代码使用选定的一组 dunder 方法如__add__支持mybag otherbag。可定义的 dunder 方法覆盖面很广包括算术类__add__、__sub__、__mul__、__truediv__、__floordiv__、__mod__、__pow__、__matmul__、__lshift__、__rshift__及对应的反射形式__radd__等与原位形式__iadd__等、比较类__eq__、__ne__、__lt__、__le__、__gt__、__ge__、位运算类__and__、__or__、__xor__及其反射/原位形式、以及__abs__、__bool__、__int__、__float__、__complex__、__len__、__contains__、__getitem__、__setitem__、__str__、__hash__、__index__、__invert__、__neg__、__pos__等。这套白名单在 boxing.py 中与装箱逻辑保持一致凡是不在白名单内的__xxx__方法在装箱时都会抛出TypeError测试 test_jitclasses.pytest_jitclass_unsupported_dunder对此有专门验证。在实现层面__getitem__/__setitem__等会被特殊处理base.py除注册__getitem__外还会把对应的operator.getitem也降低lower到同一实现从而同时捕获来自 Python 语法与 Numba 代码中的[]操作。实战示例一个完整的 jitclass综合 docs/source/user/jitclass.rst 与 test_jitclass.py 的官方示例一个兼具字段、property、静态方法与数组字段的完整 jitclass 如下import numpy as np from numba import int32, float32 # 导入 Numba 类型 from numba.experimental import jitclass spec [ (value, int32), # 标量字段 (array, float32[:]), # 数组字段 ] jitclass(spec) class Bag(object): def __init__(self, value): self.value value self.array np.zeros(value, dtypenp.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) assert mybag.value n assert mybag.size n np.testing.assert_allclose(mybag.increment(3), 3 * np.ones(n, dtypenp.float32)) np.testing.assert_allclose(mybag.add(1, 1), 2) np.testing.assert_allclose(Bag.add(1, 2), 3)在 Numba 编译函数中使用同样顺手from numba import njit njit def use_bag(b): b.increment(5) # 方法调用 total 0.0 for v in b.array: # 属性读取 total v return total性能语义与已知限制小结性能上的关键事实均有文档或源码依据编译函数内使用 jitclass 更高效属性访问只是读取 C 结构get_attr_impl短方法可被 LLVM 内联解释器中使用与调用任何 Numba 编译函数有相同的装箱/拆箱开销且解释器中的 jitclass 操作尚未优化属性值仅在访问时逐个装箱提案强调的装箱开销最小且易于摊薄的前提是把计算密集逻辑写进编译方法解释器只做薄薄一层的调度。已知限制汇总jitclass 类对象在编译函数内被当作构造函数函数不是一等公民isinstance()只在解释器可用type()返回值可能不是原始类不支持继承仅object基类、不支持classmethod、不支持 property deleter、不支持用户自定义__del__循环引用会导致内存泄漏需要开发者自行规避仅 CPU 目标含 parallel可用未初始化的字段含有垃圾数据容器类成员必须先初始化再使用。结语从 NBEP 3 到numba.experimental.jitclassNBEP 3 是 jitclass 的奠基性设计文档它确立了类对象非一等公民POD 存储模型spec 静态类型声明NRT 引用计数与自动析构解释器侧 box 包装五大支柱。十余年后这些设计在 numba/experimental/jitclass 中依然清晰可辨并在此之上演化出类型注解推断、staticmethod 支持、dunder 白名单与 typed 容器成员等增强。如果你要在自己的数值代码中引入自定义类并享受 nopython 模式编译docs/source/user/jitclass.rst 是现行行为的权威参考numba/tests/test_jitclasses.py 与 numba/tests/doc_examples/test_jitclass.py 则是可以直接运行的最完整范例集。【免费下载链接】numbaNumPy aware dynamic Python compiler using LLVM项目地址: https://gitcode.com/gh_mirrors/nu/numba创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表