多态分发

使用 jit()vectorize() 编译的函数是开放式的:它们可以用许多不同的输入类型调用,并且必须选择(可能即时编译)正确的低级特化。我们在此解释该机制是如何实现的。

要求

JIT 编译的函数可以接受多个参数,并且在选择特化时,每个参数都会被考虑在内。因此,它是一种多重分发形式,比单一分发更复杂。

每个参数根据其 Numba 类型 在选择中发挥作用。Numba 类型通常比 Python 类型更精细:例如,Numba 根据 NumPy 数组的维度和布局(C 连续等)对其进行不同的类型处理。

一旦为每个参数推断出 Numba 类型,就必须从可用的特化中选择一个;或者,如果没有找到合适的特化,则必须编译一个新的。这不是一个简单的决定:可能有多个特化与给定的具体签名兼容(例如,一个双参数函数已经为 (float64, float64)(complex64, complex64) 编译了特化,并且它被以 (float32, float32) 调用)。

因此,分发机制中有两个关键步骤

  1. 推断具体参数的 Numba 类型

  2. 为推断出的 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 编译函数的每个具体参数确定了一个整数类型码。现在剩下的是将该具体签名与函数的每个可用特化进行匹配。可能出现三种结果

  • 存在一个令人满意的最佳匹配:然后调用相应的特化(它将处理参数拆箱和其他细节)。

  • 两个或多个“最佳匹配”之间存在平局:抛出异常,拒绝解决歧义。

  • 没有令人满意的匹配:编译一个新的特化,专门针对推断出的具体参数类型。

选择通过遍历所有可用特化来工作,并计算每个具体参数类型与特化预期签名中相应类型的兼容性。具体来说,我们关注的是

  1. 具体参数类型是否允许隐式转换为特化的参数类型;

  2. 如果允许,这种转换会带来什么样的语义(用户可见)成本。

隐式转换规则

从源类型到目标类型可能存在五种隐式转换(请注意这是一种不对称关系)

  1. 精确匹配:两种类型相同;这是理想情况,因为特化将完全按照预期运行;

  2. 同类提升:两种类型属于同一“种类”(例如 int32int64 都是整数类型),并且源类型可以无损转换为目标类型(例如从 int32int64,但反之则不行);

  3. 安全转换:两种类型属于不同种类,但源类型可以合理地转换为目标类型(例如从 int32float64,但反之则不行);

  4. 不安全转换:存在从源类型到目标类型的转换,但它可能会丢失精度、数量级或其他所需质量。

  5. 无转换:两种类型之间没有正确或合理高效的转换方式(例如 int64datetime64 之间,或 C 连续数组和 Fortran 连续数组之间)。

当检查一个特化时,后两种情况会将其从最终选择中排除:即当至少一个参数没有转换或只有不安全转换到签名的参数类型时。

注意

然而,如果函数在 jit() 调用中用显式签名编译(因此不允许编译新的特化),则允许不安全转换

候选者与最佳匹配

如果一个特化没有被上述规则排除,它将进入最终选择的候选者列表。这些候选者按一个有序的四元组整数进行排名:(不安全转换的数量, 安全转换的数量, 同类提升的数量, 精确匹配的数量)(请注意,元组元素的总和等于参数的数量)。最佳匹配是按升序排序的第 1 个结果,从而优先选择精确匹配而非提升,优先选择提升而非安全转换,优先选择安全转换而非不安全转换。

实现

上述机制作用于整数类型码,而非 Numba 类型。它使用一个内部哈希表,存储每对兼容类型的可能转换类型。该内部哈希表部分在启动时构建(用于内置的简单类型,如 int32int64 等),部分动态填充(用于任意复杂类型,如数组类型:例如,允许在函数预期非连续二维数组时使用 C 连续二维数组)。

总结

选择正确的特化涉及以下步骤

  • 检查每个可用特化,并将其与具体参数类型进行匹配。

  • 排除任何至少一个参数不提供足够兼容性的特化。

  • 如果还有剩余候选者,选择在保留类型语义方面最佳的一个。

杂项

一些分发性能基准测试存在于 Numba 基准测试仓库中。

该机制特定方面的一些单元测试可在 numba.tests.test_typeinfernumba.tests.test_typeof 中找到。更高层次的分发测试在 numba.tests.test_dispatcher 中。