Skip to content

JEP 414: Vector API (Second Incubator) | 向量 API(第二个孵化器)

摘要

引入一个 API 来表达向量计算,该计算能在运行时可靠地编译成受支持 CPU 架构上的最优向量指令,从而实现比等效标量计算更优的性能。

历史

Vector API 由 JEP 338 提出,并作为 孵化 API 集成到 Java 16 中。我们在此提议纳入根据反馈进行的改进,以及性能提升和其他重要的实现改进。我们包括以下显著变更:

  • 对 API 的增强,以支持字符操作,如 UTF-8 字符解码。具体而言,我们添加了方法用于在 short 向量和 char 数组之间复制字符,以及新的向量比较运算符,用于与整数向量进行无符号比较。

  • 对 API 的增强,用于在 byte 向量和 boolean 数组之间进行转换。

  • 使用英特尔的 Short Vector Math Library (SVML),为 x64 平台上的超越函数和三角函数的逐道(lanewise)操作提供内建支持。

  • 对 Intel x64 和 ARM NEON 实现的通用性能增强。

目标

  • 清晰简洁的 API — 该 API 应能够清晰简洁地表达由循环内和可能带有控制流的向量操作序列组成的广泛向量计算。应能够表达一种相对于向量大小或每个向量的通道数具有通用性的计算,从而使这些计算能够在支持不同向量大小的硬件上实现可移植性。

  • 与平台无关 — 该 API 应与 CPU 架构无关,能够在支持向量指令的多个架构上实现。与 Java API 中常见的平台优化与可移植性冲突的情况一样,这里将偏向于使 API 具有可移植性,即使这会导致某些特定于平台的惯用法无法在可移植代码中表达。

  • 在 x64 和 AArch64 架构上的可靠运行时编译和性能 — 在功能强大的 x64 架构上,Java 运行时(特别是 HotSpot C2 编译器)应将向量操作编译为相应的高效向量指令,如 Streaming SIMD Extensions(SSE)和 Advanced Vector Extensions(AVX)所支持的指令。开发人员应确信他们所表达的向量操作将可靠地映射到相关的向量指令上。在功能强大的 ARM AArch64 架构上,C2 编译器也将类似地将向量操作编译为 NEON 所支持的向量指令。

  • 优雅降级 — 有时,向量计算无法在运行时完全表示为一系列向量指令,可能是因为架构不支持某些必需的指令。在这种情况下,Vector API 的实现应优雅降级并仍能正常工作。如果向量计算无法高效地编译为向量指令,则可能需要发出警告。在没有向量的平台上,优雅降级将产生与手动展开的循环相媲美的代码,其中展开因子是所选向量中的通道数。

非目标

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

  • 不以支持除 x64 和 AArch64 之外的 CPU 架构上的向量指令为目标。但重要的是要指出,如目标中所述,API 不得排除此类实现。

  • 不以支持 C1 编译器为目标。

  • 不以支持 Java strictfp 关键字定义的严格浮点计算为目标。对浮点标量执行的浮点运算的结果可能与对浮点标量向量执行的等效浮点运算的结果不同。但是,此目标不排除表达或控制浮点向量计算所需精度或可重复性的选项。

动机

向量计算由一系列向量操作组成。向量通常由(通常是)固定序列的标量值组成,其中标量值对应于硬件定义的向量通道数。对两个具有相同通道数的向量执行的二进制操作将针对每个通道,从每个向量中对应的两个标量值上应用等效的标量操作。这通常被称为 单指令多数据(SIMD)。

向量操作表达了一定程度的并行性,这使得单个 CPU 周期内可以执行更多工作,从而可以显著提高性能。例如,给定两个向量,每个向量包含八个整数的序列(即八个通道),则可以使用单个硬件指令将这两个向量相加。向量加法指令在十六个整数上操作,执行八个整数加法,而在通常情况下,它只能操作两个整数,执行一个整数加法。

HotSpot 已经支持 自动向量化,它将标量操作转换为超字操作,然后将这些操作映射到向量指令。可转换的标量操作集是有限的,并且对于代码形状的变化也很脆弱。此外,可能只会利用可用向量指令的一个子集,从而限制了生成代码的性能。

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

