跳到主要内容

JEP 417: Vector API(第三次孵化)

QWen Max 中英对照 JEP 417: Vector API (Third Incubator)

总结

引入一个 API,用于表达矢量计算,该 API 在支持的 CPU 架构上能够可靠地在运行时编译为最佳的矢量指令,从而实现优于等效标量计算的性能。

历史

Vector API 最早由 JEP 338 提出,并作为 孵化 API 集成到 Java 16 中。JEP 414 提议了第二轮孵化,并将其集成到 Java 17 中。

我们在此建议纳入针对反馈的改进、性能提升以及其他重要的实现增强。我们包含以下显著变更:

目标

  • 清晰简洁的 API — 该 API 应能够清晰简洁地表达广泛的矢量计算,这些计算由循环内的矢量操作序列组成,并可能包含控制流。应该可以表达出对于矢量大小或每个矢量的通道数通用的计算,从而使这些计算能够在支持不同矢量大小的硬件之间移植。

  • 平台无关性 — 该 API 应与 CPU 架构无关,从而支持在多个具有矢量指令的架构上实现。正如 Java API 中常见的那样,当平台优化和可移植性发生冲突时,偏向于使 API 可移植,即使这会导致某些特定平台的习惯用法无法在可移植代码中表达。

  • 在 x64 和 AArch64 架构上的可靠运行时编译和性能 — 在具备能力的 x64 架构上,Java 运行时(特别是 HotSpot C2 编译器)应将矢量操作编译为相应的高效且高性能的矢量指令,例如 Streaming SIMD Extensions (SSE) 和 Advanced Vector Extensions (AVX) 支持的指令。开发人员应有信心,他们表达的矢量操作将可靠地紧密映射到相关的矢量指令。在具备能力的 ARM AArch64 架构上,C2 同样会将矢量操作编译为 NEONSVE 支持的矢量指令。

  • 优雅降级 — 有时矢量计算无法在运行时完全表达为矢量指令序列,可能是因为架构不支持某些所需指令。在这种情况下,Vector API 的实现应优雅降级并仍能正常工作。这可能涉及在矢量计算无法高效编译为矢量指令时发出警告。在没有矢量支持的平台上,优雅降级将生成与手动展开循环竞争的代码,其中展开因子是所选矢量中的通道数。

非目标

  • 增强 HotSpot 中现有的自动向量化算法并不是目标。

  • 在 x64 和 AArch64 以外的 CPU 架构上支持向量指令并不是目标。然而,重要的是要声明,正如在目标中所表达的,该 API 不应排除此类实现。

  • 支持 C1 编译器并不是目标。

  • 保证支持 Java 平台对标量操作所要求的严格浮点计算并不是目标。对标量浮点数执行的浮点操作的结果可能与对标量浮点数向量执行的等效浮点操作不同。任何偏差都将被清楚地记录。这一非目标并不排除表达或控制浮点向量计算所需精度或可重复性的选项。

动机

向量计算由一系列对向量的操作组成。一个向量包含一个(通常)固定序列的标量值,其中标量值对应于硬件定义的向量通道数量。对具有相同数量通道的两个向量应用二元运算时,会在每个通道上对标量值执行等效的标量运算,这些标量值分别来自两个向量中的对应部分。这通常被称为 单指令多数据 (SIMD)。

向量操作表达了一定程度的并行性,这种并行性使得在单个 CPU 周期内可以完成更多工作,从而可能带来显著的性能提升。例如,给定两个向量,每个向量包含八个整数的序列(即八个通道),这两个向量可以通过一条硬件指令相加。向量加法指令可以在通常对两个整数进行操作、执行一次整数加法所需的时间内,对十六个整数进行操作,完成八次整数加法。

HotSpot 已经支持 自动矢量化,它可以将标量操作转换为超字(superword)操作,然后映射到矢量指令。可转换的标量操作集是有限的,并且对代码形态的变化也较为脆弱。此外,可能仅使用了可用矢量指令的一个子集,从而限制了生成代码的性能。

如今,一个希望编写能被可靠地转换为超字(superword)操作的标量操作的开发者,需要理解 HotSpot 的自动向量化算法及其局限性,以实现稳定且可持续的性能。在某些情况下,可能无法编写出可转换的标量操作。例如,HotSpot 不会将用于计算数组哈希码的简单标量操作转换为超字操作(因此 Arrays::hashCode 方法无法被向量化),也无法自动向量化用于按字典顺序比较两个数组的代码(因此我们添加了一个用于字典序比较的内联函数)。

