使用Java 18的Vector API提高性能 - martin

22-03-29 banq

Java Vector API 为现代 CPU 的数据并行功能提供了一个抽象层。
由于不同的处理器架构有不同的风格,因此没有简单的解决方案来利用软件中特定于平台的功能。通常需要以特定于平台的方式编写代码并利用平台的特定功能来获得出色的性能优势。Vector API 试图使开发人员能够以与平台无关的方式编写数据并行软件。

这篇博文试图在一些示例中探索新的 Vector API 提供的可能性,以及对于特定用例是否值得探索潜在的实现。

为了解释 Java Vector API 抽象是如何工作的,我们需要探索不同的 CPU 架构并提供对数据并行计算的基本理解。然而,这个概念并不是那么新。它在 C# 中已经存在了一段时间,并且已被证明是在现代硬件架构上利用数据并行计算的好方法。
与常规计算操作相比,如 1+1,在一次操作中添加两个“数据”,数据并行操作是在多个“数据”上执行简单的操作(例如,+)同时。这种操作模式称为 SIMD(单指令,多数据),而传统的执行方式称为 SISD(单指令,单数据)。性能加速的结果是在一个 CPU 周期内对多个“数据”应用相同的操作。
由于处理器采用不同的 uArch(x86、ARM),它们的 SIMD 实现存在显着差异。
 

简单求和
从一个简单的例子开始,我们可以详细看看下面的代码片段:

public static int[] simpleSum(int[] a, int[] b) {
    var c = new int[a.length];
    for (var i = 0; i < a.length; i++) {
        c[i] = a[i] + b[i];
    }
    return c;
}

同样的代码通过Vector API翻译成数据并行加速代码:

private static final VectorSpecies<Integer> SPECIES = IntVector.SPECIES_PREFERRED;
public static int[] vectorSum(int[] a, int[] b) {
    var c = new int[a.length];
    var upperBound = SPECIES.loopBound(a.length);

    var i = 0;
    for (; i < upperBound; i += SPECIES.length()) {
        var va = IntVector.fromArray(SPECIES, a, i);
        var vb = IntVector.fromArray(SPECIES, b, i);
        var vc = va.add(vb);
        vc.intoArray(c, i);
    }
    // Compute elements not fitting in the vector alignment.
    for (; i < a.length; i++) {
        c[i] = a[i] + b[i];
    }

    return c;

}

这段代码需要更多的说明。
在代码能够运行并利用SIMD加速之前,必须确定数据宽度。AVX兼容CPU可以处理256比特,而AVX-512可以提供512比特的数据宽度。

private static final VectorSpecies<Integer> SPECIES = IntVector.SPECIES_PREFERRED;

因此,我们需要正确配置循环的大小。简单的for-loop使用i++将索引递增1;在这个例子中,需要根据SIMD寄存器的数据宽度进行转移。
在AVX-512(512位)上运行的整数运算的情况下,我们必须以16为单位递增。
第一个迭代执行a[0]+b[0]到a[15]+b[15]的操作。
下一个操作对a[16]+b[16]到a[31]+b[31]执行同样的操作,以此类推

//Determines the last index that fits the registers and cuts off any 'overhanging' items.
var upperBound = SPECIES.loopBound(a.length);

var i = 0;
// Increment by the data-width!
for (; i < upperBound; i += SPECIES.length()) {
   ....
}


最后,需要处理所有剩余的项目,这些项目没有在数据宽度内对齐。因此,该操作必须以非并行的方式进行,从第一个未被矢量循环触及的项目开始。

for (; i < a.length; i++) { // cleanup loop
    c[i] = a[i] + b[i];
}

Vector API的代码可能看起来有点奇怪,但与我以前使用C的经验相比,设计非常相似。

详细点击标题