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:Foreign Function & Memory (FFM) API(预览版) 中定义的
MemorySegment
加载和存储向量。FFM API 已足够成熟,我们可以放心地将此依赖项添加到 Vector API 中。我们将移除操作byte[]
和ByteBuffer
的等效 API 点,因为可以为它们获取MemorySegment
。使用MemorySegment
将能够创建与向量字节长度对齐的 超对齐 内存区域。在某些架构上,这种对齐在加载和存储向量时能够实现更优的性能。添加两个新的跨通道向量操作,compress(压缩)及其逆操作 expand(展开),以及一个互补的向量掩码 compress(压缩)操作。compress(压缩)向量操作将按掩码选择的源向量的通道映射到目标向量中,按通道顺序排列;expand(展开)操作执行相反的操作。compress(压缩)操作在过滤查询结果时非常有用;例如,它可以从源向量中选择所有非零整数元素,并将它们按源通道顺序连续写入目标向量。
扩展支持的按位整数通道操作集,以包括:
- 计算一位的数量,
- 计算前导零位的数量,
- 计算尾随零位的数量,
- 反转位的顺序,
- 反转字节的顺序,以及
- 压缩和展开位。
后两个操作类似于跨通道的 compress(压缩)和 expand(展开)操作,但它们映射位而不是整个通道。(有关按位 compress(压缩)和 expand(展开)的更多信息,请参阅 Henry S. Warren 所著的《Hacker's Delight》一书的第 7-4 节和第 7-5 节,由 Addison-Wesley 于 2013 年出版。)
与此 JEP 并行但独立的是,我们将向装箱的原始
Integer
和Long
类型添加按位标量 compress(压缩)和 expand(展开)操作的方法,这些方法独立使用也很有用。我们将根据这些新的标量方法来指定按位 compress(压缩)和 expand(展开)向量操作。
目标
清晰简洁的 API — API 应能够清晰简洁地表达由循环内组合的一系列向量操作构成的广泛向量计算,并可能包含控制流。应能够表达与向量大小或每个向量的通道数无关的通用计算,从而使此类计算能够在支持不同向量大小的硬件之间移植。
平台无关性 — API 应与 CPU 架构无关,能够在支持向量指令的多个架构上实现。与 Java API 中的常规情况一样,当平台优化与可移植性发生冲突时,我们将偏向于使 API 具有可移植性,即使这会导致某些平台特定的习惯用法无法在可移植代码中表达。
在 x64 和 AArch64 架构上的可靠运行时编译和性能 — 在功能强大的 x64 架构上,Java 运行时(特别是 HotSpot C2 编译器)应将向量操作编译为相应的高效且高性能的向量指令,如 流 SIMD 扩展(SSE)和 高级向量扩展(AVX)所支持的指令。开发人员应确信他们所表达的向量操作将可靠地紧密映射到相关向量指令上。在功能强大的 ARM AArch64 架构上,C2 同样会将向量操作编译为 NEON 和 SVE 所支持的向量指令。
优雅降级 — 有时,向量计算无法完全在运行时表达为一系列向量指令,可能是因为架构不支持某些所需的指令。在这种情况下,Vector API 实现应优雅降级并仍能工作。如果向量计算无法有效编译为向量指令,则可能需要发出警告。在没有向量的平台上,优雅降级将产生与手动展开循环相竞争的代码,其中展开因子是所选向量的通道数。
非目标
不旨在增强 HotSpot 中现有的自动向量化算法。
不旨在支持除 x64 和 AArch64 之外的 CPU 架构上的向量指令。但重要的是要指出,如目标中所述,API 不得排除此类实现。
不旨在支持 C1 编译器。
不旨在保证支持 Java 平台对标量操作所要求的严格浮点计算。在浮点标量上执行的浮点运算的结果可能与在浮点标量向量上执行的等效浮点运算的结果不同。任何偏差都将明确记录。此非目标不排除表达或控制浮点向量计算所需精度或可重复性的选项。
动机
向量计算由一系列向量操作组成。向量由(通常是)固定顺序的标量值组成,其中标量值对应于硬件定义的向量通道数。对具有相同通道数的两个向量应用二元操作时,将对每个通道执行从每个向量中对应两个标量值的等效标量操作。这通常被称为 单指令多数据(SIMD)。
向量操作表达了一定程度的并行性,使得可以在单个 CPU 周期内执行更多工作,从而可以显著提高性能。例如,给定两个向量,每个向量包含八个整数的序列(即八个通道),可以使用单个硬件指令将这两个向量相加。向量加法指令在通常对两个整数执行一次整数加法所需的时间内,对十六个整数进行操作,执行八次整数加法。
HotSpot 已经支持 自动向量化,它将标量操作转换为超字操作,然后将这些操作映射到向量指令。可转换的标量操作集是有限的,并且在代码形状发生变化时也很脆弱。此外,可能只会利用可用向量指令的一个子集,从而限制了生成代码的性能。
如今,希望编写可可靠转换为超字操作的标量操作的开发人员需要了解 HotSpot 的自动向量化算法及其局限性,以实现可靠且可持续的性能。在某些情况下,可能无法编写可转换的标量操作。例如,HotSpot 不会转换用于计算数组哈希码的简单标量操作(即 Arrays::hashCode 方法),也无法自动向量化代码以按字典顺序比较两个数组(因此我们 为字典顺序比较添加了一个内部函数)。
Vector 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 位的向量大小,以及 max 位。一个 512 位的形状可以将 byte
打包成 64 个通道,或者将 int
打包成 16 个通道,并且具有这种形状的向量可以一次操作 64 个 byte
或 16 个 int
。max-bits 形状支持当前架构的最大向量大小。这支持 ARM SVE 平台,该平台实现可以支持从 128 位到 2048 位的任何固定大小,增量为 128 位。
我们认为这些简单的形状足够通用,可以在所有相关平台上使用。然而,在 API 的孵化期间,随着我们对未来平台的实验,我们可能会进一步修改形状参数的设计。这种工作不在本项目的早期范围内,但这些可能性部分地说明了形状在 Vector 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>
表示的掩码。掩码中的每个元素都是一个布尔值,对应于一个向量通道。掩码选择应用操作的通道:如果通道的掩码元素为真,则应用该操作;如果掩码为假,则执行一些替代操作。
与向量类似,VectorMask<E>
的实例是为每个元素类型和长度组合定义的非公开具体子类的实例。在操作中使用的 VectorMask<E>
实例应具有与操作中涉及的向量实例相同的类型和长度。向量比较操作产生掩码,这些掩码随后可以用作其他操作的输入,以选择性地操作某些通道,从而模拟流程控制。还可以使用 VectorMask<E>
类中的静态工厂方法来创建掩码。
我们预计掩码在开发形状通用的向量计算中将发挥重要作用。这一预期基于谓词寄存器(掩码的等效物)在 ARM 可扩展向量扩展和英特尔 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
获取一个首选的 VectorSpecies
,其形状对当前架构而言是最优的。我们将其存储在 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
这是使用 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 Valhalla 在 JEP 401(基本对象) 中所提议的那样。在此期间,Vector<E>
及其子类被视为 基于值的类,因此应避免在其实例上进行身份敏感操作。尽管向量实例在逻辑上由通道中的元素组成,但这些元素不会被 C2 标量化——向量的值被视为一个整体单元,类似于 int
或 long
,它映射到适当大小的向量寄存器。为了克服逃逸分析的限制并避免装箱,C2 对向量实例进行了特殊处理。
Intel SVML 内部函数用于超越运算
Vector API 支持对浮点向量的逐元素超越和三角函数运算。在 x64 上,我们利用 Intel 短向量数学库(SVML)为这些操作提供优化的内部实现。这些内部操作的数值属性与 java.lang.Math
中定义的相应标量操作相同。
SVML 操作的汇编源文件位于 jdk.incubator.vector
模块的源代码中,位于特定于操作系统的目录下。JDK 构建过程将这些源文件针对目标操作系统编译成 SVML 特定的共享库。该库相当大,重量接近 1MB。如果通过 jlink
构建的 JDK 映像省略了 jdk.incubator.vector
模块,则 SVML 库不会被复制到映像中。
目前,该实现仅支持 Linux 和 Windows。我们稍后会考虑 macOS 支持,因为提供带有必需指令的汇编源文件需要相当多的工作。
HotSpot 运行时将尝试加载 SVML 库,如果库存在,则将 SVML 库中的操作绑定到命名的存根例程。C2 编译器根据操作和向量种类(即元素类型和形状)生成调用相应存根例程的代码。
未来,如果 Project Panama 扩展其对本地调用约定的支持以支持向量值,那么 Vector API 实现可能能够从外部源加载 SVML 库。如果这种方法没有性能影响,那么将不再需要以源代码形式包含 SVML 并将其构建到 JDK 中。在此之前,鉴于潜在的性能提升,我们认为上述方法是可接受的。
未来工作
如上所述,我们最终期望将向量类声明为 基本类。此外,我们期望利用 Project Valhalla 的基本类泛型特化,以便
Vector<E>
的实例可以是具体类型为基本类型的基本值。这将使优化和表达向量计算变得更容易。一旦我们在基本类上实现了泛型特化,可能就不再需要Vector<E>
的特定类型子类型(如IntVector
)。我们打算在多个版本中对 API 进行孵化,并根据基本类和相关设施的可用性进行调整。我们计划改进实现,以优化包含向量化代码的循环,并随着时间的推移逐步提高性能。
我们还计划增强组合单元测试,以断言 C2 生成向量硬件指令。当前的单元测试在未经验证的情况下假设重复执行足以使 C2 生成向量硬件指令。我们将探索使用 C2 的 IR 测试框架 来跨平台断言 IR 图中存在向量节点(例如,使用 正则表达式匹配)。如果这种方法有问题,我们可能会探索使用非产品
-XX:+TraceNewVectors
标志来打印向量节点的基本方法。我们将评估合成向量形状的定义,以便更好地控制循环展开和矩阵操作,并考虑对排序和解析算法的适当支持。(有关更多详细信息,请参阅此演示文稿。)
替代方案
HotSpot 的自动向量化是一种替代方法,但需要大量工作。此外,与 Vector API 相比,它仍然脆弱且受限,因为具有复杂控制流的自动向量化很难实现。
一般来说,即使经过数十年的研究——特别是针对 FORTRAN 和 C 数组循环——似乎标量代码的自动向量化并不是优化用户编写的临时循环的可靠策略,除非用户异常小心地关注关于编译器准备自动向量化哪些循环的未成文约定。编写一个无法自动向量化的循环太容易了,而人类读者却无法察觉其原因。即使在 HotSpot 上,多年来对自动向量化的研究也留下了许多仅在特殊情况下有效的优化机制。我们希望更频繁地使用这些机制!
测试
我们将开发组合单元测试,以确保覆盖所有操作、所有支持的类型和形状,以及各种数据集。
我们还将开发性能测试,以确保达到性能目标,并将向量计算有效地映射到向量指令上。这很可能包括 JMH 微基准测试,但也需要更实际的实用算法示例。这些测试最初可能位于特定项目的存储库中。鉴于测试的比例和生成方式,在集成到主存储库之前可能需要进行整理。
风险和假设
存在 API 偏向 x64 架构上支持的 SIMD 功能的风险,但通过支持 AArch64 可以缓解这一风险。这主要适用于明确固定的支持形状集合,这些形状偏向于以形状通用的方式编码算法。我们认为 Vector API 的大多数其他操作都偏向于可移植算法。为了缓解这一风险,我们将考虑其他架构,特别是 ARM 标量向量扩展架构,其编程模型会根据硬件支持的单一固定形状动态调整。我们欢迎并鼓励在 HotSpot 的 ARM 特定区域工作的 OpenJDK 贡献者参与这项工作。
Vector API 使用装箱类型(如
Integer
)作为原始类型(如int
)的代理。这一决定是由 Java 泛型的当前限制所迫,这些限制对原始类型不友好。当 Project Valhalla 最终引入更强大的泛型时,当前的决策将显得笨拙,并可能需要更改。我们假设这些更改将可能实现,而不会导致过度的向后不兼容。