Shader變體大殺器：Specialization constants

UWATech發表於2024-06-26

【USparkle專欄】如果你深懷絕技，愛“搞點研究”，樂於分享也博採眾長，我們期待你的加入，讓智慧的火花碰撞交織，讓知識的傳遞生生不息！

Metal和Vulkan都提供了一個Specialization constants，是非常棒的API用以解決掉Shader變體過多的問題。具體實現程式碼放在最後。

什麼是Shader變體

著色器變體Shader variants，或者說著色器並列Shader permutation問題，指的是取一堆著色器程式碼，並用不同的選項編譯N次。在大多數情況下，這些排列直接與著色器支援的功能繫結，通常透過以“uber-shader”樣式編寫程式碼，具有許多不同的功能，這些功能可以獨立開啟和關閉。用一個超級簡單的例子，看看一個小的HLSL Pixel Shader程式碼，它使用前處理器宏來開啟和關閉不同的功能：

// ENABLE_NORMAL_MAP, ENABLE_EMISSIVE_MAP, ENABLE_AO_MAPPING are
// macros whose values are defined via compiler command line options
static const bool EnableNormalMap = ENABLE_NORMAL_MAP;
static const bool EnableEmissiveMap = ENABLE_EMISSIVE_MAP;
static const bool EnableAOMap = ENABLE_AO_MAPPING;

float4 PSMain(in PSInput input) : SV_Target0
{
    float3 normal = normalize(input.VtxNormal);

#if EnableNormalMap

    normal = ApplyNormalMap(normal, input.TangentFrame, input.UV);

#endif

    float3 albedo = input.Albedo;
    float3 ambientAlbedo = albedo;

#if EnableAOMap

    ambientAlbedo *= SampleAOMap(input.UV);

#endif

    float3 lighting = CalcLighting(albedo, ambientAlbedo, normal);

#if EnableEmissiveMap

    lighting += SampleEmissiveMap(input.UV);

#endif

    return float4(lighting, 1.0f);

}

　　

我們有3個功能可以在這裡啟用或禁用，如果所有這些功能都是獨立的，我們需要編譯2^3 = 8個排列。我們新增的每個新功能都意味著我們將排列計數翻倍！我們也可以將此推廣到具有2個以上不同狀態的非二進位制特徵，在這種情況下，計算間次總數的公式如下所示：

Shader變體大殺器：Specialization constants

隨著這種指數增長，不難想象當你新增越來越多的功能時，總排列計數達到數千或數百萬。你要求編譯器反覆使用這個大的單個著色器程式。這是一遍又一遍地解析和最佳化相同程式碼的大量浪費工作，除非多核CPU的機器或本地機器網路來分發編譯，比如IncredBuild，否則每次對著色器程式碼的更改都會變成幾個小時的時間。如果你編譯過Unreal Engine的原始碼並改過ush檔案，你應該深有體會。

說到主流的遊戲引擎。Unity透過允許使用者用著色器語言手寫著色器來暴露這一點，而Unreal Engine則允許使用者使用基於節點的視覺化介面建立材質網路。根據渲染器的設定方式，這些可能作為來自引擎本身的排列計數的附加乘數。例如，材質編輯器可能只允許使用者生成主要引數（如顏色、法線、粗糙度等），而引擎將其與實現實際照明計算的手寫著色器排列相結合。如果你有一個照明階段的100個排列和一個專案的100個材質，那麼恭喜你，你現在有10000個著色器了！

問題是，花費在編譯上的時間甚至不是整個問題，它只是其中的一部分。編譯所有這些著色器可以開始生成大量二進位制資料，特別是如果你生成除錯資訊，比如這樣：

glslc ARGS -g base.frag -o base.spv

　　

