DirectX11 With Windows SDK--27 計算着色器:雙調排序

前言

上一章我們用一個比較簡單的例子來嘗試使用計算着色器,但是在看這一章內容之前,你還需要了解下面的內容:

章節
26 計算着色器:入門
深入理解與使用緩衝區資源

這一章我們繼續用一個計算着色器的應用實例作爲切入點,進一步瞭解相關知識。

DirectX11 With Windows SDK完整目錄

Github項目源碼

歡迎加入QQ羣: 727623616 可以一起探討DX11,以及有什麼問題也可以在這裏彙報。

線程標識符

對於線程組(大小(ThreadDimX, ThreadDimY, ThreadDimZ))中的每一個線程,它們都有一個唯一的線程ID值。我們可以使用系統值SV_GroupThreadID來取得,它的索引範圍爲(0, 0, 0)(ThreadDimX - 1, ThreadDimY - 1, ThreadDimZ - 1)

而對於整個線程組來說,由於線程組集合也是在3D空間中排布,它們也有一個唯一的線程組ID值。我們可以使用系統值SV_GroupID來取得,線程組的索引範圍取決於調用ID3D11DeviceContext::Dispatch時提供的線程組(大小(GroupDimX, GroupDimY, GroupDimZ)),範圍爲(0, 0, 0)(GroupDimX - 1, GroupDimY - 1, GroupDimZ - 1)

緊接着就是系統值SV_GroupIndex,它是單個線程組內的線程三維索引的一維展開。若已知線程組的大小爲(ThreadDimX, ThreadDimY, ThreadDimZ),則可以確定SV_GroupIndexSV_GroupThreadID滿足下面關係:

SV_GroupIndex = SV_GroupThreadID.z * ThreadDimX * ThreadDimY + SV_GroupThreadID.y * ThreadDimX + SV_GroupThreadID.x;

最後就是系統值SV_DispatchThreadID,線程組中的每一個線程在ID3D11DeviceContext::Dispatch提供的線程組集合中都有其唯一的線程ID值。若已知線程組的大小爲 (ThreadDimX, ThreadDimY, ThreadDimZ),則可以確定SV_DispatchThreadIDSV_GroupThreadIDSV_GroupID滿足以下關係:

SV_DispatchThreadID.xyz = SV_GroupID.xyz * float3(ThreadDimX, ThreadDimY, ThreadDimZ) + SV_GroupThreadID.xyz;

共享內存和線程同步

在一個線程組內,允許設置一片共享內存區域,使得當前線程組內的所有線程都可以訪問當前的共享內存。一旦設置,那麼每個線程都會各自配備一份共享內存。共享內存的訪問速度非常快,就像寄存器訪問CPU緩存那樣。

共享內存的聲明方式如下:

groupshared float4 gCache[256];

對於每個線程組來說,它所允許分配的總空間最大爲32kb(即8192個標量,或2048個向量)。內部線程通常應該使用SV_ThreadGroupID來寫入共享內存,這樣以保證每個線程不會出現重複寫入操作,而讀取共享內存一般是線程安全的。

分配太多的共享內存會導致性能問題。假如一個多處理器支持32kb的共享內存,然後你的計算着色器需要20kb的共享內存,這意味着一個多處理器只適合處理一個線程組,因爲剩餘的共享內存不足以給新的線程組運行,這也會限制GPU的並行運算,當該線程組因爲某些原因需要等待,會導致當前的多處理器處於閒置狀態。因此保證一個多處理器至少能夠處理兩個或以上的線程組(比如每個線程組分配16kb以下的共享內存),以儘可能減少該多處理器的閒置時間。

現在來考慮下面的代碼:

Texture2D gInput : register(t0);
RWTexture2D<float4> gOutput : register(u0);

groupshared float4 gCache[256];

[numthreads(256, 1, 1)]
void CS(uint3 GTid : SV_GroupThreadID,
    uint3 DTid : SV_DispatchThreadID)
{
    // 將紋理像素值緩存到共享內存
    gCache[GTid.x] = gInput[DTid.xy];
    
    // 取出共享內存的值進行計算
    
    // 注意!!相鄰的兩個線程可能沒有完成對紋理的採樣
    // 以及存儲到共享內存的操作
    float left = gCache[GTid.x - 1];
    float right = gCache[GTid.x + 1];
    
    // ...
}

因爲多個線程同時運行,同一時間各個線程當前執行的指令有所偏差,有的線程可能已經完成了共享內存的賦值操作,有的線程可能還在進行紋理採樣操作。如果當前線程正在讀取相鄰的共享內存片段,結果將是未定義的。爲了解決這個問題,我們必須在讀取共享內存之前讓當前線程等待線程組內其它的所有線程完成寫入操作。這裏我們可以使用GroupMemoryBarrierWithGroupSync函數:

