Shader变体大杀器：Specialization constants

【USparkle专栏】如果你深怀绝技，爱“搞点研究”，乐于分享也博采众长，我们期待你的加入，让智慧的火花碰撞交织，让知识的传递生生不息！

这是侑虎科技第1612篇文章，感谢作者徐門子美供稿。欢迎转发分享，未经作者授权请勿转载。如果您有任何独到的见解或者发现也欢迎联系我们，一起探讨。（QQ群：793972859）

作者主页：

https://www.zhihu.com/people/BloodyGuys

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个以上不同状态的非二进制特征，在这种情况下，计算间次总数的公式如下所示：

随着这种指数增长，不难想象当你添加越来越多的功能时，总排列计数达到数千或数百万。你要求编译器反复使用这个大的单个着色器程序。这是一遍又一遍地解析和优化相同代码的大量浪费工作，除非多核CPU的机器或本地机器网络来分发编译，比如IncredBuild（ https://www.incredibuild.com/ ），否则每次对着色器代码的更改都会变成几个小时的时间。如果你编译过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（https://learn.microsoft.com/en-us/windows/win32/api/d3d12/nf-d3d12-id3d12pipelinestate-getcachedblob）和Vulkan（https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkPipelineCache.html）都提供了应用程序手动缓存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

https://blogs.igalia.com/itoral/2018/03/20/improving-shader-performance-with-vulkans-specialization-constants/

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

https://therealmjp.github.io/posts/shader-permutations-part1/

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

https://therealmjp.github.io/posts/shader-permutations-part2/

文末，再 次感谢 徐門子美 的分享， 作者主页： https://www.zhihu.com/people/BloodyGuys， 如果您有任何独到的见解或者发现也欢迎联系我们，一起探讨。（QQ群： 793972859 ）

近期精彩回顾