跳到主要内容

JEP 426:矢量 API(第四孵化版)

QWen Max 中英对照 JEP 426 Vector API (Fourth Incubator)

概述

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

历史

Vector API 最早由 JEP 338 提出,并作为 孵化 API 集成到 JDK 16 中。第二轮孵化由 JEP 414 提议,并集成到 JDK 17 中。第三轮孵化由 JEP 417 提议,并集成到 JDK 18 中。

我们在此建议纳入一些改进,以响应反馈以及性能改进和其他重要的实现增强。我们包含以下显著的变化:

  • 增强 API 以根据 JEP 424:外部函数与内存 (FFM) API(预览版) 定义,加载和存储向量到 MemorySegment 或从中加载和存储。FFM API 已经足够成熟,因此我们有信心将这一依赖添加到 Vector API 中。我们将移除操作于 byte[]ByteBuffer 的等效 API 点,因为可以为两者获取一个 MemorySegment。使用 MemorySegment 将能够创建 超对齐 的内存区域,这些区域对齐到向量的字节长度。在某些架构上,这种对齐可以在加载和存储向量时提供更优的性能。

  • 添加两个新的跨通道向量操作,压缩 (compress) 及其逆操作 扩展 (expand),以及一个互补的向量掩码 压缩 (compress) 操作。压缩 向量操作会根据掩码选择源向量的通道,并按通道顺序映射到目标向量;扩展 操作则执行相反的操作。压缩 操作在过滤查询结果时非常有用;例如,它可以从未零整数元素从源向量中选出,并按照源通道顺序连续写入到目标向量中。

  • 扩展支持的按位整数逐通道操作集,包括:

    • 计算一位的数量,
    • 计算前导零的数量,
    • 计算尾随零的数量,
    • 反转位的顺序,
    • 反转字节的顺序,以及
    • 压缩和扩展位。

    后两项操作类似于跨通道的 压缩扩展 操作,只是它们映射的是位而不是整个通道。(有关按位 压缩扩展 的更多信息,请参见 Henry S. Warren 的《Hacker's Delight》第 7-4 和 7-5 节,Addison-Wesley,2013 年。)

    与此 JEP 同时但分开进行,我们将为封装的基本类型 IntegerLong 类型添加按位标量 压缩 (compress)扩展 (expand) 操作的方法,这些方法本身也具有独立的价值。我们将根据这些新的标量方法来指定按位 压缩扩展 向量操作。

目标

  • 清晰简洁的 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 位形状支持当前架构的最大矢量大小。这使得对 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> 类及其类型特定子类中的工厂方法获取。这些工厂方法以所需向量实例的种类(species)作为输入,并生成各种类型的实例,例如元素为默认值的向量实例(即零向量),或者从给定数组初始化的向量实例。

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

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

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

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

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

首先,我们从 FloatVector 中获取一个形状最适合当前架构的首选种类(species)。我们将其存储在一个 static final 字段中,以便运行时编译器将该值视为常量,从而更好地优化向量计算。主循环随后以向量长度(即种类长度)为步幅迭代输入数组。它从数组 ab 中对应索引处加载给定种类的 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

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

(在本例中,HotSpot 能够自动向量化标量计算,并生成类似的向量指令序列。主要区别在于,自动向量化器会为乘以 -1.0f 生成一个向量乘法指令,而 Vector API 的实现则生成一个向量 XOR 指令来翻转符号位。然而,本例的关键点是展示 Vector 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 编译器可以进行增强,以转换循环,剥离最后一次迭代并移除循环体中的掩码。这仍然是一个有待进一步研究的领域。

运行时编译

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 库,并且如果存在该库,会将其中的操作绑定到命名的存根例程。C2 编译器生成的代码会根据操作和向量种类(即元素类型和形状)调用相应的存根例程。

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

未来工作

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

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

  • 我们还计划增强组合单元测试,以验证 C2 编译器生成矢量硬件指令的能力。目前的单元测试假设(未经验证)重复执行足以促使 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 最终引入更强大的泛型时,当前的决定可能会显得笨拙,并且可能需要更改。我们假设这些更改可以在不过度向后不兼容的情况下实现。