WaveCompactValue()について

最近、最適化で忙しいPocolです。
皆さん、お元気でしょうか?

今日は,WaveCompactValue()を勉強しようかなと思いましたので,そのメモを残しておこうと思います。
この関数は,[Drobot 2017]で紹介された手法です。

スライドに掲載されている実装は下記のよう感じです。

uint WaveCompactValue( uint checkValue )
{
    ulong mask; // lane unique compaction mask
    for ( ; ; ) // Loop until all active lanes removed
    {
        uint firstValue = WaveReadFirstLane( checkValue );
        mask = WaveBallot( firstValue == checkValue ); // mask is only updated for remaining active lanes
        if ( firstValue == checkValue ) break; // exclude all lanes with firstValue from next iteration
    }
    // At this point, each lane of mask should contain a bit mask of all other lanes with the same value.
    uint index = WavePrefixSum( mask ); // Note this is performed independently on a different mask for each lane.
    return index;
}

これをHLSLに書き直すと次のような感じになるかとおもいます。

uint WaveCompactValue(uint checkValue)
{
  // レーンのユニークなコンパクションマスク.
    uint4 mask;
    
    // すべてのアクティブレーンが取り除かれるまでループ.
    for (;;)
    {
        // アクティブレーンの最初の値を読み取る.
        uint firstValue = WaveReadLaneFirst(checkValue);

        // mask は残っているアクティブレーンに対してのみ更新される.
        mask = WaveActiveBallot(firstValue == checkValue);

        // firstValue を持つすべてのレーンを次のイテレーションから除外する。
        if (firstValue == checkValue)
             break;
    }
    // この時点で、マスクの各レーンは、同じ値を持つ他のレーンのすべてのビットマスクを含んでいなければならない。
    uint index = WavePrefixSum(mask); // これはレーンごとに異なるマスクで独立して行われる。
    return index;
}

さて,このWaveCompactValue()ですが,どういった使い道があるかというと,分類分けに使用することができます。
元々の[Drobot 2017]では色々なスレッドからAtomic操作をすると重くなるため,Atomic操作を減らす目的のために使われていました。
詳細な説明は,[Drobot 2017]のスライド51にアニメーション付きで載っていますので,そちらを参照してください。
軽く図の説明だけ載せておきます。




ちなみにグループ分けのよう番号を別途作りたい場合は,

uint2 WaveCompactValue(uint checkValue)
{
  // レーンのユニークなコンパクションマスク.
    uint4 mask;

    // グループ分け番号.
    uint groupIndex = 0;
    
    // すべてのアクティブレーンが取り除かれるまでループ.
    for (uint i=0; ; ++i)
    {
        // アクティブレーンの最初の値を読み取る.
        uint firstValue = WaveReadLaneFirst(checkValue);

        // mask は残っているアクティブレーンに対してのみ更新される.
        mask = WaveActiveBallot(firstValue == checkValue);

        // グループ分け番号を更新.
        groupIndex = i;

        // firstValue を持つすべてのレーンを次のイテレーションから除外する。
        if (firstValue == checkValue)
             break;
    }
    // この時点で、マスクの各レーンは、同じ値を持つ他のレーンのすべてのビットマスクを含んでいなければならない。
    uint index = WavePrefixSum(mask); // これはレーンごとに異なるマスクで独立して行われる。
    return uint2(index, groupIndex);
}

のように実装すると良いみたいです。
WaveCompactValue()はタイルの分類分けやマテリアルの分類分けなんかの場面で有効活用できそうな気がしています。
…というわけで,WaveCompactValue()を使って分類分けすれば,無駄なAtomic操作を減らせるので,高速化できるよ!という話でした。

参考文献

・[Drobot 2017] Michal Drobot, “Improved Culling for Tiled and Clustered Rendering Call of Duty Infinite Warfare”, SIGGRAPH 2017 Advances in Real-time Rendering and Games course, https://advances.realtimerendering.com/s2017/index.html