Texture2D gInput : register(t0);
RWTexture2D<float4> gOutput : register(u0);

groupshared float4 gCache[256];

[numthreads(256, 1, 1)]
void CS(uint3 GTid : SV_GroupThreadID,
    uint3 DTid : SV_DispatchThreadID)
{
    // 將紋理像素值緩存到共享內存
    gCache[GTid.x] = gInput[DTid.xy];
    
    // 等待所有線程完成寫入
    GroupMemoryBarrierWithGroupSync();
    
    // 現在讀取操作是線程安全的,可以開始進行計算
    float left = gCache[GTid.x - 1];
    float right = gCache[GTid.x + 1];
    
    // ...
}

雙調排序

雙調序列

所謂雙調序列(Bitonic Sequence),是指由一個非嚴格遞增序列X(允許相鄰兩個數相等)和非嚴格遞減序列Y構成的序列,比如序列\((5, 3, 2, 1, 4, 6, 6, 12)\)

定義:一個序列\(a_1 , a_2, ..., a_n\)是雙調序列,需要滿足下面條件:

  1. 存在一個\(a_k(1 <= k <= n)\),使得\(a_1 >= ... >= a_k <= ... <= a_n\)成立,或者\(a_1 <= ... <= a_k >= ... >= a_n\)成立;
  2. 序列循環移位後仍能夠滿足條件(1)

Batcher歸併網絡

Batcher歸併網絡是由一系列Batcher比較器組成的,Batcher比較器是指在兩個輸入端給定輸入值x和y,再在兩個輸出端輸出最大值\(max(x, y)\)和最小值\(min(x, y)\)

雙調歸併網絡

雙調歸併網絡是基於Batch定理而構建的。該定理是說將任意一個長爲2n的雙調序列分爲等長的兩半X和Y,將X中的元素與Y中的元素按原序比較,即\(a_i\)\(a_{i+n}(i <= n)\)比較,將較大者放入MAX序列,較小者放入MIN序列。則得到的MAX序列和MIN序列仍然是雙調序列,並且MAX序列中的任意一個元素不小於MIN序列中的任意一個元素。

根據這個原理,我們可以將一個n元素的雙調序列通過上述方式進行比較操作來得到一個MAX序列和一個MIN序列,然後對這兩個序列進行遞歸處理,直到序列不可再分割爲止。最終歸併得到的爲一個有序序列。

這裏我們用一張圖來描述雙調排序的全過程:

其中箭頭方向指的是兩個數交換後,箭頭段的數爲較大值,圓點段的數爲較小值。

我們可以總結出如下規律:

  1. 每一趟排序結束會產生連續的雙調序列,除了最後一趟排序會產生我們所需要的單調序列
  2. 對於2^k個元素的任意序列,需要進行k趟排序才能產生單調序列
  3. 對於由\(2^{k-1}\)個元素的單調遞增序列和\(2^{k-1}\)個元素的單調遞減序列組成的雙調序列,需要進行k趟交換才能產生2^k個元素的單調遞增序列
  4. 在第n趟排序中的第m趟交換,若兩個比較數中較小的索引值爲i,那麼與之進行交換的數索引爲\(i+2^{n-m}\)

雙調排序的空間複雜度爲\(O(n)\),時間複雜度爲\(O(n{(lg(n))}^2)\),看起來比\(O(nlg(n))\)系列的排序算法慢上一截,但是得益於GPU的並行計算,可以看作同一時間內有n個線程在運行,使得最終的時間複雜度可以降爲\(O({(lg(n))}^2)\),效率又上了一個檔次。

需要注意的是,雙調排序要求排序元素的數目爲\(2^k, (k>=1)\),如果元素個數爲\(2^k < n < 2^{k+1}\),則需要填充數據到\(2^{k+1}\)個。若需要進行升序排序,則需要填充足夠的最大值;若需要進行降序排序,則需要填充足夠的最小值。

排序核心代碼實現

本HLSL實現參考了directx-sdk-samples,雖然裏面的實現看起來比較簡潔,但是理解它的算法實現費了我不少的時間。個人以自己能夠理解的形式對它的實現進行了修改,因此這裏以我這邊的實現版本來講解。

首先是排序需要用到的資源和常量緩衝區,定義在BitonicSort.hlsli

// BitonicSort.hlsli
Buffer<uint> gInput : register(t0);
RWBuffer<uint> gData : register(u0);

