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

banq發表於2022-03-29

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的經驗相比,設計非常相似。

詳細點選標題

相關文章