你可能會開始遇到儲存所有資料的問題，或者在執行時載入資料並將其儲存在記憶體中。但是，情況其實變得更糟了！最新的圖形API，如D3D和Vulkan，實際上實現了兩步編譯管道，DXC、FXC和GLSLC等離線著色器編譯器實際上會生成某種中間位元組碼，而不是可以在GPU著色器核心上執行的機器程式碼。GPU實際上在硬體裡有截然不同的Instruction Set Architecture，即ISA，即使在同一GPU供應商內，指令也經常從一代硬體更改為下一代。為了處理此設定，驅動程式將使用某種JIT（just-in-time）編譯器，該編譯器將你的DXBC、DXIL或SPIR-V轉換為著色器核心可以執行的最終ISA。整個管道通常是這樣的：

  1. 著色器原始碼由著色器編譯器離線編譯，如DXC、FXC或GLSLC，此步驟的輸出是硬體無關的中間位元組碼格式，DXBC（FXC）、DXIL（DXC）或SPIR-V（DXC或GLSLC）。

  2. 編譯的位元組碼由引擎載入到記憶體中。

  3. 引擎透過傳遞所需的狀態以及任何所需的著色器階段的編譯位元組碼來建立PSO，在此步驟中，驅動程式執行自己的JIT編譯器，將位元組碼轉換為在物理著色器核心上執行的特定於硬體的ISA。

  4. 引擎將PSO繫結到命令緩衝區，併發出使用PSO引用的著色器的繪製/排程命令。

  5. 命令緩衝區被提交到GPU，在那裡實際執行著色器程式。

擁有更多的Shader排列通常意味著多次呼叫JIT編譯步驟，這確實可以加起來。這通常不是這裡發生的簡單翻譯：大多數驅動程式將啟動一個成熟的最佳化編譯器（通常是LLVM），然後透過它執行你的位元組碼，以生成最終的ISA。如果你只是期望在新場景中流式傳輸時在後臺快速建立所有PSO，這真的會很難搞：你必須預先建立這些PSO，這將消耗大量時間和CPU資源，這反過來會影響你的有效載入時間。更糟糕的是，時間可能會因不同的GPU架構、不同的驅動程式版本、你在PSO描述中提供的其他狀態和著色器階段而異。然而，D3D12和Vulkan都提供了應用程式手動快取PSO的機制：

ID3D12PipelineState::GetCachedBlob // DX12

VkPipelineCache(3) // Vulkan

　　

但這些API往往很複雜，仍然無法幫助最佳化“首次啟動”載入時間。V社甚至將分散式快取系統整合到適用於OpenGL和Vulkan的Steam中。

即使你建立了PSO，未經控制的Shader排列爆炸的缺點仍然存在。要使用這些PSO，你需要在命令緩衝區上繪製或排程之前繫結它們。這個繫結步驟消耗了CPU時間：用單個PSO發出許多畫布可能比在每次畫之間切換PSO要便宜得多。最終結果是，隨著你增加著色器/PSO數量，生成命令緩衝區所需的CPU時間會增加。擁有許多PSO還可以阻止你使用例項化或批處理等技術將東西組合成單個繪圖或排程，這也增加了CPU時間。但這還不是全部：除非你使用特定於Nvidia的擴充套件，否則D3D12和Vulkan都無法在ExecuteIndirect或DrawIndirect/DispatchIndirect呼叫中從GPU中更改PSO，這意味著，如果你希望使用GPU-Driven技術在GPU上剔除和計算LOD，你將需要最小數量的ExecuteIndirect/DrawIndirect呼叫，這些呼叫與這些繪製使用的PSO數量相等。太多的PSO/著色器也會給光線跟蹤帶來問題，因為大型著色器表（Shader Tables）可能會導致硬體利用率過低。