Vector API 旨在通过提供一种在 Java 中编写复杂向量算法的方法来改善这种情况,该方法使用现有的 HotSpot 自动向量化器,但具有一种用户模型,使向量化变得更加可预测和健壮。手动编码的向量循环可以表达高性能算法,如向量化的 hashCode 或专业的数组比较,而自动向量化器可能永远不会对其进行优化。许多领域都可以从这个显式的 Vector API 中受益,包括机器学习、线性代数、密码学、金融以及 JDK 本身中的代码。

描述

向量由抽象类 Vector<E> 表示。类型变量 E 被实例化为向量所覆盖的标量原始整数或浮点元素类型的装箱类型。向量还具有一个 形状,它定义了向量的位数大小。向量的形状决定了当 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 类的静态最终字段中定义。为了方便起见,我们为一些常见的 全服务 操作(如加法和乘法)定义了专用方法,这些方法可以代替通用方法使用。

向量上的某些操作,如转换和重新解释,本质上是 改变形状的;即,它们产生的向量的形状与输入的形状不同。在向量计算中,改变形状的操作可能会对可移植性和性能产生负面影响。因此,当适用时,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)和英特尔 AVX-512 中的核心重要性。

示例

以下是一个简单的标量计算在数组元素上的例子:

java
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 的等效向量计算示例:

java
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 字段中,以便运行时编译器将该值视为常量,从而能够更好地优化向量计算。然后,主循环以向量长度的步长(即 VectorSpecies 的长度)遍历输入数组。它会在数组 ab 的相应索引处加载给定 VectorSpeciesfloat 向量,流畅地执行算术运算,然后将结果存储到数组 c 中。如果在最后一次迭代后仍有数组元素剩余,则使用普通的标量循环计算这些“尾部”元素的结果。

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

txt
  0.43%  / │  0x0000000113d43890: vmovdqu 0x10(%r8,%rbx,4),%ymm0  ; 将从数组 a 中加载的数据移动到向量寄存器 ymm0
  7.38%  │ │  0x0000000113d43897: vmovdqu 0x10(%r10,%rbx,4),%ymm1  ; 将从数组 b 中加载的数据移动到向量寄存器 ymm1
  8.70%  │ │  0x0000000113d4389e: vmulps %ymm0,%ymm0,%ymm0  ; ymm0 中的元素自乘
  5.60%  │ │  0x0000000113d438a2: vmulps %ymm1,%ymm1,%ymm1  ; ymm1 中的元素自乘
 13.16%  │ │  0x0000000113d438a6: vaddps %ymm0,%ymm1,%ymm0  ; 将 ymm0 和 ymm1 的结果相加,结果存回 ymm0
 21.86%  │ │  0x0000000113d438aa: vxorps -0x7ad76b2(%rip),%ymm0,%ymm0  ; 使用特定的掩码对 ymm0 中的结果进行异或操作,实现取反
  7.66%  │ │  0x0000000113d438b2: vmovdqu %ymm0,0x10(%r9,%rbx,4)  ; 将结果存储回数组 c
 26.20%  │ │  0x0000000113d438b9: add    $0x8,%ebx  ; 更新索引 ebx,以指向下一个向量块
  6.44%  │ │  0x0000000113d438bc: cmp    %r11d,%ebx  ; 比较 ebx 与数组长度,以检查是否已完成所有迭代
         \ │  0x0000000113d438bf: jl     0x0000000113d43890  ; 如果没有完成,则跳回循环的开始

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

运行时编译

Vector API 有两种实现方式。第一种在 Java 中实现操作,因此它是可用的但并非最优。第二种为 HotSpot C2 运行时编译器定义了内联向量操作,以便在可用时,将向量计算编译为适当的硬件寄存器和向量指令。

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

我们期望最终将向量类声明为基本类,如 Valhalla 项目JEP 401(基本对象) 中所提出的那样。同时,Vector<E> 及其子类被视为 基于值的类,因此应避免在其实例上执行对身份敏感的操作。尽管向量实例在逻辑上由通道中的元素组成,但这些元素不会被 C2 标量化——向量的值被视为一个整体单元,类似于 intlong,它映射到适当大小的向量寄存器。C2 对向量实例进行特殊处理,以克服逃逸分析的限制并避免装箱。

