多态分发
使用 jit()
或 vectorize()
编译的函数是开放式的:它们可以用许多不同的输入类型调用,并且必须选择(可能即时编译)正确的低级特化。我们在此解释该机制是如何实现的。
要求
JIT 编译的函数可以接受多个参数,并且在选择特化时,每个参数都会被考虑在内。因此,它是一种多重分发形式,比单一分发更复杂。
每个参数根据其 Numba 类型 在选择中发挥作用。Numba 类型通常比 Python 类型更精细:例如,Numba 根据 NumPy 数组的维度和布局(C 连续等)对其进行不同的类型处理。
一旦为每个参数推断出 Numba 类型,就必须从可用的特化中选择一个;或者,如果没有找到合适的特化,则必须编译一个新的。这不是一个简单的决定:可能有多个特化与给定的具体签名兼容(例如,一个双参数函数已经为 (float64, float64)
和 (complex64, complex64)
编译了特化,并且它被以 (float32, float32)
调用)。
因此,分发机制中有两个关键步骤
推断具体参数的 Numba 类型
为推断出的 Numba 类型选择最佳可用特化(或选择编译一个新的)
编译时 vs. 运行时
本文档讨论运行时分发,即当 JIT 编译的函数从纯 Python 调用时。在这种情况下,性能很重要。为了保持在 Python 中正常函数调用开销的范围内,分发的开销应保持在微秒以下。当然,越快越好…
当一个 JIT 编译的函数被另一个 JIT 编译的函数调用时(在 nopython 模式下),多态性在编译时解析,使用非性能关键机制,运行时性能开销为零。
注意
实际上,此处描述的性能关键部分是用 C 语言编写的。
类型解析
因此,第一步是在调用时为函数的每个具体参数推断一个 Numba 类型。鉴于 Numba 类型比 Python 类型具有更细的粒度,无法简单地查找对象的类并用它作为字典的键来获取相应的 Numba 类型。
相反,有一个机制可以检查对象,并根据其 Python 类型,查询各种属性以推断出适当的 Numba 类型。这可能或多或少复杂:例如,一个 Python int
参数总是推断为 Numba intp
(一个指针大小的整数),但一个 Python tuple
参数可以推断为多个 Numba 类型(取决于元组的大小及其每个元素的具体类型)。
Numba 类型系统是高级的,并且用纯 Python 编写;有一个基于泛型函数的纯 Python 机制来执行上述推断(在 numba.typing.typeof
中)。该机制用于编译时推断,例如对常量。不幸的是,它对于运行时基于值的分发来说太慢了。它仅用作不常用(或难以推断)类型的备用方案,并表现出数微秒的开销。
类型码
Numba 类型系统级别太高,无法从 C 代码中高效操作。因此,C 分发层使用基于整数类型码的另一种表示。每个 Numba 类型在构造时都会获得一个唯一的整数类型码;此外,一个内部化系统确保不会创建同一类型的两个实例。因此,分发层能够通过使用简单的整数类型码来避免 Numba 类型系统的开销,这种类型码适用于众所周知的优化(快速哈希表等)。
类型解析步骤的目标变为:为函数的每个具体参数推断一个 Numba 类型码。理想情况下,它不再处理 Numba 类型...
硬编码的快速路径
尽管整数类型码避免了类型系统的抽象和面向对象开销,但它们仍然具有相同的概念复杂性。因此,加速推断的一个重要技术是首先对最重要的类型进行检查,并为每种类型硬编码一个快速解析。
几种类型受益于这种优化,特别是
基本 Python 标量(
bool
,int
,float
,complex
);基本 Numpy 标量(各种整型、浮点型、复数);
具有特定维度和基本元素类型的 Numpy 数组。
理想情况下,每个快速路径在经过一些简单检查后,使用硬编码的结果值或直接的表查找。
然而,我们不能将该技术应用于所有参数类型;否则会出现大量的临时内部缓存,这将变得难以维护。此外,硬编码快速路径的递归应用不一定能组合成低开销(例如,在嵌套元组情况下)。
基于指纹的类型码缓存
对于非平凡类型(例如,想象一个元组,或者一个 Numpy datetime64
数组),硬编码的快速路径不匹配。此时会启用另一种更通用的机制。
这里的原则是检查每个参数值,就像纯 Python 机制会做的那样,并明确地描述其 Numba 类型。区别在于我们实际上不计算 Numba 类型。相反,我们计算一个简单的字节字符串,它是该 Numba 类型的低级可能表示:一个指纹。指纹格式设计得简短且从 C 代码计算起来极其简单(实际上,它具有类似字节码的格式)。
一旦计算出指纹,就会在将指纹映射到类型码的缓存中查找。该缓存是一个哈希表,由于指纹通常非常短(很少超过 20 字节),因此查找速度很快。
如果缓存查找失败,则必须首先使用缓慢的纯 Python 机制计算类型码。幸运的是,这只会发生一次:在后续调用中,将为给定指纹返回缓存的类型码。
在极少数情况下,指纹无法高效计算。对于某些无法从 C 语言轻松检查的类型就是这种情况:例如 cffi
函数指针。此时,每次使用此类参数的函数调用都会调用缓慢的纯 Python 机制。
注意
两个指纹可能表示一个 Numba 类型。这并不会使机制不正确;它只会创建更多的缓存条目。
总结
函数参数的类型解析按顺序涉及以下机制
尝试一些针对常见简单类型的硬编码快速路径。
如果上述失败,则计算参数的指纹并在缓存中查找其类型码。
如果所有上述都失败,则调用纯 Python 机制,该机制将为参数确定一个 Numba 类型(并查找其类型码)。
特化选择
在上一步中,已为 JIT 编译函数的每个具体参数确定了一个整数类型码。现在剩下的是将该具体签名与函数的每个可用特化进行匹配。可能出现三种结果
存在一个令人满意的最佳匹配:然后调用相应的特化(它将处理参数拆箱和其他细节)。
两个或多个“最佳匹配”之间存在平局:抛出异常,拒绝解决歧义。
没有令人满意的匹配:编译一个新的特化,专门针对推断出的具体参数类型。
选择通过遍历所有可用特化来工作,并计算每个具体参数类型与特化预期签名中相应类型的兼容性。具体来说,我们关注的是
具体参数类型是否允许隐式转换为特化的参数类型;
如果允许,这种转换会带来什么样的语义(用户可见)成本。
隐式转换规则
从源类型到目标类型可能存在五种隐式转换(请注意这是一种不对称关系)
精确匹配:两种类型相同;这是理想情况,因为特化将完全按照预期运行;
同类提升:两种类型属于同一“种类”(例如
int32
和int64
都是整数类型),并且源类型可以无损转换为目标类型(例如从int32
到int64
,但反之则不行);安全转换:两种类型属于不同种类,但源类型可以合理地转换为目标类型(例如从
int32
到float64
,但反之则不行);不安全转换:存在从源类型到目标类型的转换,但它可能会丢失精度、数量级或其他所需质量。
无转换:两种类型之间没有正确或合理高效的转换方式(例如
int64
和datetime64
之间,或 C 连续数组和 Fortran 连续数组之间)。
当检查一个特化时,后两种情况会将其从最终选择中排除:即当至少一个参数没有转换或只有不安全转换到签名的参数类型时。
注意
然而,如果函数在 jit()
调用中用显式签名编译(因此不允许编译新的特化),则允许不安全转换。
候选者与最佳匹配
如果一个特化没有被上述规则排除,它将进入最终选择的候选者列表。这些候选者按一个有序的四元组整数进行排名:(不安全转换的数量, 安全转换的数量, 同类提升的数量, 精确匹配的数量)
(请注意,元组元素的总和等于参数的数量)。最佳匹配是按升序排序的第 1 个结果,从而优先选择精确匹配而非提升,优先选择提升而非安全转换,优先选择安全转换而非不安全转换。
实现
上述机制作用于整数类型码,而非 Numba 类型。它使用一个内部哈希表,存储每对兼容类型的可能转换类型。该内部哈希表部分在启动时构建(用于内置的简单类型,如 int32
、int64
等),部分动态填充(用于任意复杂类型,如数组类型:例如,允许在函数预期非连续二维数组时使用 C 连续二维数组)。
杂项
一些分发性能基准测试存在于 Numba 基准测试仓库中。
该机制特定方面的一些单元测试可在 numba.tests.test_typeinfer
和 numba.tests.test_typeof
中找到。更高层次的分发测试在 numba.tests.test_dispatcher
中。