Skip to content

JEP 338: Vector API (Incubator) | 向量 API(孵化器)

摘要

提供一个 孵化器模块 的初步迭代版本 jdk.incubator.vector,用于表达向量计算,这些计算在支持的 CPU 架构上可靠地编译为最优的向量硬件指令,从而实现比等效的标量计算更高的性能。

目标

  • 清晰简洁的 API:API 应能够清晰、简洁地表达一系列由向量操作组成的向量计算,这些向量操作通常组合在循环中,并可能包含控制流。应能够表达与向量大小(或每个向量的通道数)无关的通用计算,从而使这些计算能够在支持不同向量大小的硬件之间移植(如下一个目标所述)。

  • 平台无关:API 应是架构无关的,能够在支持向量硬件指令的多个 CPU 架构上支持运行时实现。在 Java API 中,当平台优化和可移植性发生冲突时,通常会偏向于使 Vector API 可移植,即使某些特定于平台的习惯用法无法直接在可移植代码中直接表达。下一个关于 x64 和 AArch64 性能的目标是 Java 所支持的所有平台上适当的性能目标的代表。ARM 可伸缩向量扩展(SVE)在这方面特别值得关注,以确保 API 能够支持这种架构,尽管在撰写本文时还没有已知的生产硬件实现。

  • 在 x64 和 AArch64 架构上可靠的运行时编译和性能:Java 运行时(特别是 HotSpot C2 编译器)应在支持的 x64 架构上,将一系列向量操作编译为相应的向量硬件指令序列,如 Streaming SIMD Extensions(SSE)和 Advanced Vector Extensions(AVX)扩展所支持的指令,从而生成高效且性能优越的代码。程序员应确信他们所表达的向量操作将可靠地映射到相关的硬件向量指令。同样的,这也应适用于能够编译到 Neon 支持的向量硬件指令序列的 ARM AArch64 架构。

  • 优雅降级:如果向量计算无法在运行时完全表达为一系列硬件向量指令序列,可能是因为架构不支持某些必需的指令,或者因为不支持其他 CPU 架构,那么 Vector API 实现应优雅地降级并仍然能够工作。这可能包括在向量计算无法充分编译为向量硬件指令时向开发者发出警告。在没有向量的平台上,优雅降级将产生与手动展开循环竞争的代码,其中展开因子是所选向量中的通道数。

非目标

  • 不旨在增强 HotSpot 中的自动矢量化支持。

  • 不旨在让 HotSpot 支持除 x64 和 AArch64 以外的 CPU 架构上的向量硬件指令。此类支持将留待后续的 JEP 处理。但是,重要的是要声明,如目标所述,API 不应排除此类实现。此外,进行的工作可能会自然地利用和扩展 HotSpot 中现有的自动矢量化向量支持抽象,使此类任务变得更容易。

  • 不旨在在此迭代或未来的迭代中支持 C1 编译器。我们预计在未来的工作中将支持 Graal 编译器。

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

动机

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

向量操作表达了一定程度的并行性,使得单个 CPU 周期内可以执行更多的工作,因此可以显著提高性能。例如,给定两个向量,每个向量覆盖八个整数的序列(八个通道),则可以使用单个硬件指令将这两个向量相加。向量加法硬件指令在十六个整数上操作,执行八个整数加法,而通常这需要两个整数上执行一个整数加法所需的时间。

HotSpot 支持 自动矢量化,其中标量操作被转换为超字操作,然后映射到向量硬件指令。可转换的标量操作集是有限的,并且容易受代码形状变化的影响。此外,可能仅利用可用向量硬件指令的子集,从而限制了生成代码的性能。

想要编写可靠地转换为超字操作的标量操作的开发者需要了解 HotSpot 的自动矢量化支持及其限制,以实现可靠且可持续的性能。

在某些情况下,开发者可能无法编写可转换的标量操作。例如,HotSpot 不会转换计算数组哈希码(参见 JDK 源代码中的 Arrays::hashCode 方法实现)的简单标量操作,也无法自动矢量化代码以按字典顺序比较两个数组(这就是为什么添加了一个内联函数来执行字典顺序比较,请参阅 8033148)。