Intel SVML 内联函数用于超越运算

Vector API 支持对浮点向量的逐元素超越和三角函数运算。在 x64 平台上,我们利用 Intel Short Vector Math Library (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 扩展其对本地调用约定的支持以支持向量值,那么 Vector API 实现可能会从外部源加载 SVML 库。如果这种方法没有性能影响,那么将不再需要以源代码形式包含 SVML 并将其构建到 JDK 中。在此之前,鉴于潜在的性能提升,我们认为上述方法是可接受的。

未来工作

  • 如上所述,我们期望最终将向量类声明为 基本类。此外,我们期望利用 Valhalla 项目对基本类的泛型特化,以便 Vector<E> 的实例可以是具体类型为基本类型的基本值。这将使得优化和表达向量计算变得更加容易。一旦我们在基本类上实现了泛型特化,可能就不再需要 Vector<E> 针对特定类型的子类型(如 IntVector)。我们打算在多个版本中孵化该 API,并根据基本类和相关设施的可用性进行调整。

  • 我们希望在支持硬件屏蔽的架构上改进接受屏蔽的向量操作的性能。如果屏蔽操作更高效,那么 上面的示例 可以编写得更简单,无需使用标量循环来处理尾部元素,同时仍能达到最佳性能:

    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);
        }
    }
  • 我们打算在 JEP 412(外部函数和内存 API) 从孵化阶段过渡后,增强该 API 以使用它来加载和存储向量。描述向量类型的内存布局可能会很有用,例如,可以在由向量元素组成的内存段上进行跨步操作。

  • 我们预计将对实现进行增强,以改进包含向量化代码的循环的优化,支持 ARM SVE 平台,并随着时间的推移逐步提高性能。

备选方案

HotSpot 的自动向量化是一种替代方法,但将需要大量的工作。此外,与 Vector API 相比,它仍然脆弱且有限,因为具有复杂控制流的自动向量化很难实现。

总的来说,即使在几十年的研究之后——特别是针对 FORTRAN 和 C 语言数组循环——似乎标量代码的自动向量化并不是优化用户编写的特定循环的可靠策略,除非用户非常仔细地关注哪些循环是编译器准备自动向量化的未成文约定。很容易编写一个由于某种人类读者无法检测的原因而无法自动向量化的循环。即使在 HotSpot 上进行了多年的自动向量化工作,我们也留下了大量仅在特殊场合下有效的优化机制。我们希望更频繁地使用这些机制!

测试

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

我们还将开发性能测试,以确保达到性能目标,并且向量计算能够高效地映射到向量指令上。这很可能包括 JMH 微基准测试,但也需要更现实的实用算法示例。这些测试最初可能存放在特定项目的存储库中。鉴于测试的比例和生成方式,在集成到主存储库之前,很可能需要进行整理。

作为性能测试的补充,我们可能会创建白盒测试,以强制 JIT 向我们报告 Vector API 源代码确实触发了向量化。

风险和假设

  • 存在 API 偏向于 x64 架构上支持的 SIMD 功能的风险,但通过使用 AArch64 支持可以减轻这一风险。这主要适用于明确固定的支持形状集合,这些形状偏向于以形状通用的方式编码算法。我们认为 Vector API 的大多数其他操作都偏向于可移植算法。为了减轻这一风险,我们将考虑其他架构,特别是 ARM 标量向量扩展架构,其编程模型会根据硬件支持的单一固定形状动态调整。我们欢迎并鼓励 OpenJDK 贡献者参与 HotSpot 中 ARM 特定区域的工作,并加入这一努力。

  • Vector API 使用装箱类型(如 Integer)作为基本类型(如 int)的代理。这一决定是由于当前 Java 泛型的局限性所迫,因为泛型对基本类型不友好。当 Valhalla 项目最终引入更强大的泛型时,当前的决策可能会显得笨拙,并可能需要更改。我们假设这些更改将可以在不造成过多向后不兼容性的情况下进行。