最後，在GPU本身上執行著色器程式。現代GPU有很多技巧，即使在許多繪製呼叫和許多著色器開關的情況下，也能保持效能可接受。然而，在一般情況下，當你可以為他們提供大批次的工作時，他們仍然會實現更高的效能，而不需要切換狀態和著色器。因此，由於過於頻繁地切換著色器，你的GPU效能可能會下降。你的CPU效能也是如此：切換著色器/PSO通常涉及驅動程式中更昂貴的操作，這些操作需要在GPU上重新配置必要的Pipeline State，並處理切換著色器的下游影響（例如著色器使用的資源繫結的低階表示）。因此，如果你真的想在幾毫秒的CPU時間內提交大量繪製呼叫，你需要儘可能少地使用PSO和PSO交換機。

單一著色器檔案

為什麼著色器必須是單一的，而不是模組化，有很多include檔案？為了回答這個問題，得看下很早之前的API。

在圖形API的早期，與D3D8一起引入的原始SM 1.0對頂點著色器的硬限制為128個指令，畫素著色器基本上只是直接暴露了當時硬體的（非常有限的）暫存器組合器和紋理功能。在這一點上，它們只是少數指令！直到D3D9和後來，人們才用實際的高階著色器語言而不是手寫“彙編”來創作它們。

但由於語言的簡單性，著色器可以處於另一個水平：畢竟沒有建構函式，沒有副作用，直到SM 5.0才新增一般記憶體寫入！看看另一個簡單的HLSL畫素著色器：

float3 ComputeLighting(in float3 albedo, in float3 normal, in Light light)
{
    return saturate(dot(normal, light.Dir)) * albedo * light.Color;
}

float PSMain(in float3 albedo : ALBEDO, in float3 normal : NORMAL) : SV_Target0
{
    return ComputeLighting(albedo, normalize(normal), CBuffer.Light).x;
}

　　

如果期望編譯器獨立發出ComputeLighting指令，你可能希望它使用float3向量執行乘法，從而在計算飽和點積後產生6次總乘法。還希望PSMain收到float3結果，忽略Y/Z元件，然後返回X元件。在現實中，著色器編譯器從未以這種方式工作：相反，他們會完全內聯/扁平化函式呼叫call和死條dead-strip，導致生成等同於此程式碼的彙編：

float PSMain(in float3 albedo : ALBEDO, in float3 normal : NORMAL) : SV_Target0
{
    return saturate(dot(normal, light.Dir)) * albedo.x * light.Color.x;
}

　　

你可以看到C或C++編譯器在類似的地方做同樣的事情。但使用著色器是不同的，因為由於程式設計模型的簡單性，你幾乎可以保證獲得最佳化和內聯的結果。如果你再次回到早期的SM 2.0著色器時代，真的需要捏一撮指令，以避免達到指令限制，當你在數千個畫素上執行這些程式時，每個小加法或乘法都可以快速加起來。“壓平一切！”程式碼生成的風格也特別適合早期的可程式設計GPU，因為它們要麼不能進行動態流控制，要麼非常不擅長。因此，能夠以分支或迴圈風格編寫的著色器程式碼，並儘可能將其轉換為完全扁平和展開的東西，這對你有利。

這些漸進式微最佳化在現代GPU仍然有用，它們的branch和loop也比以前要好得多（至少只要流量控制是“均勻的”，這意味著wave/subgroup中的所有執行緒都採取相同的路徑），但流量控制永遠不會100%自由。再一次，如果多次完成一些用於回覆迴圈控制變數和檢查條件的額外說明可以加起來，最好讓編譯器展開指令。Flow Control還可以抑制編譯器喜歡在重疊無關程式碼的記憶體或數學指令的地方進行某些型別的最佳化，以提取額外的指令級並行性（簡稱ILP），這可能會減少孤立的波的延遲。

即使你忽略了最佳化方面，仍然有一個事實，即GPU著色器核心傳統上並沒有真正設定為處理真正的函式呼叫或動態排程。出於各種原因，這些著色器核心傾向於使用簡單的僅限暫存器的SIMD執行模型，沒有真正的堆疊。著色器核心通常還使用一個模型，其中波必須靜態分配它們在整個程式期間所需的所有暫存器，而不是可能涉及溢位到堆疊的更動態的模型。

