「Shaderでif文を使ったら遅い」は正しくない🧐

前書き

この記事は、2023のUnityアドカレの12/18の記事です。
今年は、完走賞に挑戦してみたいと思います。Qiita君ぬい欲しい！

あと8記事、1/3……頑張るゾイ！！🔥🔥🔥

はじめに

Shaderを書くときによく、「if文を使うな！」と言われます。それは本当でしょうか？（タイトルでオチがついていますが…）

今回は、その根拠に迫りたいと思います。

分岐の種類

GPUは複数のスレッドが一つの命令制御機の元で並列に動いており、スレッドごとに別の命令を実行させることができません。ゆえに、分岐Aと分岐Bがあった場合に、分岐Aのスレッドが終わってから、分岐Bのスレッドを実行し、自分の番ではない間は待機させられるということです。

それでは、スレッド全員がAに行く場合はどうなるのでしょうか？Bの処理は誰一人、必要としていないはずですが、実行されてしまうのでしょうか？

Static分岐（静的）

コンパイル時に分岐させてしまう方法です。要は、分岐Aの場合のShaderと、分岐Bの場合のShaderを両方用意しておき、CPUからどちらのShaderを使うかを切り替えてあげるという戦略です。分岐しているのはCPUであり、GPUは実は分岐していません。ですから、分岐のペナルティは皆無であるということは、言うまでもないでしょう。

Unityの場合は、shader_featureやmulti_compileなどで指定するShaderKeywords機能として提供されています。

#pragma multi_compile BRANHC_A BRNACH_B

#if defined(BRABCH_A)
    return 0;
#elif defined(BRANCH_B)
    return 1;
#endif

プリプロセッサなので、可読性に難があります。分岐の組み合わせパターン数だけShaderを出力するので、Shaderファイルが巨大化しやすいという問題を孕んでいます。調子に乗ってmulti_compileを増やしたりしていると、組み合わせは10万、100万に上り、数百MBや数GBになってしまったShaderファイルも見てきました。

Divegentな動的分岐（発散的）

動的分岐は普通のif文です。

分岐条件に使う変数（conditinal）が、スレッドによって異なる場合がある場合、「発散」するので、divegentと呼ばれます。

if(vertexData.position.x > 100)
{
    return 0;
}
else
{
    return 1;
}

Unformな動的分岐（統一的）

こちらはdivegentと異なり、uniformな変数で分岐をします。HLSLではuniformな変数というのは、ConstantBufferのことです。すべてのスレッドが同じ方向（AorB）に分岐するかが、Shaderをコア上で実行する前にわかるはずです。

if(constantBufferData.hoge > 100)
{
    return 0;
}
else
{
    return 1;
}

一次資料

公式の見解を見てみましょう。

しかし、「分岐がよくない」ということは述べているものの、内部的な振る舞いについて具体的な言及がされた資料は見つかりませんでした。GPUのモデルによっても変わってくるため、明言したくないといったところでしょうか。

NVIDIAによるPTX（Parallel Thread Execution）命令セットのドキュメント

これは、NVIDIAのCUDAの吐き出す中間命令（PTX）のアセンブリのドキュメントです。

.uniサフィックスをつけていないフロー制御命令（if文のbraなど）は分岐である…的なことが書かれています。.uniというのは、すべてのスレッドが同じ方向にすすむ場合、フロー制御命令につけることができる属性です。つまりuniformな動的分岐であることを示すものです。

「divegentが良くないぜ」って主張はわかるのですが、uniformなら良いのかは曖昧です。

7.5. Divergence of Threads in Control Constructs

Threads in a CTA execute together, at least in appearance, until they come to a conditional control construct such as a conditional branch, conditional function call, or conditional return. If threads execute down different control flow paths, the threads are called divergent.
If all of the threads act in unison and follow a single control flow path, the threads are called uniform. Both situations occur often in programs.

