WebGPU的計算著色器實現冒泡排序

大家好~本文使用WebGPU的計算著色器，實現了奇偶排序。
奇偶排序是冒泡排序的並行版本，在1996年由J Kornerup提出。它解除了每輪冒泡間的串列依賴以及每輪冒泡內部的串列依賴，使得冒泡操作可以並行執行

目錄

最終版本的程式碼在這裡

介紹奇偶排序演算法

假設待排序的數組為Arr1
在奇數步中，Arr1中奇數項與相鄰的右邊一項比較和交換；
在偶數步中，Arr1中奇數項與相鄰的左邊一項比較和交換;
直到一步中沒有交換項，則停止

舉例來說的話，如下圖所示：
image

在每步中，紅框內的兩項進行比較和交換；
直到一步中沒有交換項，則停止

分析時間複雜度

與冒泡排序一樣，總的比較次數不變，依然為O(n^2)次
但因為為並行執行，所以時間複雜度降低為O(log2(n^2))=O(n)

需求

排序的需求如下所示:

  • 對一個包含128個數字的數組進行升序的排序

初步設計

因為數組可以兩兩分為64個組，每個組並行執行操作，所以計算著色器只使用一個workgroup，包含64個局部單位，每個局部單位對應一個組；

在每個局部單位中：
啟動一個while循環，執行每步操作，然後同步，最後判斷所有局部單位在該步驟中是否有交換操作，如果都沒有的話則停止循環

「執行每步操作」時判斷該步驟是奇數還是偶數步，從而取對應的兩項來比較和交換

程式碼實現

經過上面的設計後，現在我們來實現程式碼
計算著色器程式碼如下所示：

//64個局部單位
const workgroupSize = 64;

// 局部單位之間的共享變數，用於存放128個數字
var<workgroup> sharedData: array<f32,128>;
// 局部單位之間的共享變數，用於標誌所有局部單位在該步驟中是否有交換操作（只要有任意一個局部單位在該步驟中有交換操作，則該標誌為true）
var<workgroup> isSwap: bool;
// 局部單位之間的共享變數，用於記錄步驟數，從而判斷是奇數還是偶數步
var<workgroup> stepCount: u32;

struct BeforeSortData {
  data : array<f32, 128>
}
struct AfterSortData {
  data : array<f32, 128>
}

//待排序的數組
@binding(0) @group(0) var<storage, read> beforeSortData : BeforeSortData;
//排序後的數組
@binding(1) @group(0) var<storage, read_write> afterSortData :  AfterSortData;

@compute @workgroup_size(workgroupSize, 1, 1)
fn main(
@builtin(global_invocation_id) GlobalInvocationID : vec3<u32>,
) {
//將待排序的數據讀取到共享變數中

var index = GlobalInvocationID.x * 2;
sharedData[index] = beforeSortData.data[index];
sharedData[index+ 1 ] = beforeSortData.data[index + 1];

//初始化共享變數

isSwap = false;
stepCount = 0;

//同步
workgroupBarrier();

//開始循環
while(true){
    var firstIndex:u32;
    var secondIndex:u32;

    //判斷該步驟是奇數還是偶數步，從而得到對應的兩項的序號

    //偶數步
    if(stepCount % 2 == 0){
        firstIndex = index + 1;
        secondIndex = index + 2;
    }
    //奇數步
    else{
        firstIndex = index;
        secondIndex = index + 1;
    }

    //確保沒超過邊界
    if(secondIndex < 128){
        //將大的一項交換到後面，從而實現升序
        if(sharedData[firstIndex] > sharedData[secondIndex]){
            var temp = sharedData[firstIndex];
            sharedData[firstIndex] = sharedData[secondIndex];
            sharedData[secondIndex] = temp;

            isSwap = true;
        }
    }

    stepCount += 1;

    workgroupBarrier();

    //如果該步驟中沒有交換操作的話則停止循環
    if(!isSwap){
        break;
    }
}

//將排序後的數據傳給返回給CPU端的Storage Buffer中，從而可在CPU端得到排序後的結果
afterSortData.data[index] = sharedData[index];
afterSortData.data[index + 1] = sharedData[index + 1];
}

本來我是想像設置workgroupSize一樣，將128設為const的，如下所示：

const workgroupSize = 64;
const itemCount = 128;

var<workgroup> sharedData: array<f32,itemCount>;

但是運行時會報錯！照理說根據WGSL的文檔，是不應該報錯的！所以不清楚是我沒搞清楚還是WGSL目前的bug？
（另外，如果使用override而不是const，也會報錯！）

如果改為使用workgroupSize，則不會報錯。程式碼如下所示：