如此多的排列實際上會使程式設計師更難對著色器原始碼進行實際手動最佳化，這既因為增加了迭代時間，也因為可能有太多的排列需要關注！編譯器執行的一些繁重的最佳化也可能對整體效能出人意料地糟糕。特別是aggressive unrolling和指令重疊往往會很快吞噬暫存器，這可能會導致佔用率降低。換句話說，這些最佳化可能會小幅減少單波的延遲，但可能會對GPU透過迴圈不同的活動波（active waves）來提高整體吞吐量的最佳化相違背。處理這個問題可能麻煩，因為有時感覺你必須騙過編譯器生成使用更少暫存器的程式碼，因為它不知道在著色器執行時GPU上還會發生什麼，它也不知道記憶體操作的預期延遲（又可能它在快取中，也可能沒有！）。

如何解決Shader變體

只編譯你需要的東西
這是減少排列計數的最簡單、最古老、可能最不有效的方法。一般的想法是，在你的2^N可能的排列中，它們中的一些子集要麼是多餘的，要麼是無效的，要麼永遠不會被實際使用。因此，如果你能去除不必要的排列，你將減少需要編譯和載入的著色器集。

優點：離線編譯更少，不需要對渲染器進行更改
缺點：沒有減少PSO，可能無法充分減少著色器數量，可能需要離線分析

延遲渲染或者Runtime VT
延遲渲染把大量資訊儲存在GBuffer中光照階段一併處理，透過Runtime VT把大量貼圖合併，這些都是直接從根源上減少大量的Shader檔案。

優點：直接減少了PSO和Shader數量
缺點：不能使用延遲渲染或者Runtime VT系統怎麼辦？

真正的函式呼叫和動態排程
將執行時函式呼叫與動態排程相結合可能比連結步驟更好，但它也更像是一個根本性的變化。雖然連結可以離線進行，沒有驅動程式輸入，但動態排程肯定需要驅動程式和硬體支援。大多數GPU使用的“將所有東西都塞在靜態分配的暫存器塊中”模型當然不適合真正的動態排程，而且很容易想象，程式設計的各種限制可能是必要的。

好訊息是，在PC上，我們現在有這個！壞訊息是，它非常受限，只適用於光線追蹤。在D3D12/DXR中，它基本上透過一種“執行時連結器”步驟來工作。基本上，你使用“lib”目標將一堆函式編譯成DXIL二進位制檔案（就像離線連結一樣），然後在執行時將所有部分組裝成一個組合狀態物件。稍後，當你呼叫DispatchRays時，驅動程式能夠動態執行正確的hit/miss/anyhit/etc.著色器，因為它已連結到狀態物件。有一個可呼叫的著色器功能，可以在沒有任何實際光線跟蹤的情況下使用，但它仍然需要從DispatchRays啟動的光線生成著色器中使用。換句話說：它現在可用於類似計算的著色器，但目前無法在圖形管道中使用。

Metal目前也透過提供可以從任何著色器階段呼叫的功能指標。也許這可以作為PC和Vulkan的啟發！

優點：離線編譯和 PSO 數量顯著減少，為其他新技術開啟了大門
缺點：GPU 效能可能更差，需要更改引擎處理著色器和 PSO 的方式

主角：Specialization constants

Vulkan和Metal都支援一個功能，稱為Specialization constants ，Metal稱該功能為Function specialization。

基本想法是這樣的：

  • 使用著色器程式碼中使用的全域性統一值（基本上就像UniformBuffer中的值）編譯著色器。

  • 建立PSO時，為該Uniform傳遞一個值，該值對於使用PSO的所有繪製和排程都是恆定的。

  • 驅動程式以某種方式確保傳遞的值在著色器程式中使用。這可能包括：

    • 將該值視為“推送常數”，基本上是從命令緩衝區設定的小型統一/恆定緩衝區

    • 將值修補到編譯的中間位元組碼或特定於供應商的ISA中

  • 當驅動程式進行JIT編譯時，將該值視為編譯時常量，並根據該值執行完全最佳化，包括常量摺疊constant folding和無效程式碼消除dead code elimination。

