第4章. 面向移动端的着色器优化：超越函数的测量与轻量化

Shader "Benchmark/Transcendental"
{
    Properties
    {
        // uGUI（CanvasRenderer）会把纹理绑定到 "_MainTex"，因此名称固定
        [PerRendererData][MainTexture] _MainTex("Sprite Texture", 2D) = "white" {}
        [MainColor] _BaseColor("Tint", Color) = (1, 1, 1, 1)

        // 用于 Mask（UI/Mask）的模板设置。Mask 组件会替换 Material 并写入
        _StencilComp("Stencil Comparison", Float) = 8
        _Stencil("Stencil ID", Float) = 0
        _StencilOp("Stencil Operation", Float) = 0
        _StencilWriteMask("Stencil Write Mask", Float) = 255
        _StencilReadMask("Stencil Read Mask", Float) = 255

        _ColorMask("Color Mask", Float) = 15

        [Toggle(UNITY_UI_ALPHACLIP)] _UseUIAlphaClip("Use Alpha Clip", Float) = 0
    }

    SubShader
    {
        Tags
        {
            "Queue" = "Transparent"
            "RenderType" = "Transparent"
            "RenderPipeline" = "UniversalPipeline"
            "IgnoreProjector" = "True"
            "PreviewType" = "Plane"
            "CanUseSpriteAtlas" = "True"
        }

        Stencil
        {
            Ref [_Stencil]
            Comp [_StencilComp]
            Pass [_StencilOp]
            ReadMask [_StencilReadMask]
            WriteMask [_StencilWriteMask]
        }

        Cull Off
        ZWrite Off
        ZTest [unity_GUIZTestMode]
        Blend SrcAlpha OneMinusSrcAlpha
        ColorMask [_ColorMask]

        Pass
        {
            HLSLPROGRAM

            #pragma vertex vert
            #pragma fragment frag
            // 循环次数是 uniform（动态循环），因此需要 3.0 而不是 2.0
            #pragma target 3.0

            // 支持 RectMask2D / Mask 的 Alpha Clip
            #pragma multi_compile_local _ UNITY_UI_CLIP_RECT
            #pragma multi_compile_local _ UNITY_UI_ALPHACLIP

            // 切换要测量的超越函数。由于希望在运行时通过 Shader.EnableKeyword 切换，
            // 这里使用 multi_compile 而不是 shader_feature（未使用的 variant 会被剥离）。
            // 每个关键字对应一个独立的 variant，因此生成的代码中只会出现所选的函数（没有分支开销）。
            // 由于是全局关键字，可以从 TranscendentalBenchmark 一侧统一控制。
            // 如果没有启用任何关键字，则使用开头的 _FUNC_NONE variant。
            #pragma multi_compile_fragment _FUNC_NONE _FUNC_SIN _FUNC_SIN_APPROX _FUNC_COS _FUNC_TAN _FUNC_TAN_TAYLOR _FUNC_TAN_PADE _FUNC_ASIN _FUNC_ACOS _FUNC_ATAN _FUNC_EXP _FUNC_EXP2 _FUNC_POW _FUNC_LOG _FUNC_LOG2 _FUNC_LOG10

            #include "Packages/com.unity.render-pipelines.universal/ShaderLibrary/Core.hlsl"

            struct Attributes
            {
                float4 positionOS : POSITION;
                float4 color : COLOR;
                float2 uv : TEXCOORD0;
                UNITY_VERTEX_INPUT_INSTANCE_ID
            };

            struct Varyings
            {
                float4 positionHCS : SV_POSITION;
                float2 uv : TEXCOORD0;
                float4 positionOS : TEXCOORD1; // 以与 _ClipRect 相同的空间（Canvas 本地空间）保存
            };

            TEXTURE2D(_MainTex);
            SAMPLER(sampler_MainTex);

            CBUFFER_START(UnityPerMaterial)
                half4 _BaseColor;
            CBUFFER_END

            // 由 TranscendentalBenchmark 设置的全局值。
            // 不是按材质的值，因此不要放进 UnityPerMaterial（放进去会破坏 SRP Batcher）
            int _TranscendentalLoopCount;
            float _PowExponent;

            // 相当于 UnityUI.cginc 的 UnityGet2DClipping（用 HLSL 自行定义）
            float Get2DClipping(float2 position, float4 clipRect)
            {
                float2 inside = step(clipRect.xy, position.xy) * step(position.xy, clipRect.zw);
                return inside.x * inside.y;
            }

            // 将传给超越函数的值限制在 (0.05, 0.95) 范围内。
            // 一并规避 asin/acos 的定义域 |x| <= 1、log 系的定义域 x > 0，以及 exp 系的溢出。
            // 由于这是所有 variant 共用的前处理，其成本会在求差值时被抵消。
            float PrepareInput(float seed)
            {
                return seed;
                // return frac(seed) * 0.9 + 0.05;
            }

            float sinApprox(float x)
            {
                // sin 的抛物线近似。只需乘法和 abs，用来观察与真正的 sin 指令之间的成本差异。
                // 定义域为 |x| <= PI（PrepareInput 的输出落在此范围内）。
                // 系数会被常量折叠，因此运行时没有除法开销。
                return (4.0 / PI) * x - (4.0 / (PI * PI)) * x * abs(x);
            }

            float tanTaylor(float x)
            {
                // tan 的泰勒展开取到 3 次项（仅乘法）。用来观察与真正的 tan 指令之间的成本差异。
                // 定义域为 |x| < PI/2，超过 π/4 后精度会迅速下降。
                float x2 = x * x;
                return x * (1.0 + x2 * 0.3333333); // 1.0/3.0
            }

            float tanPade(float x)
            {
                // tan 的帕德近似（乘法 + 1 次除法）。定义域为 |x| < PI/2。
                // 分母的零点 x ≈ ±1.581 靠近 tan 的极点 ±PI/2，因此在极点附近也能跟随。
                float x2 = x * x;
                return x * (15.0 - x2) / (15.0 - 6.0 * x2);
            }

            // 被测量的超越函数。每个关键字只保留一个，
            // 因此生成的代码中只会出现所选函数的展开结果。
            float Transcendental(float x)
            {
                #if defined(_FUNC_SIN)
                    return sin(x);
                #elif defined(_FUNC_SIN_APPROX)
                    return sinApprox(x);
                #elif defined(_FUNC_COS)
                    return cos(x);
                #elif defined(_FUNC_TAN)
                    return tan(x);
                #elif defined(_FUNC_TAN_TAYLOR)
                    return tanTaylor(x);
                #elif defined(_FUNC_TAN_PADE)
                    return tanPade(x);
                #elif defined(_FUNC_ASIN)
                    return asin(x);
                #elif defined(_FUNC_ACOS)
                    return acos(x);
                #elif defined(_FUNC_ATAN)
                    return atan(x);
                #elif defined(_FUNC_EXP)
                    return exp(x);
                #elif defined(_FUNC_EXP2)
                    return exp2(x);
                #elif defined(_FUNC_POW)
                    // 如果把指数设为常量，会被折叠成 pow(x, 2) -> x * x 的形式，
                    // 因此传入 uniform，让其走通用的 pow 路径
                    return pow(x, _PowExponent);
                #elif defined(_FUNC_LOG)
                    return log(x);
                #elif defined(_FUNC_LOG2)
                    return log2(x);
                #elif defined(_FUNC_LOG10)
                    return log10(x);
                #else // _FUNC_NONE：不含超越函数的基线
                    return x;
                #endif
            }

            Varyings vert(Attributes IN)
            {
                Varyings OUT;
                OUT.positionOS = IN.positionOS;
                OUT.positionHCS = TransformObjectToHClip(IN.positionOS.xyz);
                OUT.uv = IN.uv * 300.0;
                return OUT;
            }

            half4 frag(Varyings IN) : SV_Target
            {
                // 这里是测量对象（显示效果没有意义）。
                // 用 UNITY_LOOP 抑制循环展开，使 _TranscendentalLoopCount 次的成本原样累加。
                // 除超越函数的内容之外，所有 variant 的形式完全相同，
                // 因此与 _FUNC_NONE 的差值就是该函数的成本。
                float accum = 0.0;

                UNITY_LOOP
                for (int i = 0; i < _TranscendentalLoopCount; i++)
                {
                    // 如果写成循环不变的表达式，会被外提到循环外，因此把 i 混入参数
                    float x = IN.uv.x + i * 0.017;

                    accum += Transcendental(x);
                }
                return half4(saturate(accum.xxx) * 0.02, 1);
            }
            ENDHLSL
        }
    }
}

float sinApprox = (4.0 / PI) * x - (4.0 / (PI * PI)) * x * abs(x);

float tanTaylor(float x) {
    float x2 = x * x;
    return x * (1.0 + x2 * 0.3333333);
}

float tanPade(float x) {
    float x2 = x * x;
    return x * (15.0 - x2) / (15.0 - 6.0 * x2);
}