cbuffer CB : register(b0)
{
    uint gLevel;        // 2^需要排序趟數
    uint gDescendMask;  // 下降序列掩碼
    uint gMatrixWidth;  // 矩陣寬度(要求寬度>=高度且都爲2的倍數)
    uint gMatrixHeight; // 矩陣高度
}

然後是核心的排序算法:

// BitonicSort_CS.hlsl
#include "BitonicSort.hlsli"

#define BITONIC_BLOCK_SIZE 512

groupshared uint shared_data[BITONIC_BLOCK_SIZE];

[numthreads(BITONIC_BLOCK_SIZE, 1, 1)]
void CS(uint3 Gid : SV_GroupID,
    uint3 DTid : SV_DispatchThreadID,
    uint3 GTid : SV_GroupThreadID,
    uint GI : SV_GroupIndex)
{
    // 寫入共享數據
    shared_data[GI] = gData[DTid.x];
    GroupMemoryBarrierWithGroupSync();
    
    // 進行排序
    for (uint j = gLevel >> 1; j > 0; j >>= 1)
    {
        uint smallerIndex = GI & ~j;
        uint largerIndex = GI | j;
        bool isDescending = (bool) (gDescendMask & DTid.x);
        bool isSmallerIndex = (GI == smallerIndex);
        uint result = ((shared_data[smallerIndex] <= shared_data[largerIndex]) == (isDescending == isSmallerIndex)) ?
            shared_data[largerIndex] : shared_data[smallerIndex];
        GroupMemoryBarrierWithGroupSync();

        shared_data[GI] = result;
        GroupMemoryBarrierWithGroupSync();
    }
    
    // 保存結果
    gData[DTid.x] = shared_data[GI];
}

可以看到,我們實際上可以將遞歸過程轉化成迭代來實現。

現在我們先從核心排序算法講起,由於收到線程組的線程數目、共享內存大小限制,這裏定義一個線程組包含512個線程,即一個線程組最大允許排序的元素數目爲512。這裏共享內存的作用在這裏是臨時緩存中間排序的結果。

首先,我們需要將數據寫入共享內存中:

// 寫入共享數據
shared_data[GI] = gData[DTid.x];
GroupMemoryBarrierWithGroupSync();

接着就是要開始遞歸排序的過程,其中gLevel的含義爲單個雙調序列的長度,它也說明了需要對該序列進行\(lg(gLevel)\)趟遞歸交換。

在一個線程中,我們僅知道該線程對應的元素,但現在我們還需要做兩件事情:

  1. 找到需要與該線程對應元素進行Batcher比較的另一個元素
  2. 判斷當前線程對應元素與另一個待比較元素相比,是較小索引還是較大索引

這裏用到了位運算的魔法。先舉個例子,當前j爲4,則待比較兩個元素的索引分別爲2和6,這兩個索引值的區別在於索引2(二進制010)和索引6(二進制110),前者二進制第三位爲0,後者二進制第三位爲1.

但只要我們知道上述其中的一個索引,就可以求出另一個索引。較小索引值的索引可以通過屏蔽二進制的第三位得到,而較大索引值的索引可以通過按位或運算使得第三位爲1來得到:

uint smallerIndex = GI & ~j;
uint largerIndex = GI | j;
bool isSmallerIndex = (GI == smallerIndex);

然後就是判斷當前元素是位於當前趟排序完成後的遞增序列還是遞減序列,比如序列\((4, 6, 4, 3, 5, 7, 2, 1)\),現在要進行第二趟排序,那麼前後4個數將分別生成遞增序列和遞減序列,我們可以設置gDescendMask的值爲4(二進制100),這樣二進制索引範圍在100到111的值(對應十進制4-7)處在遞減序列,如果這個雙調序列長度爲16,那麼索引4-7和12-15的兩段序列都可以通過gDescendMask來判斷出處在遞減序列:

bool isDescending = (bool) (gDescendMask & DTid.x);

最後就是要確定當前線程對應的共享內存元素需要得到較小值,還是較大值了。這裏又以一個雙調序列\((2, 5, 7, 4)\)爲例,待比較的兩個元素爲5和4,當前趟排序會將它變爲單調遞增序列,即所處的序列爲遞增序列,當前線程對應的元素爲5,shared_data[smallerIndex] <= shared_data[largerIndex]的比較結果爲>,那麼它將拿到(較小值)較大索引的值。經過第一趟交換後將變成\((2, 4, 7, 5)\),第二趟交換就不討論了。

根據對元素所處序列、元素當前索引和比較結果的討論,可以產生出八種情況:

