JEP 438: Vector API (Fifth Incubator) | 向量 API(第五次孵化)
摘要
引入一个 API 来表达向量计算,该 API 在运行时可靠地编译为受支持的 CPU 架构上的最佳向量指令,从而实现优于等效标量计算的性能。
历史
向量 API 最初由 JEP 338 提出,并作为 孵化 API 集成到 JDK 16 中。第二轮孵化由 JEP 414 提出,并集成到 JDK 17 中。第三轮孵化由 JEP 417 提出,并集成到 JDK 18 中。第四轮孵化由 JEP 426 提出,并集成到 JDK 19 中。
这个 JEP 提议在 JDK 20 中重新孵化该 API,相对于 JDK 19,API 没有变化。该实现包括少量的错误修复和性能增强。这个 JEP 还明确指出,与 Project Valhalla 保持一致是完成向量 API 的关键部分。
目标
- 清晰简洁的 API——该 API 应该能够清晰简洁地表达广泛的向量计算,包括在循环内以及可能带有控制流的由向量操作序列组成的计算。应该可以表达对向量大小(即每个向量的通道数)通用的计算,从而使这样的计算能够在支持不同向量大小的硬件上可移植。
- 平台无关性——该 API 应该与 CPU 架构无关,能够在支持向量指令的多个架构上实现。就像在 Java API 中通常的情况一样,如果平台优化和可移植性发生冲突,那么我们将倾向于使 API 可移植,即使这会导致一些特定于平台的习惯用法在可移植代码中无法表达。
- 在 x64 和 AArch64 架构上可靠的运行时编译和性能——在有能力的 x64 架构上,Java 运行时,特别是 HotSpot C2 编译器,应该将向量操作编译为相应的高效且高性能的向量指令,例如 流式 SIMD 扩展(SSE)和 高级向量扩展(AVX)支持的指令。开发人员应该相信他们表达的向量操作将可靠地紧密映射到相关的向量指令。在有能力的 ARM AArch64 架构上,C2 同样会将向量操作编译为 NEON 和 SVE 支持的向量指令。
- 优雅降级——有时向量计算在运行时不能完全表示为向量指令序列,可能是因为架构不支持某些所需的指令。在这种情况下,向量 API 实现应该优雅地降级并且仍然能够正常工作。如果向量计算不能有效地编译为向量指令,这可能会发出警告。在没有向量的平台上,优雅降级将生成与手动展开循环具有竞争力的代码,其中展开因子是所选向量中的通道数。
- 与 Project Valhalla 保持一致——向量 API 的长期目标是利用 Project Valhalla 对 Java 对象模型的增强。主要这将意味着将向量 API 当前的 基于值的类 更改为值类,以便程序可以使用值对象,即没有对象标识的类实例。因此,向量 API 将在多个版本中进行孵化,直到 Project Valhalla 的必要功能作为 预览功能 可用。一旦这些 Valhalla 功能可用,我们将调整向量 API 和实现以使用它们,然后将向量 API 本身提升为预览功能。有关更多详细信息,请参阅 运行时编译 和 未来工作 部分。
非目标
- 增强 HotSpot 中现有的自动向量化算法不是目标。
- 支持除 x64 和 AArch64 以外的 CPU 架构上的向量指令不是目标。然而,重要的是要像目标中所表达的那样,该 API 绝不能排除这样的实现。
- 支持 C1 编译器不是目标。
- 保证支持严格的浮点计算不是目标,就像 Java 平台对标量操作所要求的那样。在浮点标量上执行的浮点操作的结果可能与在浮点标量向量上执行的等效浮点操作的结果不同。任何偏差都将被清楚地记录下来。这个非目标并不排除表达或控制浮点向量计算所需精度或可重复性的选项。
动机
向量计算由向量上的一系列操作组成。向量由(通常)固定序列的标量值组成,其中标量值对应于硬件定义的向量通道数。应用于具有相同通道数的两个向量的二元操作将对每个通道在每个向量的相应两个标量值上应用等效的标量操作。这通常被称为 单指令多数据(SIMD)。
向量操作表达了一定程度的并行性,使得在单个 CPU 周期内可以执行更多的工作,从而可以带来显著的性能提升。例如,给定两个向量,每个向量包含八个整数的序列(即八个通道),这两个向量可以使用单个硬件指令相加。向量加法指令在通常需要对两个整数进行操作(执行一次整数加法)的时间内对十六个整数进行操作,执行八次整数加法。
HotSpot 已经支持 自动向量化,它将标量操作转换为超字操作,然后将其映射到向量指令。可转换的标量操作集是有限的,并且对于代码形状的变化也很脆弱。此外,可能只使用了可用向量指令的一个子集,限制了生成代码的性能。
如今,希望编写可靠地转换为超字操作的标量操作的开发人员需要了解 HotSpot 的自动向量化算法及其局限性,以实现可靠且可持续的性能。在某些情况下,可能无法编写可转换的标量操作。例如,HotSpot 不会转换计算数组哈希码的简单标量操作(Arrays::hashCode 方法),也不能自动向量化按字典顺序比较两个数组的代码(因此我们 添加了一个用于字典顺序比较的内在函数)。
向量 API 的目的是通过提供一种在 Java 中编写复杂向量算法的方法来改善这种情况,使用现有的 HotSpot 自动向量化器,但具有使向量化更加可预测和稳健的用户模型。手写的向量循环可以表达高性能算法,例如向量化的 hashCode
或专门的数组比较,自动向量化器可能永远不会对其进行优化。许多领域可以从这个显式向量 API 中受益,包括机器学习、线性代数、密码学、金融以及 JDK 本身中的代码。
描述
向量由抽象类 Vector<E>
表示。类型变量 E
被实例化为由向量覆盖的标量基本整数或浮点元素类型的装箱类型。向量还有一个 形状,它定义了向量的大小(以位为单位)。向量的形状决定了当 HotSpot C2 编译器编译向量计算时 Vector<E>
的实例如何映射到硬件向量寄存器。向量的长度,即通道数或元素数,是向量大小除以元素大小。
支持的元素类型(E
)集是 Byte
、Short
、Integer
、Long
、Float
和 Double
,分别对应于标量基本类型 byte
、short
、int
、long
、float
和 double
。
支持的形状集对应于 64、128、256 和 512 位的向量大小,以及 最大 位。512 位形状可以将 byte
打包到 64 个通道中,或者将 int
打包到 16 个通道中,并且这种形状的向量可以一次对 64 个 byte
或 16 个 int
进行操作。最大 位形状支持当前架构的最大向量大小。这使得能够支持 ARM SVE 平台,其中平台实现可以支持从 128 到 2048 位的任何固定大小,以 128 位为增量。
我们相信这些简单的形状在所有相关平台上都足够通用且有用。然而,在这个 API 的孵化期间,当我们对未来平台进行实验时,我们可能会进一步修改形状参数的设计。这样的工作不在这个项目的早期范围之内,但这些可能性部分地影响了当前形状在向量 API 中的作用。(有关进一步的讨论,请参见下面的 未来工作 部分。)
元素类型和形状的组合确定了向量的 种类,由 VectorSpecies<E>
表示。
向量上的操作分为 通道级 或 跨通道级。
- 通道级 操作并行地将标量运算符(如加法)应用于一个或多个向量的每个通道。通道级操作通常(但不总是)产生具有相同长度和形状的向量。通道级操作进一步分为一元、二元、三元、测试或转换操作。
- 跨通道级 操作在整个向量上应用操作。跨通道级操作产生标量或可能具有不同形状的向量。跨通道级操作进一步分为置换或归约操作。
为了减少 API 的表面,我们为每类操作定义了集合方法。这些方法将运算符常量作为输入;这些常量是 VectorOperator.Operator
类的实例,并在 VectorOperators
类的静态最终字段中定义。为了方便起见,我们为一些常见的 全功能 操作(如加法和乘法)定义了专用方法,可以替代通用方法使用。
向量上的某些操作,如转换和重新解释,本质上是 形状改变 的;即,它们产生的向量的形状与输入向量的形状不同。向量计算中的形状改变操作可能会对可移植性和性能产生负面影响。因此,在适用的情况下,API 为每个形状改变操作定义了 形状不变 的变体。为了获得最佳性能,开发人员应尽可能使用形状不变的操作编写形状不变的代码。在 API 规范中,形状改变的操作被标识为这样的操作。
Vector<E>
类声明了一组由所有元素类型支持的常见向量操作的方法。对于特定于元素类型的操作,有 Vector<E>
的六个抽象子类,每个支持的元素类型一个:ByteVector
、ShortVector
、IntVector
、LongVector
、FloatVector
和 DoubleVector
。这些特定类型的子类定义了由于方法签名引用元素类型或相关数组类型而与元素类型绑定的其他操作。这样的操作的例子包括归约(例如,将所有通道求和为一个标量值),以及将向量的元素复制到数组中。这些子类还定义了特定于整数子类型的其他全功能操作(例如,按位操作如逻辑或),以及特定于浮点类型的操作(例如,超越数学函数如指数运算)。
作为实现问题,Vector<E>
的这些特定类型的子类由针对不同向量形状的具体子类进一步扩展。这些具体子类不是公共的,因为没有必要提供特定于类型和形状的操作。这将 API 表面减少为关注点的总和,而不是乘积。具体 Vector
类的实例通过在基类 Vector<E>
及其特定类型的子类中定义的工厂方法获得。这些工厂方法以所需向量实例的种类作为输入,并产生各种实例,例如其元素为默认值的向量实例(即零向量),或从给定数组初始化的向量实例。
为了支持控制流,一些向量操作可选地接受由公共抽象类 VectorMask<E>
表示的掩码。掩码中的每个元素是对应于向量通道的布尔值。掩码选择应用操作的通道:如果通道的掩码元素为 true,则应用该操作,如果掩码为 false,则采取一些替代操作。
与向量类似,VectorMask<E>
的实例是为每个元素类型和长度组合定义的非公共具体子类的实例。在操作中使用的 VectorMask<E>
的实例应该与参与操作的向量实例具有相同的类型和长度。向量比较操作产生掩码,然后可以将其用作其他操作的输入,以有选择地对某些通道进行操作,从而模拟流控制。掩码也可以使用 VectorMask<E>
类中的静态工厂方法创建。
我们预计掩码将在开发对形状通用的向量计算中发挥重要作用。这个预期是基于谓词寄存器(与掩码等效)在 ARM 可扩展向量扩展和 Intel 的 AVX-512 中的核心重要性。
在这样的平台上,VectorMask<E>
的实例被映射到谓词寄存器,并且接受掩码的操作被编译为接受谓词寄存器的向量指令。在不支持谓词寄存器的平台上,应用一种效率较低的方法:VectorMask<E>
的实例在可能的情况下被映射到兼容的向量寄存器,并且通常接受掩码的操作由等效的未掩码操作和混合操作组成。
为了支持跨通道置换操作,一些向量操作接受由公共抽象类 VectorShuffle<E>
表示的混洗。混洗中的每个元素是对应于通道索引的 int
值。混洗是通道索引的映射,描述了通道元素从给定向量到结果向量的移动。
与向量和掩码类似,VectorShuffle<E>
的实例是为每个元素类型和长度组合定义的非公共具体子类的实例。在操作中使用的 VectorShuffle<E>
的实例应该与参与操作的向量实例具有相同的类型和长度。
示例
以下是对数组元素进行的简单标量计算:
void scalarComputation(float[] a, float[] b, float[] c) {
for (int i = 0; i < a.length; i++) {
c[i] = (a[i] * a[i] + b[i] * b[i]) * -1.0f;
}
}
(我们假设数组参数的长度相同。)
以下是使用向量 API 的等效向量计算:
static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;
void vectorComputation(float[] a, float[] b, float[] c) {
int i = 0;
int upperBound = SPECIES.loopBound(a.length);
for (; i < upperBound; i += SPECIES.length()) {
// FloatVector va, vb, vc;
var va = FloatVector.fromArray(SPECIES, a, i);
var vb = FloatVector.fromArray(SPECIES, b, i);
var vc = va.mul(va)
.add(vb.mul(vb))
.neg();
vc.intoArray(c, i);
}
for (; i < a.length; i++) {
c[i] = (a[i] * a[i] + b[i] * b[i]) * -1.0f;
}
}
首先,我们从 FloatVector
获取一个对当前架构最优的首选种类。我们将其存储在一个 static final
字段中,以便运行时编译器将该值视为常量,从而可以更好地优化向量计算。然后主循环以向量长度(即种类长度)的步长遍历输入数组。它从数组 a
和 b
的相应索引处加载给定种类的 float
向量,流畅地执行算术运算,然后将结果存储到数组 c
中。如果在最后一次迭代后还有任何数组元素剩余,则使用普通的标量循环计算这些 尾部 元素的结果。
此实现对大型数组实现了最佳性能。在支持 AVX 的 Intel x64 处理器上,HotSpot C2 编译器生成类似于以下的机器代码:
0.43% / │ 0x0000000113d43890: vmovdqu 0x10(%r8,%rbx,4),%ymm0
7.38% │ │ 0x0000000113d43897: vmovdqu 0x10(%r10,%rbx,4),%ymm1
8.70% │ │ 0x0000000113d4389e: vmulps %ymm0,%ymm0,%ymm0
5.60% │ │ 0x0000000113d438a2: vmulps %ymm1,%ymm1,%ymm1
13.16% │ │ 0x0000000113d438a6: vaddps %ymm0,%ymm1,%ymm0
21.86% │ │ 0x0000000113d438aa: vxorps -0x7ad76b2(%rip),%ymm0,%ymm0
7.66% │ │ 0x0000000113d438b2: vmovdqu %ymm0,0x10(%r9,%rbx,4)
26.20% │ │ 0x0000000113d438b9: add $0x8,%ebx
6.44% │ │ 0x0000000113d438bc: cmp %r11d,%ebx
\ │ 0x0000000113d438bf: jl 0x0000000113d43890
这是使用向量 API 的原型和在 Project Panama 的开发仓库的 vectorIntrinsics
分支 上找到的实现对上述代码进行 JMH 微基准测试的输出。这些生成的机器代码的热点区域显示了对向量寄存器和向量指令的清晰翻译。我们禁用了循环展开(通过 HotSpot 选项 -XX:LoopUnrollLimit=0
)以便使翻译更清晰;否则,HotSpot 将使用现有的 C2 循环优化来展开此代码。所有 Java 对象分配都被省略。
(在这个特定的例子中,HotSpot 能够自动向量化标量计算,并且它将生成类似的向量指令序列。主要区别在于自动向量化器为乘以 -1.0f
生成一个向量乘法指令,而向量 API 实现生成一个向量异或指令来翻转符号位。然而,这个例子的关键点是展示向量 API 并显示其实现如何生成向量指令,而不是与自动向量化器进行比较。)
在支持谓词寄存器的平台上,上面的例子可以写得更简单,无需标量循环来处理尾部元素,同时仍然实现最佳性能:
void vectorComputation(float[] a, float[] b, float[] c) {
for (int i = 0; i < a.length; i += SPECIES.length()) {
// VectorMask<Float> m;
var m = SPECIES.indexInRange(i, a.length);
// FloatVector va, vb, vc;
var va = FloatVector.fromArray(SPECIES, a, i, m);
var vb = FloatVector.fromArray(SPECIES, b, i, m);
var vc = va.mul(va)
.add(vb.mul(vb))
.neg();
vc.intoArray(c, i, m);
}
}
在循环体中,我们获取一个依赖于循环的掩码,用于输入到加载和存储操作中。当 i < SPECIES.loopBound(a.length)
时,掩码 m
声明所有通道都被设置。对于循环的最后一次迭代,当 SPECIES.loopBound(a.length) <= i < a.length
且 (a.length - i) <= SPECIES.length()
时,掩码可能声明一个未设置通道的后缀。加载和存储操作不会抛出越界异常,因为掩码阻止访问超出数组长度的部分。
我们希望开发人员以上述风格为所有支持的平台编写代码并实现最佳性能,但在当今没有谓词寄存器的平台上,上述方法不是最优的。理论上,C2 编译器可以进行增强以转换循环,剥离最后一次迭代并从循环体中移除掩码。这仍然是一个需要进一步研究的领域。
运行时编译
向量 API 有两个实现。第一个在 Java 中实现操作,因此它是功能性的但不是最优的。第二个为 HotSpot C2 运行时编译器定义内在向量操作,以便在可用时将向量计算编译为适当的硬件寄存器和向量指令。
为了避免 C2 内在函数的爆炸式增长,我们定义了对应于各种操作类型(如一元、二元、转换等)的通用内在函数,这些内在函数接受一个描述要执行的特定操作的参数。大约二十个新的内在函数支持整个 API 的内在化。
我们最终期望将向量类声明为值类,如 Project Valhalla 所提议的(见 值对象 和 JEP 401(基本类型类))。同时,Vector<E>
及其子类被视为 基于值的类,因此应避免对其实例进行与标识敏感的操作。虽然向量实例抽象地由通道中的元素组成,但那些元素不会被 C2 标量化——向量的值被视为一个整体单元,就像一个 int
或一个 long
,它映射到适当大小的向量寄存器。C2 对向量实例进行特殊处理,以克服逃逸分析的限制并避免装箱。在未来,我们将使这种特殊处理与 Valhalla 的值对象保持一致。
Intel SVML 内在函数用于超越运算
向量 API 支持对浮点向量进行超越和三角函数的通道级操作。在 x64 上,我们利用 Intel 短向量数学库(SVML)为这些操作提供优化的内在实现。内在操作具有与 java.lang.Math
中定义的相应标量操作相同的数值属性。
SVML 操作的汇编源文件在 jdk.incubator.vector
模块的源代码中,位于特定于操作系统的目录下。JDK 构建过程将这些针对目标操作系统的源文件编译为一个特定于 SVML 的共享库。这个库相当大,接近一兆字节。如果通过 jlink
构建的 JDK 映像省略了 jdk.incubator.vector
模块,那么 SVML 库就不会被复制到映像中。
目前该实现仅支持 Linux 和 Windows。我们稍后会考虑 macOS 支持,因为提供带有所需指令的汇编源文件是一项不小的工作量。
HotSpot 运行时将尝试加载 SVML 库,如果存在,将 SVML 库中的操作绑定到命名的存根例程。C2 编译器生成根据操作和向量种类(即元素类型和形状)调用适当存根例程的代码。
在未来,如果 Project Panama 扩展其对本机调用约定的支持以支持向量值,那么向量 API 实现可能能够从外部源加载 SVML 库。如果这种方法没有性能影响,那么就不再需要以源代码形式包含 SVML 并将其构建到 JDK 中。在此之前,考虑到潜在的性能提升,我们认为上述方法是可以接受的。
未来工作
- 如上文所述,我们最终期望将向量类声明为 值类。有关使向量 API 与 Valhalla 保持一致的持续努力,请参阅 Project Valhalla 的代码仓库的
lworld+vector
分支。我们进一步期望利用 Project Valhalla 的值类通用特化,以便Vector<E>
的实例是值对象,其中E
是一个基本类型类,如int
,而不是其装箱类Integer
。一旦我们对基本类型类有了通用特化,特定类型的Vector<E>
的子类型,如IntVector
,可能就不再需要了。 - 我们可能会添加对 IEEE 浮点 binary16 值(float16 值)向量的支持。这也依赖于 Project Valhalla,需要我们将 float16 值表示为值对象,并增强向量 API 实现以利用 float16 值向量的向量硬件指令。有关探索性工作,请参阅 Project Panama 的向量 API 代码仓库的
vectorIntrinsics+fp16
分支。 - 我们预计将增强实现以改进包含向量化代码的循环的优化,并随着时间的推移逐步提高性能。
- 我们还预计将增强组合单元测试,以断言 C2 生成向量硬件指令。目前,单元测试在没有验证的情况下假设重复执行足以使 C2 生成向量硬件指令。我们将探索使用 C2 的 IR 测试框架 来跨平台断言向量节点存在于 IR 图中(例如,使用 正则表达式匹配)。如果这种方法有问题,我们可能会探索一种基本方法,使用非产品级的
-XX:+TraceNewVectors
标志来打印向量节点。 - 我们将评估合成向量形状的定义,以更好地控制循环展开和矩阵操作,并考虑对排序和解析算法的适当支持。(有关更多详细信息,请参阅 此演示文稿。)
替代方案
HotSpot 的自动向量化是一种替代方法,但它需要大量的工作。此外,与向量 API 相比,它仍然很脆弱且有限,因为具有复杂控制流的标量代码自动向量化非常难以实现。
一般来说,即使经过几十年的研究——特别是对于 FORTRAN 和 C 数组循环——似乎标量代码的自动向量化对于优化临时用户编写的循环不是一个可靠的策略,除非用户非常小心地注意关于编译器准备自动向量化哪些循环的未写明的契约。很容易编写一个无法自动向量化的循环,原因是人类读者无法察觉的。即使在 HotSpot 中,多年的自动向量化工作也给我们留下了很多仅在特殊情况下起作用的优化机制。我们希望更频繁地使用这些机制!
测试
我们将开发组合单元测试,以确保对所有操作、所有支持的类型和形状在各种数据集上的覆盖。
我们还将开发性能测试,以确保满足性能目标并且向量计算有效地映射到向量指令。这可能包括 JMH 微基准测试,但也需要更现实的有用算法示例。这样的测试可能最初位于项目特定的存储库中。考虑到测试的比例和生成方式,在集成到主存储库之前可能需要进行整理。
风险和假设
- 存在 API 偏向于 x64 架构上支持的 SIMD 功能的风险,但通过对 AArch64 的支持可以减轻这种风险。这主要适用于明确固定的支持形状集,这对以形状通用的方式编码算法有偏见。我们认为向量 API 的大多数其他操作偏向于可移植算法。为了减轻这种风险,我们将考虑其他架构,特别是 ARM 标量向量扩展架构,其编程模型动态调整以适应硬件支持的单一固定形状。我们欢迎并鼓励在 HotSpot 的 ARM 特定区域工作的 OpenJDK 贡献者参与这项工作。
- 向量 API 使用装箱类型(例如
Integer
)作为基本类型(例如int
)的代理。这个决定是由 Java 泛型的当前限制所迫,Java 泛型对基本类型不友好。当 Project Valhalla 最终引入更强大的泛型时,当前的决定将显得很尴尬,并且可能需要改变。我们假设在不过度破坏向后兼容性的情况下,这样的改变是可能的。