A CTA with divergent threads may have lower performance than a CTA with uniformly executing threads, so it is important to have divergent threads re-converge as soon as possible. All control constructs are assumed to be divergent points unless the control-flow instruction is marked as uniform, using the .uni suffix. For divergent control flow, the optimizing code generator automatically determines points of re-convergence. Therefore, a compiler or code author targeting PTX can ignore the issue of divergent threads, but has the opportunity to improve performance by marking branch points as uniform when the compiler or author can guarantee that the branch point is non-divergent.

Qualcomm® Adreno™ GPU Best Practices Shaders

国内スマホの大半が積んでいるSnapdragonのメーカーの最適化ドキュメントです。

パフォーマンスの良い順に、Static分岐（Constant）、Uniform分岐、Divegent分岐で、static分岐なら許容範囲のパフォーマンスが得られるというようなことを言っています。

静的分岐がパフォーマンス的に問題を起こさないのは、それはそうに決まっているのですが、uniformやdivegentでどんなハードウェア的差が生じるのか、影響度合いについては語られていません。

Threads in flight/dynamic branching

Branching is crucial for the performance of the shader. Every time the branch encounters divergence, or where some elements of the thread branch one way and some elements branch in another, both branches will be taken with predicates using NULL out operations for the elements that do not take a given branch. This is true only if the data is aligned in a way that follows those conditions, which is rarely the case for fragment shaders. There are three types of branches, listed in order from best performance to worst on Adreno GPUs:

• Branching on a constant, known at compile time
• Branching on a uniform variable
• Branching on a variable modified inside the shader

Branching on a constant may yield acceptable performance.

三項演算cond?a:bなら使っていいという話

ifはだめだから、三項演算を使う…というコードを時々見かけます。

三項演算は、aとbをそれぞれ評価してから、条件を見てどちらかを選ぶという命令に展開されます。結局両方やっているわけですね。if文には[branch]/[flatten]という属性があり、flattenをつけると、三項演算と同じ両方評価、branchをつけると正真正銘の分岐になります。どちらもつけないと、コンパイラのゴキゲン次第ですが、DXCだとflattenの方に行くっぽいですね。さすがに頻用はしないと思うので、明示しておくのがよいかと思います。

分岐で待機を作るよりも、両方評価したほうが早いという可能性はありうるので、使いどころはありますが、可読性に問題がある場合は、flattenを明示的につけたif文にしましょう。

計測

測っちゃうのが早かろうということで、ベンチマークをとってみましょう。

実装

以下のような処理を、condをStatic（完全定数）/uniform/divegent（頂点データ）から与えるようにして、変化を見ます。

[branch] // or [flatten]
if(cond)
    // Fastパス
else
    // Slowパス

実装コード（C#/HLSL）

trueパスを軽い処理、falseパスを重い処理としたComputeShaderを用意しました。

_Input配列から、Expを計算し、_Output配列に保存します。

StructuredBuffer<float> _Input;
RWStructuredBuffer<float> _Output;
uint _Uniform;
uint _DivergentFastInterval;

float CalcExp(float x);