Vector API 旨在通过提供一种使用现有的 HotSpot 自动矢量化器编写复杂矢量算法的方法来改善这种情况,同时其用户模型使矢量化更加可预测且健壮。手写的矢量循环可以表达高性能算法,例如矢量化的 hashCode 或专门的数组比较,而自动矢量化器可能永远不会优化这些算法。众多领域都可以从这种显式的 Vector API 中受益,包括机器学习、线性代数、密码学、金融以及 JDK 自身内部的代码。

描述

向量由抽象类 Vector<E> 表示。类型变量 E 被实例化为向量所涵盖的标量基本整数或浮点元素类型的装箱类型。向量还有一个 形状(shape),它定义了向量的大小(以位为单位)。向量的形状决定了当 HotSpot C2 编译器编译向量计算时,Vector<E> 的一个实例如何映射到硬件向量寄存器。向量的长度,即通道数或元素数,是向量大小除以元素大小的结果。

支持的元素类型 (E) 集合为 ByteShortIntegerLongFloatDouble,分别对应于标量基本类型 byteshortintlongfloatdouble

所支持的形状集对应于 64、128、256 和 512 位的向量大小,以及 max 位。512 位的形状可以将 byte 打包到 64 条车道,或将 int 打包到 16 条车道,这种形状的向量可以一次操作 64 个 byte 或一次操作 16 个 intmax-bits 形状支持当前架构的最大向量大小。这使得支持 ARM SVE 平台成为可能,该平台的实现可以支持从 128 位到 2048 位之间的任何固定大小,以 128 位为增量。

我们相信,这些简单的形状具有足够的通用性,可以在所有相关平台上发挥作用。然而,在孵化此 API 期间,当我们尝试未来的平台时,我们可能会进一步修改形状参数的设计。这样的工作并不在该项目的早期范围内,但这些可能性部分地影响了形状在 Vector API 中的当前角色。(更多讨论请参见下面的未来工作部分。)

元素类型和形状的组合决定了向量的 种类,由 VectorSpecies<E> 表示。

对矢量的操作分为通道内跨通道

  • 一个 按通道 操作会将一个标量运算符(例如加法)并行应用于一个或多个向量的每个通道。按通道操作通常但不总是生成相同长度和形状的向量。按通道操作进一步分为一元、二元、三元、测试或转换操作。

  • 一个 跨通道 操作会在整个向量上应用一个操作。跨通道操作会产生一个标量或者可能具有不同形状的向量。跨通道操作进一步分为排列或归约操作。

为了减少 API 的表面积,我们为每类操作定义了集合方法。这些方法将运算符常量作为输入;这些常量是 VectorOperator.Operator 类的实例,并在 VectorOperators 类中的静态 final 字段中定义。为了方便起见,我们为一些常见的全服务操作(如加法和乘法)定义了专用方法,这些方法可以用来替代通用方法。

向量的某些操作(例如转换和重新解释)本质上是改变形状的;也就是说,它们生成的向量形状与其输入的形状不同。向量计算中的改变形状的操作可能会对可移植性和性能产生负面影响。因此,API 在适用的情况下定义了每种改变形状操作的形状不变版本。为了获得最佳性能,开发人员应尽可能使用形状不变的操作编写形状不变的代码。在 API 规范中,改变形状的操作被明确标识为如此。

Vector<E> 类声明了一组所有元素类型都支持的常用向量操作方法。对于特定于元素类型的操作,Vector<E> 有六个抽象子类,每个支持的元素类型各一个:ByteVectorShortVectorIntVectorLongVectorFloatVectorDoubleVector。这些特定类型的子类定义了额外的操作,这些操作绑定到元素类型,因为方法签名引用的是元素类型或相关的数组类型。此类操作的示例包括归约(例如,将所有通道求和为一个标量值)以及将向量的元素复制到数组中。这些子类还定义了针对整数子类型的额外全服务操作(例如,按位操作如逻辑或),以及针对浮点类型的特定操作(例如,超越数学函数如指数运算)。

作为一个实现问题,这些 Vector<E> 的类型特定子类会进一步由针对不同向量形状的具体子类进行扩展。这些具体子类并不公开,因为没有必要提供针对类型和形状的特定操作。这将 API 的覆盖范围减少为关注点的总和,而非乘积。具体 Vector 类的实例是通过在基础 Vector<E> 类及其类型特定子类中定义的工厂方法获取的。这些工厂方法以所需向量实例的种类作为输入,并生成各种类型的实例,例如元素为默认值的向量实例(即零向量),或者从给定数组初始化的向量实例。

为了支持控制流,某些矢量操作可选地接受由公共抽象类 VectorMask<E> 表示的掩码。掩码中的每个元素都是一个布尔值,对应于矢量通道。掩码选择操作所应用的通道:如果该通道的掩码元素为 true,则应用该操作;如果掩码为 false,则采取某种替代动作。

