JEP 338: Vector API(孵化中)
总结
提供一个 孵化模块 的初始迭代版本,jdk.incubator.vector
,用于表达矢量计算,这些计算在支持的 CPU 架构上能够可靠地在运行时编译为最佳的矢量硬件指令,从而实现比等效标量计算更优越的性能。
目标
-
清晰简洁的 API: 该 API 应能够清晰简洁地表达广泛包含一系列矢量运算的矢量计算,这些矢量运算通常在循环内进行组合,可能带有控制流。应该可以表达对矢量大小(或每个矢量的通道数)通用的计算,从而使这些计算能够在支持不同矢量大小的硬件之间具有可移植性(详情见下一个目标)。
-
平台无关性: 该 API 应与架构无关,从而支持在多个支持矢量硬件指令的 CPU 架构上运行时实现。正如 Java API 中常见的那样,在平台优化和可移植性发生冲突时,偏向于使矢量 API 具有可移植性,即使某些特定平台的习惯用法无法直接在可移植代码中表达。x64 和 AArch64 性能的下一个目标代表了所有支持 Java 的平台上的适当性能目标。在这方面,ARM 可扩展矢量扩展 (SVE) 特别值得关注,以确保 API 能够支持此架构,尽管截至目前尚无已知的生产硬件实现。
-
在 x64 和 AArch64 架构上的可靠运行时编译和性能: Java 运行时,特别是 HotSpot C2 编译器,应在有能力的 x64 架构上将一系列矢量操作编译为相应的硬件矢量指令序列,例如由 流式 SIMD 扩展 (SSE) 和 高级矢量扩展 (AVX) 扩展支持的指令,从而生成高效且高性能的代码。程序员应确信他们表达的矢量操作会可靠地紧密映射到相关的硬件矢量指令。同样的要求也适用于有能力的 ARM AArch64 架构,将其编译为由 Neon 支持的一系列硬件矢量指令。
-
优雅降级: 如果矢量计算无法在运行时完全表达为一系列硬件矢量指令,无论是因为架构不支持某些所需指令还是因为另一个 CPU 架构不受支持,那么矢量 API 实现应优雅降级并仍然正常工作。这可能包括在矢量计算无法充分编译为矢量硬件指令时向开发者发出警告。在没有矢量支持的平台上,优雅降级应生成与手动展开循环相媲美的代码,其中展开因子是所选矢量中的通道数。
非目标
-
增强 HotSpot 中的自动矢量化支持并不是目标。
-
让 HotSpot 支持 x64 和 AArch64 以外的 CPU 架构上的矢量硬件指令并不是目标。这种支持留待以后的 JEP 来实现。然而,重要的是要声明,正如目标中所表达的那样,API 不应排除此类实现。此外,执行的工作可能会自然地利用并扩展现有的 HotSpot 抽象以支持自动矢量化,从而使这样的任务更加容易。
-
在当前或未来的迭代中支持 C1 编译器并不是目标。我们期望 Graal 编译器能在未来的工作中得到支持。
-
支持由 Java
strictfp
关键字定义的严格浮点计算并不是目标。在浮点标量上执行的浮点运算结果可能与在浮点标量向量上执行的等效浮点运算不同。然而,这一目标并未排除表达或控制浮点向量计算所需精度或可再现性的选项。
动机
向量计算由一系列对向量的操作组成。一个向量包含一个(通常)固定序列的标量值,其中标量值对应于硬件定义的向量通道数量。对具有相同数量通道的两个向量应用二元操作时,会在每个通道上对标量值执行等效的标量操作,这些标量值分别来自每个向量中的对应位置。这通常被称为 单指令多数据(SIMD)。
向量操作表达了一定程度的并行性,这种并行性使得在单个 CPU 周期内可以完成更多工作,从而可能带来显著的性能提升。例如,给定两个向量,每个向量覆盖八个整数的序列(八个通道),那么这两个向量可以使用单个硬件指令相加。向量加法硬件指令在通常处理两个整数并执行一次整数加法所需的时间内,对十六个整数进行操作,完成八次整数加法。
HotSpot 支持 自动矢量化,其中标量操作被转换为超字(superword)操作,然后映射到矢量硬件指令。可转换的标量操作集是有限的,并且对代码结构的变化非常敏感。此外,可能仅使用了部分可用的矢量硬件指令,从而限制了生成代码的性能。
希望编写能可靠地转换为超字操作的标量操作的开发者需要理解 HotSpot 的自动矢量化支持及其局限性,以实现可靠且可持续的性能。
在某些情况下,开发人员可能无法编写可转换的标量操作。例如,HotSpot 不会转换用于计算数组哈希码的简单标量操作(参见 JDK 源代码中的 Arrays::hashCode
方法实现),也无法自动向量化代码来按字典顺序比较两个数组(这就是为什么添加了一个内在函数来进行字典顺序比较,参见 8033148)。
Vector API 旨在通过提供一种使用 HotSpot 中已有的矢量化支持来在 Java 中编写复杂矢量算法的机制,从而解决这些问题,但其用户模型使矢量化变得更加可预测和稳健。手写矢量循环可以表达高性能算法(例如矢量化的 hashCode
或专门的数组比较),而自动矢量化器可能永远无法优化这些算法。这种显式矢量化 API 可应用的领域众多,例如机器学习、线性代数、密码学、金融以及 JDK 本身的用法。
描述
向量将由抽象类 Vector<E>
表示。类型变量 E
对应于向量所涵盖的标量基本整数或浮点元素类型的装箱类型。向量还有一个 形状(shape),它定义了向量的大小(以位为单位)。当 HotSpot C2 编译器编译向量计算时,向量的形状将决定 Vector<E>
的实例如何映射到向量硬件寄存器(稍后会介绍从实例到 x64 向量寄存器的映射)。向量的长度(通道数或元素数)将是向量大小除以元素大小的结果。
所支持的元素类型 (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
。
注意: 我们认为这些简单的形状足够通用,可以在所有支持 Vector API 的平台上发挥作用。然而,随着我们在该 JEP 孵化期间对未来的平台进行实验,我们可能会进一步修改形状参数的设计。这类工作并不在该 JEP 的早期范围内,但这些可能性部分影响了形状在 Vector API 中的当前角色。请参见下面的未来工作部分。
元素类型和形状的组合决定了向量的种类,用 VectorSpecies<E>
表示。
Vector<E>
的实例是不可变的,它是一种基于值的类型,默认情况下保留对象标识不变量(有关这些不变量的放宽,请参见后文)。
对矢量的操作可以分为 lane-wise 和 cross-lane 两类。lane-wise 操作可以进一步分为一元操作、二元操作、三元操作和比较操作。cross-lane 操作可以分为排列操作、转换操作和归约操作。为了减少 API 的复杂性,我们将为每类操作定义集体方法,这些方法将接受一个操作符作为输入。支持的操作符是 Operator
类的实例,并在 VectorOperators
类中定义为静态最终字段。一些常见的操作(例如加法、乘法),称为全服务操作,将拥有专用方法,可以直接替代通用方法使用。
某些对向量的操作,例如按车道转换(lane-wise cast)和重新解释(reinterpret),可以说是天生具有改变形状的特性。在向量计算中包含改变形状的操作可能会对可移植性和性能产生意外影响。因此,在适用的情况下,API 将定义此类操作的额外一种形状不变版本。鼓励用户使用操作的形状不变版本编写形状不变的代码。此外,Javadoc 中将明确标注改变形状的操作。
Vector<E>
声明了一组适用于所有元素类型的基本向量操作方法。为了支持特定于元素类型的操作,Vector<E>
有六个抽象子类,每个支持的元素类型各一个:ByteVector
、ShortVector
、IntVector
、LongVector
、FloatVector
和 DoubleVector
。这些子类定义了额外的操作,这些操作与元素类型绑定,因为方法签名引用了元素类型(或等效的数组类型),例如归约操作(例如,将所有元素求和为标量值)或将向量元素存储到数组中。它们还定义了一些针对整数子类型的全功能操作,例如按位操作(例如,逻辑或),以及针对浮点类型的操作,例如数学运算(例如,超越函数如 pow())。
这些类通过为不同形状(大小)的 Vectors 定义的具体子类进一步扩展。
具体的子类是非公开的,因为没有必要提供针对类型和形状的特定操作。这将 API 的暴露面减少为关注点的总和,而非乘积。因此,具体的 Vector
类的实例无法直接构造。相反,实例是通过基类 Vector<E>
及其特定类型子类中定义的工厂方法获取的。这些方法以所需向量实例的种类作为输入。工厂方法提供了多种获取向量实例的方式,例如元素被初始化为默认值的向量实例(零向量),或从数组创建的向量,此外还提供了在不同类型的向量之间转换的标准支持(例如,类型转换)。
为了支持控制流,相关的矢量操作将可选地接受由公共抽象类 VectorMask<E>
表示的掩码。掩码中的每个元素(一个布尔值或位)对应于一个矢量通道。当掩码作为操作的输入时,它决定该操作是否应用于每个通道;如果通道的掩码位被设置(为真),则应用该操作。如果掩码位未被设置(为假),则会有替代行为。与矢量类似,VectorMask<E>
的实例是为每种元素类型和长度组合定义的(私有)具体子类的实例。在操作中使用的 VectorMask<E>
实例应与参与该操作的 Vector<E>
实例具有相同的类型和长度。比较操作会产生掩码,然后可以将这些掩码输入到其他操作中,以选择性地在某些通道上禁用该操作,从而模拟流控制。创建掩码的另一种方法是使用 VectorMask<E>
中的静态工厂方法。
我们预计,掩码在开发对形状通用的矢量计算中可能会发挥重要作用。(这一预期基于谓词寄存器(即掩码的等价物)在 ARM 可扩展矢量扩展 (Scalable Vector Extensions) 以及 Intel 的 AVX-512 中的核心重要性。)
示例
下面是一个对数组元素进行简单标量计算的例子:
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 实现等效向量计算的一种显式方法如下:
// Example 1
static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_256;
void vectorComputation(float[] a, float[] b, float[] c) {
for (int i = 0; i < a.length; i += SPECIES.length()) {
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);
}
}
在此示例中,从 FloatVector
获取浮点数的 256 位宽向量的种类。该种类存储在一个 static final
字段中,因此运行时编译器会将该字段的值视为常量,从而能够更好地优化向量计算。
矢量计算的主循环内核以矢量长度(即种类长度)为步幅迭代数组。静态方法 fromArray()
从数组 a
和 b
中加载给定种类的 float
矢量到相应的索引。然后流畅地执行操作,最后将结果存储到数组 c
中。
我们使用由 indexInRange()
生成的掩码,以防止读取或写入超出数组长度。前 floor(a.length / SPECIES.length())
次迭代中,掩码的所有通道都会被设置。只有最后一次迭代时,如果 a.length
不是 SPECIES.length()
的倍数,掩码才会设置前 a.length % SPECIES.length()
个通道。
由于在所有迭代中都使用了掩码,因此对于较大的数组长度,上述实现可能无法达到最佳性能。可以按如下方式在不使用掩码的情况下实现相同的计算:
// Example 2
static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_256;
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;
}
}
尾部 元素的长度小于物种长度,这些元素在向量计算之后使用标量计算进行处理。处理尾部元素的另一种方法是使用单一的掩码向量计算。
在对大型数组进行操作时,上面的实现达到了最佳性能。
对于第二个示例,HotSpot 编译器应在支持 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
这是使用 Vector API 和实现的原型(Project Panama 开发仓库的 vectorIntrinsics
分支)对测试示例代码进行 JMH 微基准测试的实际输出。这显示了 C2 生成的机器代码的热点区域。可以清楚地看到向量寄存器和向量硬件指令的转换。(为了使转换更加清晰,循环展开被禁用,否则 HotSpot 应该能够使用现有的 C2 循环优化技术进行展开。)所有 Java 对象分配都被省略了。
支持更复杂的非平凡矢量计算,并将其清晰地转换为生成的机器代码,这是一个重要的目标。
然而,这个特定的向量计算存在一些问题:
-
循环被硬编码为具体的矢量形状,因此计算无法动态适应架构支持的最大形状,该形状可能小于或大于 256 位。所以代码的可移植性较差,性能也可能较低。
-
虽然这里的循环上限计算很简单,但这是编程错误的常见来源。
-
最后需要一个标量循环,这会导致代码重复。
我们将在本 JEP 中解决前两个问题。可以获得一个形状对当前架构最优的首选物种,然后可以用通用形状编写向量计算,并且物种上的方法可以减少数组长度,例如:
static final VectorSpecies<Float> SPECIES = FloatVector.SPECIES_PREFERRED;
void vectorComputation(float[] a, float[] b, float[] c,
VectorSpecies<Float> species) {
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;
}
}
vectorComputation(a, b, c, SPECIES);
第三个问题不会被这个 JEP 完全解决,将是未来工作的一个主题。如第一个例子所示,你可以使用掩码来实现没有尾部处理的向量计算。我们预计这种带掩码的循环将适用于一系列架构,包括 x64 和 ARM,但需要额外的运行时编译器支持以生成最高效的代码。尽管对带掩码循环的工作很重要,但超出了此 JEP 的范围。
HotSpot C2 编译器详细信息
为了实现此 JEP 的目标,Vector API 有两种实现。第一种是用 Java 实现操作,因此它是功能性的,但不是最优的。第二种是针对 HotSpot C2 编译器的内在实现,对 Vector API 类型的操作进行特殊处理。这允许在存在架构支持和转换实现的情况下,正确地转换为硬件寄存器和指令。
为了避免 C2 中添加的内联函数激增,将定义一组与操作类型(如二元、一元、比较等)相对应的内联函数,其中传递的常量参数用于描述操作的具体细节。为了支持 API 各个部分的内联化,大约需要二十个新的内联函数。
Vector
实例是基于值的,也就是说,在道德上应避免使用对身份敏感的操作。此外,尽管向量实例在抽象上由通道中的元素组成,但这些元素并不会被 C2 标量化。向量值被视为一个整体单元,类似于 int
或 long
,映射到适当大小的硬件向量寄存器。内联类型将需要一些相关的增强功能,以确保向量值被视为一个整体单元。
在内联类型可用之前,C2 将对 Vector
实例进行特殊处理,以克服逃逸分析的局限性并避免装箱。因此,应避免对向量进行依赖标识的操作。
未来工作
一旦准备就绪,Vector API 将从值类型中显著受益(参见 Project Valhalla)。Vector<E>
的实例可以是值,其具体类为内联类型。这将使优化和表达向量计算变得更加容易。对于特定类型的 Vector<E>
子类型(如 IntVector
),在支持内联类型的泛型特化和类型特定方法声明的情况下,可能不再需要。
因此,正如上面所提到的,Vector API 的未来版本将会使用内联类型和增强的泛型。由此,我们将随着内联类型的可用性,在多个 JDK 版本中孵化该 API 并进行适配。
当 JEP 370 Foreign-Memory Access API 从孵化阶段的 API 转变为正式 API 后,我们将增强 API,以利用其特性来加载和存储向量。此外,用于描述向量种类的内存布局可能会很有用,例如,可以用于跨步遍历由元素组成的内存段。
我们预计将以以下方式加强实施:
-
包括对矢量化超越运算(例如对数和三角函数)的支持,
-
改进包含矢量化代码的循环优化,
-
在支持的平台上优化带掩码的矢量运算,以及
-
针对大型矢量尺寸进行调整(例如,ARM SVE 所支持的)。
当我们对实现进行增量改进时,性能工作将持续进行。
替代方案
HotSpot 的自动向量化是一种替代方法,但需要进行大量工作。此外,与使用 Vector API 相比,它可能仍然很脆弱且有局限性,因为对于复杂控制流进行自动向量化非常难以实现。
总体而言,即使经过数十年的研究(特别是针对 FORTRAN 和 C 数组循环),标量代码的自动向量化似乎并不是优化用户随意编写的循环的可靠策略,除非用户异常小心地关注编译器准备自动向量化的确切循环的那些不成文约定。写出一个无法自动向量化的循环太容易了,其原因只有优化器能检测到,而人类读者却无法察觉。多年以来,即使在 HotSpot 中,关于自动向量化的研究工作为我们留下了许多仅在特殊场合下才有效的优化机制。我们希望能够更频繁地利用这些机制!
测试
我们将开发组合单元测试,以确保覆盖所有操作、所有支持的类型和形状,以及各种数据集。
我们还将开发性能测试,以确保达到性能目标,并且矢量计算能够高效映射到矢量硬件指令。这可能包括 JMH 微基准测试,但还需要更多实际有用的算法示例。此类测试最初可能会存放在特定于项目的存储库中。考虑到测试的比例以及它们的生成方式,在集成到主存储库之前可能需要进行整理。
作为性能测试的备份,我们可以创建白盒测试,以强制 JIT 向我们报告矢量 API 源代码确实触发了矢量化。
风险与假设
存在 API 偏向于 x64 架构上支持的 SIMD 功能的风险,但通过支持 AArch64 可以减轻这种风险。这主要适用于显式固定的受支持形状集,这些形状不利于以形状通用的方式编写算法。我们认为 Vector API 的大多数其他操作偏向于可移植算法。为了降低该风险,我们会将其他架构考虑在内,特别是 ARM Scalar Vector Extension 架构,其编程模型会动态调整为硬件支持的单一固定形状。我们欢迎并鼓励从事 HotSpot 的 ARM 特定领域的 OpenJDK 贡献者参与此项工作。
Vector API 使用包装类型(例如 Integer
)作为基本类型(例如 int
)的代理。这一决策是由 Java 泛型当前的限制所决定的,因为 Java 泛型对基本类型并不友好。当 Project Vahalla 最终引入更强大的泛型时,当前的决策可能会显得笨拙,并可能需要更改。我们假设这样的更改将可以在不引起过度向后不兼容的情况下实现。