所處序列 當前索引 比較結果 取值結果
遞減 小索引 <= (較大值)較大索引的值
遞減 大索引 <= (較小值)較大索引的值
遞增 小索引 <= (較小值)較小索引的值
遞增 大索引 <= (較大值)較小索引的值
遞減 小索引 > (較大值)較小索引的值
遞減 大索引 > (較小值)較小索引的值
遞增 小索引 > (較小值)較大索引的值
遞增 大索引 > (較大值)較大索引的值

顯然現有的變量判斷較大/較小索引值比判斷較大值/較小值容易得多。上述結果表可以整理成下面的代碼:

uint result = ((shared_data[smallerIndex] <= shared_data[largerIndex]) == (isDescending == isSmallerIndex)) ?
    shared_data[largerIndex] : shared_data[smallerIndex];
GroupMemoryBarrierWithGroupSync();

shared_data[GI] = result;
GroupMemoryBarrierWithGroupSync();

在C++中,現在有如下資源和着色器:

ComPtr<ID3D11Buffer> mConstantBuffer;           // 常量緩衝區
ComPtr<ID3D11Buffer> mTypedBuffer1;             // 有類型緩衝區1
ComPtr<ID3D11Buffer> mTypedBuffer2;             // 有類型緩衝區2

ComPtr<ID3D11UnorderedAccessView> mDataUAV1;    // 有類型緩衝區1對應的無序訪問視圖
ComPtr<ID3D11UnorderedAccessView> mDataUAV2;    // 有類型緩衝區2對應的無序訪問視圖
ComPtr<ID3D11ShaderResourceView> mDataSRV1;     // 有類型緩衝區1對應的着色器資源視圖
ComPtr<ID3D11ShaderResourceView> mDataSRV2;     // 有類型緩衝區2對應的着色器資源視圖

然後就是對512個元素進行排序的部分代碼(size爲2的次冪):

void GameApp::SetConstants(UINT level, UINT descendMask, UINT matrixWidth, UINT matrixHeight);

//
// GameApp::GPUSort
//

md3dImmediateContext->CSSetShader(mBitonicSort_CS.Get(), nullptr, 0);
md3dImmediateContext->CSSetUnorderedAccessViews(0, 1, mDataUAV1.GetAddressOf(), nullptr);

// 按行數據進行排序,先排序level <= BLOCK_SIZE 的所有情況
for (UINT level = 2; level <= size && level <= BITONIC_BLOCK_SIZE; level *= 2)
{
    SetConstants(level, level, 0, 0);
    md3dImmediateContext->Dispatch((size + BITONIC_BLOCK_SIZE - 1) / BITONIC_BLOCK_SIZE, 1, 1);
}

給更多的數據排序

上述代碼允許我們對元素個數爲2到512的序列進行排序,但緩衝區的元素數目必須爲2的次冪。由於在CS4.0中,一個線程組最多允許一個線程組包含768個線程,這意味着雙調排序僅允許在一個線程組中對最多512個元素進行排序。

這裏我們看一個例子,假如有一個16元素的序列,然而線程組僅允許包含最多4個線程,那我們將其放置在一個4x4的矩陣內:

然後對矩陣轉置:

可以看到,通過轉置後,列數據變換到行數據的位置,這樣我們就可以進行跨度更大的交換操作了。處理完大跨度的交換後,我們再轉置回來,處理行數據即可。

現在假定我們已經對行數據排完序,下圖演示了剩餘的排序過程:

但是在線程組允許最大線程數爲4的情況下,通過二維矩陣最多也只能排序16個數。。。。也許可以考慮三維矩陣轉置法,這樣就可以排序64個數了哈哈哈。。。

不過還有一個情況我們要考慮,就是元素數目不爲(2x2)的倍數,無法構成一個方陣,但我們也可以把它變成對兩個方陣轉置。這時矩陣的寬是高的兩倍:

由於元素個數爲32,它的最大索引跨度爲16,轉置後的索引跨度爲2,不會越界到另一個方陣進行比較。但是當gLevel到32時,此時進行的是單調排序,gDescendMask也必須設爲最大值32(而不是4),避免產生雙調序列。

負責轉置的着色器實現如下:

// MatrixTranspose_CS.hlsl
#include "BitonicSort.hlsli"

#define TRANSPOSE_BLOCK_SIZE 16

groupshared uint shared_data[TRANSPOSE_BLOCK_SIZE * TRANSPOSE_BLOCK_SIZE];