Vector API 旨在通过提供在 Java 中编写复杂向量算法的机制来解决这些问题,使用 HotSpot 中预先存在的矢量化支持,但用户模型使矢量化变得更加可预测和健壮。手动编码的向量循环可以表达高性能算法(如矢量化 hashCode 或特定的数组比较),这些算法可能永远不会被自动矢量化器优化。在机器学习、线性代数、密码学、金融以及 JDK 本身内部的用途等许多领域,这种显式矢量化 API 可能都是适用的。

描述

向量将由抽象类 Vector<E> 表示。类型变量 E 对应于向量所涵盖的标量原始整数或浮点元素类型的装箱类型。向量还有一个“形状”,它定义了向量的大小(以位为单位)。向量的形状将决定当 HotSpot C2 编译器编译向量计算时,Vector<E> 的实例如何映射到向量硬件寄存器(稍后将介绍从实例到 x64 向量寄存器的映射)。向量的长度(通道数或元素数)将是向量大小除以元素大小。

支持的元素类型(E)集将是 ByteShortIntegerLongFloatDouble,分别对应于标量原始类型 byteshortintlongfloatdouble

支持的形状集将对应于 64、128、256 和 512 位的向量大小。对应于 512 位大小的形状可以将 byte 打包成 64 个通道或将 int 打包成 16 个通道,具有这种形状的向量可以一次处理 64 个 byte 或 16 个 int

注意: 我们认为这些简单的形状足够通用,可以在支持 Vector API 的所有平台上使用。但是,随着我们在本 JEP 孵化期间对未来平台进行实验,我们可能会进一步修改形状参数的设计。此类工作不在本 JEP 的初期范围内,但这些可能性部分地决定了 Vector API 中形状的当前角色。请参阅下面的未来工作部分。

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

Vector<E> 的实例是不可变的,并且是基于值的类型,默认情况下保留对象身份的不变性(稍后将放宽这些不变性)。

向量操作可以分为逐通道操作和跨通道操作。逐通道操作可以进一步分为一元、二元、三元和比较操作。跨通道操作可以分为置换、转换和归约操作。为了减小 API 的表面积,我们将为每类操作定义集体方法,然后将运算符作为输入。支持的运算符是 Operator 类的实例,并在 VectorOperators 类中定义为静态最终字段。一些常见的操作(例如,加、乘),称为全服务操作,将具有专用的方法,可以代替通用方法使用。

某些向量操作,如逐通道转换和重新解释,可以说是固有的 形状改变 操作。在向量计算中包含形状改变操作可能会对可移植性和性能产生意想不到的影响。因此,在适用的情况下,API 将定义此类操作的附加形状不变版本。鼓励用户使用形状不变的操作版本编写形状不变的代码。此外,将在 Javadoc 中明确调用形状改变操作。

Vector<E> 声明了一组所有元素类型都支持的常见向量操作的方法。为了支持特定于元素类型的操作,Vector<E> 有六个抽象子类,每个支持的元素类型都有一个:ByteVectorShortVectorIntVectorLongVectorFloatVectorDoubleVector。这些子类定义了额外的操作,这些操作与元素类型绑定,因为方法签名引用了元素类型(或等效的数组类型),例如归约操作(例如,将所有元素求和为一个标量值)或将向量元素存储到数组中。它们还定义了针对整数子类型的额外全服务操作,如位操作(例如,逻辑或),以及针对浮点类型的特定操作,如数学运算(例如,超越函数如 pow())。

这些类进一步由为不同形状(大小)的向量定义的具体子类扩展。

具体子类是非公开的,因为不需要提供特定于类型和形状的操作。这减少了 API 的表面积,使其成为一个关注点的总和,而不是一个乘积。因此,无法直接构造具体的 Vector 类实例。相反,实例是通过在基本 Vector<E> 类及其特定于类型的子类中定义的工厂方法获得的。这些方法将所需向量实例的种类作为输入。工厂方法提供了获取向量实例的不同方式,例如元素初始化为默认值的向量实例(零向量),或从数组中获取的向量,此外还提供了在不同类型或形状的向量之间进行转换的规范支持(例如,强制类型转换)。

为了支持控制流,相关的向量操作将可选地接受由公共抽象类 VectorMask<E> 表示的掩码。掩码中的每个元素,即布尔值或位,对应于一个向量通道。当掩码是操作的输入时,它决定是否将操作应用于每个通道;如果通道的掩码位被设置(为 true),则应用该操作。如果掩码位未设置(为 false),则会发生替代行为。与向量类似,VectorMask<E> 的实例是针对每种元素类型和长度组合定义的(私有)具体子类的实例。在操作中使用的 VectorMask<E> 的实例应具有与参与操作的 Vector<E> 实例相同的类型和长度。比较操作生成掩码,然后可以将这些掩码作为其他操作的输入,以选择性地禁用对特定通道的操作,从而模拟控制流。创建掩码的另一种方法是使用 VectorMask<E> 中的静态工厂方法。