[numthreads(32, 1, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
    bool fast;
    {
#if defined(PATH_STATIC_FAST)
        fast = true;
#elif defined(PATH_STATIC_SLOW)
        fast = false;
#endif

#if defined(PATH_UNIFORM)
        fast = _Uniform == 0;
#endif

#if defined(PATH_DIVERGENT)
        // _DivergentFastIntervalスレッドに一つだけSlowパスを通す
        fast = id.x % _DivergentFastInterval != (_DivergentFastInterval - 1);
#endif
    }

    float v = _Input[id.x];
    {
#if defined(BRANCH)
        [branch]
#elif defined(FLATTEN)
        [flatten]
#endif
        if(fast)
            v = exp(v); // fast
        else
            v = CalcExp(v); // slow
    }
    _Output[id.x] = v;
}

重たい処理として、10000次までのマクローリン展開でexp計算をする関数を用意しました。

float CalcExp(float x)
{
    float powX = 1;
    float denom = 1;
    float e_0 = 1;
    float e = e_0;
    for(uint i = 1; i < 10000 < i++)
    {
        powX *= x;
        denom *= i;
        e = e0 * powX / denom;
    }
    return e;
}

組み合わせ呼び出しができるように、カーネルを宣言します。static/uniform/divegentと、Static分岐はFastパス/Slowパス、そのほかはBranch/Flattenの組み合わせを網羅します。

#pragma kernel STATIC_FAST CSMain=STATIC_FAST PATH_STATIC_FAST
#pragma kernel STATIC_SLOW CSMain=STATIC_SLOW PATH_STATIC_SLOW
#pragma kernel UNIFORM_BRANCH CSMain=UNIFORM_BRANCH PATH_UNIFORM BRANCH
#pragma kernel UNIFORM_FLATTEN CSMain=UNIFORM_FLATTEN PATH_UNIFORM FLATTEN
#pragma kernel DIVERGENT_BRANCH CSMain=DIVERGENT_BRANCH PATH_DIVERGENT BRANCH
#pragma kernel DIVERGENT_FLATTEN CSMain=DIVERGENT_FLATTEN PATH_DIVERGENT FLATTEN

C#側で、組み合わせを定義します。

enum BraType { STATIC, UNIFORM, DIVERGENT, }
enum BranchType { FAST, SLOW, FLATTEN_FAST, }
class Entry
{
    public CustomSampler samp;
    public List<float> history = new List<float>();
    public float us = -1;
    public Entry(string proName) => samp = CustomSampler.Create(proName, collectGpuData: true);
}
Dictionary<(BraType mode, BranchType path), Entry> averageUs = new ()
{
    [(BraType.STATIC, BranchType.FAST)] = new ("STATIC_FAST"),
    [(BraType.STATIC, BranchType.SLOW)] = new ("STATIC_SLOW"),
    [(BraType.UNIFORM, BranchType.FAST)] = new ("UNIFORM_FAST"),
    [(BraType.UNIFORM, BranchType.SLOW)] = new ("UNIFORM_SLOW"),
    [(BraType.UNIFORM, BranchType.FLATTEN_FAST)] = new ("UNIFORM_FLATTEN_FAST"),
    [(BraType.DIVERGENT, BranchType.FAST)] = new ("DIVERGENT_FAST"),
    [(BraType.DIVERGENT, BranchType.SLOW)] = new ("DIVERGENT_SLOW"),
    [(BraType.DIVERGENT, BranchType.FLATTEN_FAST)] = new ("DIVERGENT_FLATTEN_FAST"),
};

入力バッファを用意します。

const int N = 16 * 1024;

if(input == null || output == null
    || input.count != N || output.count != N)
{
    input?.Dispose();
    output?.Dispose();
    input = new GraphicsBuffer(GraphicsBuffer.Target.Structured, N, sizeof(float));
    input.SetData(Enumerable.Range(0, N)
        .Select(i => (float)i / N)
        .ToArray());
    output = new GraphicsBuffer(GraphicsBuffer.Target.Structured, N, sizeof(float));
}

プロファイリングはCustomSamplerを使い、コマンドバッファに積み流します。

foreach (var ((type, branchType), entry) in averageUs)
{
    var kernel = type switch
    {
        BraType.STATIC => branchType switch
        {
            BranchType.FAST => cs.FindKernel("STATIC_FAST"),
            BranchType.SLOW => cs.FindKernel("STATIC_SLOW"),
            BranchType.FLATTEN_FAST => throw new NotImplementedException(),
            _ => throw new NotImplementedException(),
        },
        BraType.UNIFORM => branchType switch
        {
            BranchType.FLATTEN_FAST => cs.FindKernel("UNIFORM_FLATTEN"),
            _ => cs.FindKernel("UNIFORM_BRANCH"),
        },
        BraType.DIVERGENT => branchType switch
        {
            BranchType.FLATTEN_FAST => cs.FindKernel("DIVERGENT_FLATTEN"),
            _ => cs.FindKernel("DIVERGENT_BRANCH"),
        },
        _ => throw new NotImplementedException(),
    };
    cmd.SetComputeIntParam(cs, "_Uniform", branchType == BranchType.FAST ? 0 : 1);
    cmd.SetComputeIntParam(cs, "_DivergentFastInterval", branchType == BranchType.FAST ? 1000000000 : Slow_DivergentFastInterval);
    cmd.SetComputeBufferParam(cs, kernel, "_Input", input);
    cmd.SetComputeBufferParam(cs, kernel, "_Output", output);
    cmd.BeginSample(entry.samp);
    cmd.DispatchCompute(cs, kernel, N / 32, 1, 1);
    cmd.EndSample(entry.samp);
            
    entry.history.Add(entry.samp.GetRecorder().gpuElapsedNanoseconds / 1000f);
    if(entry.history.Count >= 60)
    {
        entry.us = SystemInfo.supportsGpuRecorder ? entry.history.Average() : -1;
        entry.history.Clear();
    }
}
Graphics.ExecuteCommandBuffer(cmd);

プロジェクト

結果

以下は、60フレームのGPU平均所要時間です。単位はμsです

Slowパスでは、当然ほぼ差はありませんでした。

GPU	API	static	uniform	divergent
RTX 4070 ti	DX11	195.21	195.43	195.38
Adreno730（SM8475）	Vulkan	8516.12	8548.35	6741.53
Adreno619（SM6375）	Vulkan	56793.53	50966.15	22148.02

Fastパスでは、uniformだけでなく、divegentであってもstaticと同等近くまで速くなりました。（Branch属性）また、Flatten属性を使用した場合、FastパスであってもSlowパスとほぼ同じだけの時間がかかってしまいました。

GPU	API	static	uniform	divegent	uniform (Flatten)	divegent (Flatten)
RTX 4070 ti	DX11	3.60	4.51	4.01	195.55	193.83
Adreno730 (Snapdragon8+Gen1)	Vulkan	10.77	10.78	11.05	6796.79	6554.32
Adreno619 (Snapdragon 695)	Vulkan	15.74	18.71	18.13	20966.15	22148.02

ちなみに、計測結果に入れていませんが、divegentで交互のスレッドでSlowとFastのパスを選ばせた場合（グループに混在）Slowパスと同じだけ遅くなりました。

考察

uniform及び、divegentのFastパスの結果が、staticのFastパスの結果とほとんど変わらないことから、すべてのスレッドが同じパスを通る場合、もう一方のパスは「待機」ではなく「スキップ」されるということが確かめられました。これは、divegentの結果から、コンパイラレベルの最適化ではなく、実行時にスレッドの命令をつかさどる部分で管理下スレッド全体の分岐状況を見て、待機させるか全員でジャンプするかを判断しているということが伺えます。

気になるのは、Slowパスにおいて、なぜかdivegentが最も速く、FlattenもStaticのSlowパスより速いという点です。計測誤差かとも思いましたが、何度図りなおしてもこれが逆転することはありませんでした。

まとめ

結論としては、現代においてGPU上でのif文はそこまで目の敵にする必要はないと言えます。少なくとも、static分岐によるUberShaderの肥大化が問題になっている場合、HLSLによる分岐を適切に使うことをお勧めします。

スレッドグループ内で、スレッドごとに別々のパスへ進んでしまう場合は遅くなってしまいますが、そもそもGPUで処理するデータ自体、近所のスレッド（≒近所のピクセル）の類似性を前提とした処理であるしょう。境界部分のグループの犠牲は仕方ないとしても全体として偏りのない分岐になってしまう場面は少ないのであろうかと思います。