深入URP之Shader篇16: UNITY_BRANCH和UNITY_FLATTEN

Shader中的if分支

我们在shader中写if语句，例如：

if(a>0){
    //do some cool thing
}else{
    //do other cool thing
}

实际上，编译器会进行优化，以及处理成多种不同的情况。比如编译器会将if和else展开，分别执行其中的代码，然后根据测试条件选择一个结果，这被称为Flatten。或者编译器会真的编译出if指令，真正的动态分支。这取决于平台是否支持。比如在老的OpenGL ES 2.0平台，编译器总是进行Flatten，所以if分支才被视作性能杀手。当然即便是没有Flatten，真正的动态分支由于会让同一个Warp中的线程不能同步执行而互相等待，也会造成性能问题。在DX11中，编译器支持[branch]和[flatten]指令，这样就可以主动选择是flatten还是真动态分支了:

[branch] if(x)
{
    x = sqrt(x);
}

[flatten] if(x) 
{ 
    x = sqrt(x); 
}

具体参考：https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-if

• 另外，如果if中的测试条件是uniform，这样编译器也是可以安全的优化为只执行一个分支的代码的，效率和没有分支一样。

• 关于shader分支优化还有很多的说法，比如使用lerp，以及使用一些特殊的函数计算来代替分支，还有三元运算符等等，但是随着硬件平台和编译器的发展，一些方法已经不能优化，甚至反而是负优化了，这儿就不展开了（其实是我也没完全弄明白）。本文只想说一下Unity shader中，对于flatten和branch的处理。

UNITY_BRANCH和UNITY_FLATTEN

由于Unity Shader是跨平台的，不仅仅是DX系列的HLSL，还支持OpenGL, Metal, Vulkan等平台的Shader，因此对branch和flatten进行了包装，也就是UNITY_BRANCH和UNITY_FLATTEN。例如：

UNITY_BRANCH
if (screenPos.x < 0.5) {
// run expensive code A
} else {
// run expensive code B
}

UNITY_FLATTEN
if (screenPos.x < 0.5) {
// run expensive code A
} else {
// run expensive code B
}

• 使用UNITY_BRANCH时，if分支只会执行满足条件的那个分支，而使用UNITY_FLATTEN两种情况都需要执行。看起来我们应该总使用branch，但是branch指令本身也是有消耗的，如果if两边的操作都比较简单，使用branch反而可能得不偿失。

• 在OpenGL中，其实是没有branch指令的，编译器会自己选择是否flatten。以及上面说过的GLES 2.0，直接就只有flatten。所以这个指令只是一个hint，具体还是要看平台。

branch不能使用的一种情况

fixed4 color = fixed4(0,0,0,0);
UNITY_BRANCH
if (screenPos.x > 0.5) {
    float2 uv = screenPos * 2.0;
    color = tex2D(_MyTex, uv);
}

这儿调用tex2D会产生错误，因为 tex2D 会使用uv的导数(即一个像素到另一个像素的变化值）去决定mip map的级别，但是由于这儿的 uv 只在if分支中计算，这不能保证可以计算出导数（因为导数需要邻近像素的数据）。修改方法如下：

fixed4 color = fixed4(0,0,0,0);
float2 uv = screenPos * 2.0;
if (screenPos.x > 0.5) {
    color = tex2D(_MyTex, uv);
}

URP中使用UNITY_BRANCH的地方

color grading

        UNITY_BRANCH
        if (userLutContrib > 0.0)
        {
            input.rgb = LinearToSRGB(input.rgb); // In LDR do the lookup in sRGB for the user LUT
            half3 outLut = ApplyLut2D(TEXTURE2D_ARGS(userLutTex, userLutSampler), input, userLutParams);
            input = lerp(input, outLut, userLutContrib);
            input.rgb = SRGBToLinear(input.rgb);
        }

因为if中需要采样LUT贴图，如果不是动态分支，不使用 userLut的情况下也会执行这个代码，代价就太大了。

附加光realtime阴影

    // We have to branch here as otherwise we would sample buffer with lightIndex == -1.
    // However this should be ok for platforms that store light in SSBO.
    UNITY_BRANCH
    if (lightIndex < 0)
        return 1.0;

    float4 shadowCoord = mul(_AdditionalShadowsBuffer[lightIndex].worldToShadowMatrix, float4(positionWS, 1.0));
#else
    float4 shadowCoord = mul(_AdditionalLightsWorldToShadow[lightIndex], float4(positionWS, 1.0));
#endif

    half4 shadowParams = GetAdditionalLightShadowParams(lightIndex);
    return SampleShadowmap(TEXTURE2D_ARGS(_AdditionalLightsShadowmapTexture, sampler_AdditionalLightsShadowmapTexture), shadowCoord, shadowSamplingData, shadowParams, true);

这个代码注释也说了，必须branch，否则就会多采样shadow map。

参考资料

• https://forum.unity.com/threads/correct-use-of-unity_branch.476804/
• https://zhuanlan.zhihu.com/p/122467342
• https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-if
• https://zhuanlan.zhihu.com/p/525931936