我们预计,掩码在开发具有通用形状的向量计算中可能会发挥重要作用。(这一预期基于谓词寄存器在 ARM 可伸缩向量扩展以及 Intel 的 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
// 示例 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 位宽浮点数的种类(species)。该种类存储在 static final 字段中,因此运行时编译器会将字段的值视为常量,从而能够更好地优化向量计算。

向量计算包含一个主循环核心,它以向量长度(即种类长度)的步长遍历数组。静态方法 fromArray() 从数组 ab 的对应索引处加载给定种类的 float 向量。然后流畅地执行操作,最后将结果存储到数组 c 中。

我们使用由 indexInRange() 生成的掩码来防止读取 / 写入超过数组长度。前 floor(a.length / SPECIES.length()) 次迭代将具有所有通道都设置的掩码。只有当 a.length 不是 SPECIES.length() 的倍数时,最后一次迭代才会具有前 a.length % SPECIES.length() 个通道设置的掩码。

由于在所有迭代中都使用了掩码,因此上述实现对于较大的数组长度可能无法实现最佳性能。可以如下方式实现没有掩码的相同计算:

java
// 示例 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;
    }
}

在向量计算之后,长度小于种类长度的“尾部”元素通过标量计算进行处理。处理尾部元素的另一种方法是使用单个带掩码的向量计算。

当处理大型数组时,上述实现达到了最佳性能。

对于第二个示例,在支持 AVX 的 Intel x64 处理器上,HotSpot 编译器应生成类似于以下的机器代码:

java
  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 原型及其实现(Panama 项目开发存储库的 vectorIntrinsics 分支)对示例测试代码进行 JMH 微基准测试的实际输出。它显示了 C2 生成的机器代码的热点区域。这些区域明确地转译为了向量寄存器和向量硬件指令。(为了使转译更加清晰,禁用了循环展开,否则 HotSpot 应该能够使用现有的 C2 循环优化技术来进行展开。)所有的 Java 对象分配都被消除了。

支持更加复杂的、非平凡的向量计算,这些计算能够明确地转译为生成的机器代码,是一个重要的目标。

然而,这个特定的向量计算存在几个问题:

  1. 循环被硬编码为具体的向量形状,因此计算无法动态地适应架构所支持的最大形状,这个最大形状可能小于或大于 256 位。因此,代码的可移植性较差,性能也可能较低。

  2. 循环上限的计算虽然在这里很简单,但可能是编程错误的常见来源。

  3. 在末尾需要一个标量循环,这导致了代码的重复。

在这个 JEP 中,我们将解决前两个问题。可以获取一个优选的类型,其形状对当前架构是最优的,然后可以使用通用形状编写向量计算,并且可以在类型上调用一个方法来向下取整数组长度,例如:

java
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 不会将这些元素标量化。向量值被视为一个整体单元,就像 intlong 一样,映射到适当大小的硬件向量寄存器。内联类型将需要一些相关的增强来确保将向量值视为一个整体单元。

在内联类型可用之前,C2 将对 Vector 实例进行特殊处理,以克服逃逸分析的限制并避免装箱。因此,应避免对向量进行身份敏感的操作。

未来工作

一旦值类型准备就绪(参见 Project Valhalla),Vector API 将从中获得显著好处。Vector<E> 的实例可以是值,其具体类是内联类型。这将使优化和表达向量计算变得更加容易。对于特定类型的 Vector<E> 的子类型(如 IntVector),使用内联类型的泛型特化和特定类型的方法声明可能不再需要。

因此,Vector API 的未来版本将如上所述利用内联类型和增强的泛型。因此,我们将在 JDK 的多个版本中孵化该 API,并随着内联类型的可用性进行调整。

当 JEP 370 Foreign-Memory Access 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 架构,其编程模型动态地调整以匹配硬件支持的单一固定形状。我们欢迎并鼓励 OpenJDK 贡献者参与 HotSpot 的 ARM 特定区域的工作。

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