注意，這個常量真正起作用的時間，是在驅動程式進行JIT編譯時，不是在編譯Shader中間位元組碼或者建立PSO時，所以這是一個需要驅動層支援的功能。

優點：顯著減少離線編譯數量和時間，無需更改渲染器
缺點：可能會增加PSO建立時間，可能無法獲得驅動程式側的最佳化

這個方法，並不會減少PSO的數量，但是會減少Shader變體的數量從而減少ShaderCache的記憶體，簡單理解為，一個Shader檔案，在不增加變體的情況下，可以在PSO建立階段，根據輸入的Specialization constants的常數值生成多份PSO，並且在Shader繫結階段會最佳化掉無效branch的程式碼或者unroll迴圈的程式碼來減少Shader的實際指令數，提高Shader執行的效率。

具體實現是這樣的，在Shader裡建立一個常量：

layout (constant_id = 0) const int SPEC_CONSTANTS = 0;

　　

在建立GraphicPipline時，就是PSO那個階段，這個SPEC_CONSTANTS常量有多少個值，就建立多少個PSO，每個PSO對應一個常量的定值，比如：

// Use specialization constants 最佳化著色器變體
for (uint32_t i = 0; i < specConstantsCount; i++)
{
    uint32_t specConstants = i;
    VkSpecializationMapEntry specializationMapEntry = VkSpecializationMapEntry{};
    specializationMapEntry.constantID = 0;
    specializationMapEntry.offset = 0;
    specializationMapEntry.size = sizeof(uint32_t);
    VkSpecializationInfo specializationInfo = VkSpecializationInfo();
    specializationInfo.mapEntryCount = 1;
    specializationInfo.pMapEntries = &specializationMapEntry;
    specializationInfo.dataSize = sizeof(uint32_t);
    specializationInfo.pData = &specConstants;
    shaderStages[1].pSpecializationInfo = &specializationInfo;
    if (vkCreateGraphicsPipelines(device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &outPipelines[i]) != VK_SUCCESS)
    {
        throw std::runtime_error("failed to create graphics pipeline!");
    }

　　

然後，著色器的SPEC_CONSTANTS需要變化時，我們在CPU端提交渲染指令時，直接切換到對應的SPEC_CONSTANTS的PSO即可：

uint32_t specConstants = globalConstants.specConstants;

VkPipeline baseScenePassPipeline = baseScenePass.pipelines[specConstants];

vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, baseScenePassPipeline);

　　

在Shader裡面，但我們使用SPEC_CONSTANTS這個常量時，程式碼看上去是這樣的：

if (SPEC_CONSTANTS == 0)
{
    ...
}

else if (SPEC_CONSTANTS == 1)
{
    ...
}
...

　　

實際上，因為SPEC_CONSTANTS是常量，所以如果SPEC_CONSTANTS是0，那麼這段程式碼實際指令時就是這樣：

if (0 == 0)
{
    ...
}

// else if (0 == 1)
// {
//     ...
// }
// ...

　　

else if後面的程式碼會被直接最佳化掉，是不是非常的Nice！

參考文章

Improving shader performance with Vulkan’s specialization constants

The Shader Permutation Problem - Part 1: How Did We Get Here?

The Shader Permutation Problem - Part 2: How Do We Fix It?

這是侑虎科技第1612篇文章，感謝作者徐門子美供稿。歡迎轉發分享，未經作者授權請勿轉載。如果您有任何獨到的見解或者發現也歡迎聯絡我們，一起探討。（QQ群：793972859）

作者主頁：https://www.zhihu.com/people/BloodyGuys

再次感謝徐門子美的分享，如果您有任何獨到的見解或者發現也歡迎聯絡我們，一起探討。（QQ群：793972859）

相關文章