[numthreads(TRANSPOSE_BLOCK_SIZE, TRANSPOSE_BLOCK_SIZE, 1)]
void CS(uint3 Gid : SV_GroupID,
    uint3 DTid : SV_DispatchThreadID,
    uint3 GTid : SV_GroupThreadID,
    uint GI : SV_GroupIndex)
{
    uint index = DTid.y * gMatrixWidth + DTid.x;
    shared_data[GI] = gInput[index];
    GroupMemoryBarrierWithGroupSync();
    
    uint2 outPos = DTid.yx % gMatrixHeight + DTid.xy / gMatrixHeight * gMatrixHeight;
    gData[outPos.y * gMatrixWidth + outPos.x] = shared_data[GI];
}

最後是GPU排序用的函數:

#define BITONIC_BLOCK_SIZE 512

#define TRANSPOSE_BLOCK_SIZE 16

void GameApp::GPUSort()
{
    UINT size = (UINT)mRandomNums.size();

    md3dImmediateContext->CSSetShader(mBitonicSort_CS.Get(), nullptr, 0);
    md3dImmediateContext->CSSetUnorderedAccessViews(0, 1, mDataUAV1.GetAddressOf(), nullptr);

    // 按行數據進行排序,先排序level <= BLOCK_SIZE 的所有情況
    for (UINT level = 2; level <= size && level <= BITONIC_BLOCK_SIZE; level *= 2)
    {
        SetConstants(level, level, 0, 0);
        md3dImmediateContext->Dispatch((size + BITONIC_BLOCK_SIZE - 1) / BITONIC_BLOCK_SIZE, 1, 1);
    }
    
    // 計算相近的矩陣寬高(寬>=高且需要都爲2的次冪)
    UINT matrixWidth = 2, matrixHeight = 2;
    while (matrixWidth * matrixWidth < size)
    {
        matrixWidth *= 2;
    }
    matrixHeight = size / matrixWidth;

    // 排序level > BLOCK_SIZE 的所有情況
    ComPtr<ID3D11ShaderResourceView> pNullSRV;
    for (UINT level = BITONIC_BLOCK_SIZE * 2; level <= size; level *= 2)
    {
        // 如果達到最高等級,則爲全遞增序列
        if (level == size)
        {
            SetConstants(level / matrixWidth, level, matrixWidth, matrixHeight);
        }
        else
        {
            SetConstants(level / matrixWidth, level / matrixWidth, matrixWidth, matrixHeight);
        }
        // 先進行轉置,並把數據輸出到Buffer2
        md3dImmediateContext->CSSetShader(mMatrixTranspose_CS.Get(), nullptr, 0);
        md3dImmediateContext->CSSetShaderResources(0, 1, pNullSRV.GetAddressOf());
        md3dImmediateContext->CSSetUnorderedAccessViews(0, 1, mDataUAV2.GetAddressOf(), nullptr);
        md3dImmediateContext->CSSetShaderResources(0, 1, mDataSRV1.GetAddressOf());
        md3dImmediateContext->Dispatch(matrixWidth / TRANSPOSE_BLOCK_SIZE, 
            matrixHeight / TRANSPOSE_BLOCK_SIZE, 1);

        // 對Buffer2排序列數據
        md3dImmediateContext->CSSetShader(mBitonicSort_CS.Get(), nullptr, 0);
        md3dImmediateContext->Dispatch(size / BITONIC_BLOCK_SIZE, 1, 1);

        // 接着轉置回來,並把數據輸出到Buffer1
        SetConstants(matrixWidth, level, matrixWidth, matrixHeight);
        md3dImmediateContext->CSSetShader(mMatrixTranspose_CS.Get(), nullptr, 0);
        md3dImmediateContext->CSSetShaderResources(0, 1, pNullSRV.GetAddressOf());
        md3dImmediateContext->CSSetUnorderedAccessViews(0, 1, mDataUAV1.GetAddressOf(), nullptr);
        md3dImmediateContext->CSSetShaderResources(0, 1, mDataSRV2.GetAddressOf());
        md3dImmediateContext->Dispatch(matrixWidth / TRANSPOSE_BLOCK_SIZE,
            matrixHeight / TRANSPOSE_BLOCK_SIZE, 1);

        // 對Buffer1排序剩餘行數據
        md3dImmediateContext->CSSetShader(mBitonicSort_CS.Get(), nullptr, 0);
        md3dImmediateContext->Dispatch(size / BITONIC_BLOCK_SIZE, 1, 1);
    }
}

最後是std::sort和雙調排序的比較結果:

可以初步看到雙調排序的排序用時比較穩定,而快速排序明顯隨元素數目增長而變慢。

DirectX11 With Windows SDK完整目錄

Github項目源碼

歡迎加入QQ羣: 727623616 可以一起探討DX11,以及有什麼問題也可以在這裏彙報。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章