与向量类似,VectorMask<E> 的实例是为每种元素类型和长度组合定义的非公共具体子类的实例。在操作中使用的 VectorMask<E> 实例应与参与操作的向量实例具有相同的类型和长度。向量比较操作会产生掩码,然后可以将这些掩码用作其他操作的输入,以选择性地对某些通道进行操作,从而模拟流控制。还可以使用 VectorMask<E> 类中的静态工厂方法来创建掩码。

我们预计,掩码在对于形状通用的矢量计算的发展中将发挥重要作用。这一预期基于谓词寄存器(即掩码的等价物)在 ARM 可扩展矢量扩展(Scalable Vector Extensions)和 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;
}
}
java

(我们假设数组参数的长度相同。)

下面是一个使用 Vector 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;
}
}
java

首先,我们从 FloatVector 中获取一个形状最适合当前架构的首选物种(species)。我们将其存储在一个 static final 字段中,以便运行时编译器将该值视为常量,从而更好地优化向量计算。主循环随后以向量长度(即物种长度)为步长遍历输入数组。它从数组 ab 中对应索引处加载给定物种的 float 向量,流畅地执行算术运算,然后将结果存储到数组 c 中。如果在最后一次迭代后有任何数组元素剩余,则这些尾部元素的结果将通过普通的标量循环计算。

该实现能够在大型数组上达到最佳性能。HotSpot C2 编译器在支持 AVX 的 Intel x64 处理器上生成类似于以下内容的机器代码:

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

这是使用 Project Panama 开发仓库的 vectorIntrinsics 分支 中 Vector API 的原型和实现,针对上述代码运行 JMH 微基准测试的输出结果。这些生成的机器代码的热点区域清楚地展示了向量寄存器和向量指令的转换。为了使转换更加清晰,我们禁用了循环展开;否则,HotSpot 会使用现有的 C2 循环优化来展开此代码。所有的 Java 对象分配均被省略。

在支持谓词寄存器的平台上,上面的例子可以写得更简单,无需使用标量循环来处理尾部元素,同时仍然能够实现最佳性能:

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);
}
}
java

在循环体中,我们为输入到加载和存储操作获取一个循环相关的掩码。当 i < SPECIES.loopBound(a.length) 时,掩码 m 声明所有通道均被设置。在循环的最后一次迭代中,当 SPECIES.loopBound(a.length) <= i < a.length(a.length - i) <= SPECIES.length() 时,掩码可能会声明一个未设置通道的后缀。由于掩码阻止了对超出数组长度部分的访问,因此加载和存储操作不会抛出越界异常。

我们更希望开发者为所有支持的平台都采用上述风格编写代码,以实现最佳性能,但在当前没有谓词寄存器的平台上,上述方法并非最优。理论上,C2 编译器可以进一步增强,以转换循环,剥离最后一次迭代并移除循环体中的掩码。这仍然是一个值得进一步研究的领域。

运行时编译

Vector API 有两种实现。第一种是用 Java 实现操作,因此它是可用的,但不是最优的。第二种是为 HotSpot C2 运行时编译器定义了内在的矢量操作,这样当硬件寄存器和矢量指令可用时,它就可以将矢量计算编译为合适的硬件寄存器和矢量指令。

为避免 C2 内部函数爆炸式增长,我们定义了与各种操作(如一元、二元、转换等)相对应的广义内部函数,这些函数接受一个描述要执行的具体操作的参数。大约二十个新的内部函数支持整个 API 的内部函数化。

我们最终期望将矢量类声明为基本类,正如 Project ValhallaJEP 401(基本对象) 中所提出的那样。在此期间,Vector<E> 及其子类被视为 基于值的类,因此应避免对其实例进行依赖标识的操作。尽管矢量实例在抽象上由通道中的元素组成,但这些元素并未被 C2 标量化 —— 矢量的值被视为一个整体单元,类似于 intlong,映射到适当大小的矢量寄存器。C2 对矢量实例进行了特殊处理,以克服逃逸分析的局限性并避免装箱操作。

用于超越运算的 Intel SVML 内部函数

Vector API 支持对浮点向量进行超越函数和三角函数的按通道操作。在 x64 平台上,我们利用 Intel 短向量数学库(SVML)为这些操作提供优化的内联函数实现。这些内联操作具有与 java.lang.Math 中定义的相应标量操作相同的数值属性。