const workgroupSize = 64;

var<workgroup> sharedData: array<f32,workgroupSize>;

這是因為workgroupSize在@workgroup_size中使用了。程式碼如下所示：

@compute @workgroup_size(workgroupSize, 1, 1)

發現問題

運行程式碼後，會報錯：

1 warning(s) generated while compiling the shader:
:74:5 warning: 'workgroupBarrier' must only be called from uniform control flow
    workgroupBarrier();
    ^^^^^^^^^^^^^^^^

:77:5 note: control flow depends on non-uniform value
    if(!isSwap){
    ^^

:77:9 note: reading from workgroup storage variable 'isSwap' may result in a non-uniform value
    if(!isSwap){
        ^^^^^^

這是因為WGSL會進行Uniformity analysis檢查，確保像「workgroupBarrier」這種barries是在uniform control flow中安全地調用
WGSL在檢查時發現：因為isSwap被多個局部單位讀寫，所以為”non-uniform value”，導致所在的control flow為non-uniform，從而報錯

更多關於Uniformity的資料在這裡：
uniformity
Add the uniformity analysis to the WGSL spec
uniformity issues

改進設計

現在需要去掉isSwap的if判斷
因為isSwap的if判斷是用來結束循環的，那麼在去掉它之後我們就需要新的結束條件
因為總共有128個數字要排序，所以最多進行128步即可完成所有的排序

相關程式碼實現

所以去掉isSwap，把循環終止條件修改下，並且重構一下程式碼
相關程式碼改為：


fn _swap(firstIndex:u32, secondIndex:u32){
    var temp = sharedData[firstIndex];
    sharedData[firstIndex] = sharedData[secondIndex];
    sharedData[secondIndex] = temp;
}

fn _oddSort(index:u32) {
    var firstIndex = index;
    var secondIndex = index + 1;

    if(sharedData[firstIndex] > sharedData[secondIndex]){
        _swap(firstIndex, secondIndex);
    }
}

fn _evenSort(index:u32) {
    var firstIndex = index + 1;
    var secondIndex = index + 2;

    if(secondIndex <128 && sharedData[firstIndex] > sharedData[secondIndex]){
        _swap(firstIndex, secondIndex);
    }
}


@compute @workgroup_size(workgroupSize, 1, 1)
fn main(
@builtin(global_invocation_id) GlobalInvocationID : vec3<u32>,
) {
    ...

    var firstIndex:u32;
    var secondIndex:u32;

    for (var i: u32 = 0; i < 128; i += 1) {
        //偶數步
        if(stepCount % 2 == 0){
            _evenSort(index);
        }
        //奇數步
        else{
            _oddSort(index);
        }

        stepCount +=1 ;

        workgroupBarrier();

    }
    ...
}

現在就能夠正確運行了

改進設計

現在我們通過判斷步驟數是偶數還是奇數來進行對應的排序，這會造成「Warp Divergence」的優化問題：
不同的局部單位會進入不同的分支（偶數步或者奇數步），造成同一時刻除了正在執行的分支以外，其餘分支都被阻塞了，十分影響性能
如下圖所示：
image

參考資料：
CUDA編程——Warp Divergence

所以我們可以去掉stepCount和相關的判斷，改為在每次循環中分別執行奇數排序和偶數排序，並把循環次數減半

相關程式碼實現

修改後的相關程式碼為：

for (var i: u32 = 0; i < 64; i += 1) {
    _oddSort(index);
    workgroupBarrier();

    _evenSort(index);
    workgroupBarrier();
}

限制

現在程式有如下的限制：

  • 只有一個workgroup，意味著只能排序最多「局部單位數量*2」個數字

總結

感謝大家的學習~

計算著色器是SIMT架構，也就是說每條指令都是並行執行的。這與CPU的串列思維不同，所以我們要切換為並行的思維，並在需要同步的時候同步

另外，計算著色器的程式碼因為處在並行環境下，所以需要仔細優化
有下面的一些優化建議：

  • 減少Warp Divergence
    • 盡量減少if判斷
    • if中盡量使用常數判斷，如if(const value < 2)
  • 減少Bank Conflict
    • 對共享變數記憶體盡量連續地讀寫
      如對本文的共享數組sharedData，盡量連續的讀寫相鄰的序號
  • 減少同步操作
    • 如果已知確定的循環次數，可以展開循環，這樣可以減少同步操作和循環的指令開銷

參考資料

啥是Parallel Reduction
CUDA(六). 從並行排序方法理解並行化思維——冒泡、歸併、雙調排序的GPU實現
uniformity
CUDA編程——Warp Divergence

Tags:

VirMach 便宜 VPS

QNews

QNews