SVML 操作的汇编源文件位于 jdk.incubator.vector 模块的源代码中,存放在特定操作系统的目录下。JDK 的构建过程会将这些源文件针对目标操作系统编译为一个 SVML 专用的共享库。该库相当大,大小接近 1 MB。如果通过 jlink 构建的 JDK 镜像中省略了 jdk.incubator.vector 模块,那么 SVML 库将不会被复制到该镜像中。

该实现目前仅支持 Linux 和 Windows。由于提供带有必要指令的汇编源文件需要大量工作,我们将在稍后考虑 macOS 支持。

HotSpot 运行时会尝试加载 SVML 库,如果存在,则将 SVML 库中的操作绑定到命名的存根例程。C2 编译器生成的代码会根据操作和矢量种类(即元素类型和形状)调用相应的存根例程。

未来,如果 Project Panama 扩展了其对本地调用约定的支持以支持向量值,那么 Vector API 的实现可能会从外部源加载 SVML 库。如果这种方法没有性能影响,那么就不再需要以源代码形式包含 SVML 并将其构建到 JDK 中。在此之前,考虑到潜在的性能提升,我们认为上述方法是可以接受的。

未来工作

  • 如上所述,我们最终期望将矢量类声明为原始类。此外,我们期望利用 Project Valhalla 的原始类泛型特化功能,使得 Vector<E> 的实例可以是具体类型为原始类型的原始值。这将使优化和表达矢量计算变得更加容易。一旦我们能够对原始类进行泛型特化,像 IntVector 这样的特定类型的 Vector<E> 子类型可能就不再需要了。我们计划在多个版本中孵化此 API,并随着原始类及相关功能的可用性对其进行调整。

  • 我们计划在 JEP 412(外部函数与内存 API) 脱离孵化阶段后,增强该 API 以使用它加载和存储矢量。描述矢量种类的内存布局可能会被证明是有用的,例如用于跨由矢量元素组成的内存段进行步进。

  • 我们预计会改进实现,以优化包含矢量化代码的循环性能,并随着时间推移逐步提升整体性能。

  • 我们还计划增强组合单元测试,以验证 C2 编译器生成矢量硬件指令。目前的单元测试假设通过足够多次重复执行就能生成矢量硬件指令,但并未加以验证。我们将探索使用 C2 的IR 测试框架,跨平台断言 IR 图中存在矢量节点(例如,使用正则匹配)。如果这种方法存在问题,我们可能会探索一种更基础的方法,并使用非生产性的 -XX:+TraceNewVectors 标志来打印矢量节点。

替代方案

HotSpot 的自动矢量化是一种替代方法,但需要大量工作。而且,与 Vector API 相比,它仍然很脆弱且有局限性,因为使用复杂控制流进行自动矢量化非常难以执行。

总体而言,即使经过数十年的研究 —— 尤其是对于 FORTRAN 和 C 数组循环 —— 除非用户异常仔细地关注编译器准备自动向量化的确切循环的不成文约定,否则标量代码的自动向量化似乎并不是优化用户随意编写的循环的可靠策略。编写一个无法自动向量化的循环太容易了,原因可能是人类读者无法察觉的。即使在 HotSpot 中,多年来的自动向量化工作只给我们留下了许多仅在特殊场合工作的优化机制。我们希望更频繁地利用这些机制!

测试

我们将开发组合单元测试,以确保覆盖所有操作、所有支持的类型和形状,以及各种数据集。

我们还将开发性能测试,以确保达到性能目标,并且矢量计算能够高效映射到矢量指令。这可能包括 JMH 微基准测试,但还需要更多实际有用的算法示例。此类测试最初可能会存放在特定于项目的代码库中。鉴于测试的比例以及它们的生成方式,在集成到主代码库之前可能需要进行整理。

风险与假设

  • API 可能会对 x64 架构上支持的 SIMD 功能产生偏向性风险,但通过支持 AArch64 可以减轻这种风险。这主要适用于显式固定的支持形状集,这些形状不利于以形状通用的方式编写算法。我们认为 Vector API 的大多数其他操作偏向于可移植算法。为了降低这种风险,我们将考虑其他架构,特别是 ARM Scalar Vector Extension 架构,其编程模型会动态调整为硬件支持的单一固定形状。我们欢迎并鼓励在 HotSpot 的 ARM 特定领域工作的 OpenJDK 贡献者参与这一工作。

  • Vector API 使用装箱类型(例如 Integer)作为基本类型的代理(例如 int)。这一决定是由于 Java 泛型的当前限制导致的,因为 Java 泛型对基本类型并不友好。当 Project Valhalla 最终引入功能更强大的泛型时,当前的决定可能会显得不够理想,并且可能需要更改。我们假设这些更改可以在不会导致过度向后不兼容的情况下实现。