Scan with Decoupled Look-back and Onesweep Radix Sort

レイトレ合宿10 ( https://sites.google.com/view/rtcamp10 )のセミナーで発表した資料です。

Decoupled Look-backと呼ぶテクニックを用いることで並列性を活かしつつVRAMアクセス量を削減したスキャンの実装、そしてそれに基づくRadix Sortの新手法Onesweepについて紹介しています。

Bluesky: @bsky.rayspace.xyz
X/Twitter: https://twitter.com/Shocker_0x15

Transcript

1. Scan with Decoupled Look-back & Onesweep Radix Sort レイトレ合宿10 https://sites.google.com/view/rtcamp10/home @Shocker_0x15 @bsky.rayspace.xyz 1
2. 2 1. 並列プリミティブ 2. Scan with Decoupled Look-back 3. Onesweep Radix Sort
3. 3 1. 並列プリミティブ 2. Scan with Decoupled Look-back 3. Onesweep Radix Sort
4. 36 並列プリミティブ(Parallel Primitives)とは リダクションやスキャン、ソートといった汎用的な計算を並列計算を活用して実装したもの リダクション (Reduce, Reduction) 配列要素全体に対して1つ結果を返す 7 2 5 8 1 3 4 6 36 Sum スキャン (Scan, Prefix Scan) 配列要素それぞれにおいて自身より前の要素列に対するリダクション 7 2 5 8 1 3 4 6 0 7 9 14 22 23 26 30 7 2 5 8 1 3 4 6 8 Max ExclusivePrefixSum 7 2 5 8 1 3 4 6 7 2 2 2 1 1 1 1 InclusivePrefixMin ソート (Sort) 配列要素全体を何らかの大小関係で並べ替える 7 2 5 8 1 3 4 6 Ascending 1 2 3 4 5 6 7 8 4
5. 並列プリミティブの応用 5 例: スレッドごとに動的に決まる数の要素を配列に連続して格納したい 3 1 0 0 4 2 1 1 要素数 0 1 2 3 4 5 6 7 スレッド ExclusivePrefixSum 0 3 4 4 4 8 10 11 オフセット 要素の配列 例: Linear BVH 三角形などのプリミティブの位置をモートンコード化、 1次元の配列としてソートすることで隣接要素が空間的に近傍になる
6. 並列プリミティブの実装・使用上の注意点 6 結合性 (Associativity) A ⊕ (B ⊕ C) = (A ⊕ B) ⊕ C ?? (⊕は何らかの二項演算子) 例: 浮動小数点数はこれを満たさない (Pseudo-Associative) 可換性 (Commutativity) A ⊕ B = B ⊕ A ?? (⊕は何らかの二項演算子) 例: 減算や除算、ベクトルの外積、行列の積などはこれを満たさない 決定性 (Determinism) 実行の度や計算デバイスを変えた際に結果が一致するか否か 結合性にも強く依存する Run-to-Run Determinism: 実行の度に結果が変わらない保証 プリミティブで扱う対象は単純な数値以外にも拡張することができる 例: ベクトル型、AABBなど リダクションやスキャンにおいては主に二項演算子を定義することで任意の型に対応する
7. 7 1. 並列プリミティブ 2. Scan with Decoupled Look-back 3. Onesweep Radix Sort
8. + + + + + + + + + + + + + + スレッドグループ内リダクション・スキャン実装例 GPGPUにおける並列プリミティブの実装は階層的 スレッドグループやWaveなどの小さい範囲ではLDSを活用して高速化 7 2 5 8 1 3 4 6 7 9 5 13 1 4 4 10 7 9 5 22 1 4 4 14 7 9 5 22 1 4 4 36 7 2 5 8 1 3 4 6 7 9 5 22 1 4 4 0 7 9 5 0 1 4 4 22 7 0 5 9 1 22 4 26 0 7 9 14 22 23 26 30 0 7 9 14 22 23 26 30 36 LDS リダクション結果 (リダクションだけならAtomicもあり、ただし結合性に注意) スキャン結果 (Exclusive Prefix Sum) Copy In/Out ネイティブ型(intやfloat)に関してはWave内命令(Wave Intrinsics)が最初から用意されていたりする カスタム型に関しては自分で実装 Barrier/Sync 0 (or Offset) 8
9. 9 従来手法1: Reduce-Then-Scan VRAMへのアクセスが入力サイズnに対しておおよそ3n (2n reads, 1n writes) • 入力を”パーティション”に分割 • パーティション数 < 1つのスレッドグループで高速に扱えるサイズ • パーティションごとのリダクション、別バッファーに結果を書き込む • パーティションあたりスレッドグループが何度かループ • リダクション結果のスキャン、パーティションごとのオフセットを得る • パーティションごとのスキャン、オフセットとスキャン結果を加算(⊕)して結果出力 • パーティションあたりスレッドグループが何度かループ Input Buffer Dispatch 1: Partition-wise Reduction Reduction Buffer Dispatch 2: Scan b 1 b 2 b 3 b 4 b B-1 b B p 1 p 2 p P Input Buffer p 1 p 2 p P Dispatch 3: Partition-wise Scan + Offset
10. 10 従来手法2: Chained Scan VRAMへのアクセスが入力サイズnに対して2n (1n reads, 1n writes) 後段のパーティションは前段のパーティションが処理完了するまで絶対に待ち続ける • パーティションごとにスキャン • 前パーティションまでのリダクション結果(Inclusive Prefix)が届くまで待ち、 自パーティションのリダクションと加算(⊕)して次パーティションに知らせる • 前パーティションまでのリダクションと自パーティションのスキャン結果を足して結果出力 • 遅延を抑えるためパーティションサイズ=スレッドグループサイズ Input Buffer Dispatch 1: Partition-wise Scan + Wait for Previous Partition’s Inclusive Prefix
11. 11 Decoupled Look-back (1) • Chained Scanは直前パーティションのInclusive Prefixに依存 Part p-1 Part p Part p+1 Part p+2 • リダクションと Inclusive Prefix計算済 • グローバルスキャン 計算済 • リダクションと Inclusive Prefix計算中 • グローバルスキャン 計算可能 • リダクションは計算可 • Inclusive Prefixと グローバルスキャンは 前ブロック待ち • リダクションは計算可 • Inclusive Prefixと グローバルスキャンは 前ブロック待ち • パーティションごとのリダクションは即座に計算できる • パーティションごとに欲しいのはPartition Exclusive Prefixのみ →パーティションをさらに遡れば直前パーティションのInclusive Prefixに非依存で計算できる Part p-1 Part p Part p+1 Part p+2 • リダクション計算済 • Inclusive Prefix計算済 • グローバルスキャン 計算可能済 • リダクション計算済 • Inclusive Prefix計算中 • グローバルスキャン 計算可能 • リダクション計算済 • Inclusive Prefix 計算可能 • グローバルスキャン 計算可能 • リダクション計算済 • Inclusive Prefix 計算可能 • グローバルスキャン 計算可能
12. 12 Decoupled Look-back (2) パーティションごとに”パーティションデスクリプター”を持つ Decoupled Look-backの手順 1. パーティションのExclusive Prefixを初期化(例: ゼロ) 2. パーティションのリダクション計算、デスクリプターに値を書き込みステータスをRに更新 (最初のパーティションのみ書き込み先はInclusive PrefixかつステータスはP) 3. 前パーティションのステータスがXでなくなるまで待つ a. ステータスがPなら前パーティションのデスクリプターのInclusive Prefix値を読み出し Exclusive Prefixに累積(⊕)して終了 b. ステータスがRなら前パーティションのデスクリプターのリダクション値を読み出し Exclusive Prefixに累積(⊕)、さらに前のパーティションに関してステップ3を繰り返す 4. Inclusive Prefix計算(Exclusive Prefix⊕自パーティションのリダクション値) 、 デスクリプターに値を書き込みステータスをPに更新 • パーティションのリダクション値 • パーティションのInclusive Prefix値 (前パーティションのInclusive Prefix ⊕ 自リダクション) • ステータスフラグ • X: 初期値 • → R: リダクション計算済み • → P: Inclusive Prefix計算済み
13. 13 Decoupled Look-back (3) Decoupled Look-backが完了次第、 パーティション内スキャンとExclusive Prefixの演算結果をバッファーに書き出せば処理完了 Prefix Sumなどの単純な二項演算であればスキャンはかなりメモリバウンドな処理 Reduce-Then-Scan (3nアクセス)に比べて問題サイズがそれなりになれば Decoupled Look-back Scan (2nアクセス)は1.5倍程度高速、単なるmemcpyにも匹敵する
14. 14 実装・使用上の注意点 • GPUカーネルのスレッドグループ番号を直接パーティション位置に対応付けないこと • 若い番号のスレッドグループが先に起動・実行されるとは限らない • アトミックカウンターを使って早く起動したスレッドグループから順に先頭から割り当てる • デスクリプターの読み書きはGPU全体で共通のメモリレベルに行う必要がある • CUDA: volatile, HLSL: globallycoherent, GLSL: coherent • デスクリプターへの値の書き込みとステータスフラグの更新の間にはメモリフェンスが必要 • プログラム上で「値→ステータス」の順で書いてもフェンスが無いと ステータスの更新のほうが他パーティションに対して先に可視になるかもしれない • CUDA: __threadfence(), HLSL: DeviceMemoryBarrier(), GLSL: memoryBarrierBuffer() • float値の加算などNon-Associativeな二項演算の場合Run-to-Run Determinismも保証されない
15. 15 最適化: SIMD-parallelized Look-back 単純なLook-backはパーティション全体でひとつ前のパーティションだけチェックする → SIMDを活用して複数の前パーティションをまとめてチェック P P P R R R R X p-8 p-7 p-6 p-5 p-4 p-3 p-2 p-1 p Partition Thread 7 6 5 4 3 2 1 0 スレッドごとにひとつのデスクリプターのステータスを読む • ひとつでもXがあれば待機 • すべてRなら Exclusive Prefixにすべてのリダクション値を累積してさらに前方の複数パーティションを読む • ひとつでもPがあれば 直近のPに該当するInclusive PrefixとそこまでのRに該当するリダクション値の累積 Exclusive Prefixにセットして終了 (これもWave内命令を使って高速に処理できる) P P P R R R R R p-8 p-7 p-6 p-5 p-4 p-3 p-2 p-1 p Partition Thread 7 6 5 4 3 2 1 0 (p-6の値) + (p-5の値) + … + (p-1の値) = (pにおけるExclusive Prefix)
16. 16 「値のサイズ + ステータスフラグサイズ」がサポートされるメモリ命令の幅に収まるなら 「値の更新→メモリフェンス→ステータスの更新」の代わりに一回のメモリ命令に集約できる 最適化: フェンス不要なデスクリプター更新 リダクション値とInclusive Prefix値は 同時には必要にならないため重ねられる
17. 17 応用 Decoupled Look-back Scanによりシングルパスでコンパクションを行える スキャン結果をバッファーに直接書き出す以外にも様々な応用の高速化に寄与 • Select-If • trueとして評価される要素のみ書き出す • Partition-If • true/falseの評価値に応じてバッファーの内容を前後に振り分ける • Reduce-By-Key • キー値ペアの配列を受け取り、同じキーごとにリダクション • Run-Length-Encode • 配列を先頭から見て「◯◯がa個、✕ ✕がb個…」といったエンコードで出力
18. 18 1. 並列プリミティブ 2. Scan with Decoupled Look-back 3. Onesweep Radix Sort
19. Radix Sort 数値をd-bitのRadix Digitに分割して位ごとのソートを繰り返す LSD (Least Significant Digit) Radix Sort: 下の位から MSD (Most Significant Digit) Radix Sort: 上の位から 71 231 5 18 51 162 32 127 例: 8-bitの数値を2-bit LSD Radix Sort Dec LSD 1 1 1 0 1 0 0 1 1 1 0 1 1 1 0 1 1 1 1 0 0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 1 1 0 0 1 0 1 0 0 1 1 1 1 1 1 0 0 0 0 0 1 0 1 0 0 0 1 0 0 MSD 32 5 18 162 71 231 51 127 0 1 0 0 1 1 1 1 0 0 1 1 1 1 1 1 0 1 0 0 1 1 0 1 0 0 0 0 0 0 0 1 0 0 1 0 0 0 1 1 1 0 0 1 0 1 1 1 0 0 0 0 1 1 0 1 0 0 0 1 0 1 0 0 32 18 162 51 5 71 231 127 0 1 0 1 1 1 1 1 0 0 1 1 0 1 1 1 0 0 0 0 1 1 1 1 0 0 0 0 0 0 0 1 0 1 0 1 0 0 0 1 1 0 1 1 0 0 1 1 0 0 0 0 0 1 1 1 0 0 1 0 0 0 1 0 8/2=4 ”Radix Digits”なので4回のソートパスで完了する 19
20. 従来手法: Upsweep-Downsweep (1) Upsweepパス: パーティションごとのRadix Digitヒストグラムを求める d-bit Radix Digitの場合、(LDSに)2^d個のバケツを用意してカウント 71 231 5 18 51 162 32 127 ソート対象の値 LSD 1 1 1 0 1 0 0 1 1 1 0 1 1 1 0 1 1 1 1 0 0 0 0 1 0 0 0 0 0 0 0 1 18 96 26 101 73 27 68 47 0 0 0 1 1 1 0 1 1 0 1 0 0 1 0 1 0 0 0 1 0 0 1 1 0 0 1 0 1 1 0 1 19 124 193 82 67 83 49 7 1 0 1 0 1 1 1 1 1 0 0 1 1 1 0 1 0 1 0 0 0 0 0 1 0 1 0 0 0 0 0 0 Partition 3 Partition 2 Partition 4 00: 1 01: 1 10: 2 11: 4 00: 2 01: 2 10: 2 11: 2 00: 1 01: 2 10: 1 11: 4 00: 01: 10: 11: Local Histograms Global Histogram パーティション数 (例: 7個) 2^d 2^d ローカルヒストグラムを グローバルヒストグラムに パーティション数分のストライドを伴ってコピー 2 3 1 2 1 4 1 2 1 1 2 2 0 3 3 2 2 2 1 3 3 1 2 4 2 4 1 1 20
21. 従来手法: Upsweep-Downsweep (2) グローバルヒストグラムのスキャン: パーティションごと、”位の値”ごとのグローバルなオフセットを求める 0 2 5 6 8 9 13 14 16 17 18 20 22 22 25 28 30 32 34 35 38 41 42 44 48 50 54 55 56 00: 01: 10: 11: Global Histogram パーティション数 (例: 7個) 2^d 2 3 1 2 1 4 1 2 1 1 2 2 0 3 3 2 2 2 1 3 3 1 2 4 2 4 1 1 Reduce-Then-Scanのスキャン同様、 「パーティション数 < 1つのスレッドグループで高速に扱えるサイズ」 としておけば一回のディスパッチでスキャンは完了する 21
22. 従来手法: Upsweep-Downsweep (3) Downsweepパス: パーティションごと”位の値”ごとのオフセットとローカルな位置を使ってスキャター 0 2 5 6 8 9 13 14 16 17 18 20 22 22 25 28 30 32 34 35 38 41 42 44 48 50 54 55 00: 01: 10: 11: Global Offsets 71 231 5 18 51 162 32 127 ソート対象の値 LSD 1 1 1 0 1 0 0 1 1 1 0 1 1 1 0 1 1 1 1 0 0 0 0 1 0 0 0 0 0 0 0 1 18 96 26 101 73 27 68 47 0 0 0 1 1 1 0 1 1 0 1 0 0 1 0 1 0 0 0 1 0 0 1 1 0 0 1 0 1 1 0 1 19 124 193 82 67 83 49 7 1 0 1 0 1 1 1 1 1 0 0 1 1 1 0 1 0 1 0 0 0 0 0 1 0 1 0 0 0 0 0 0 Partition 3 Partition 2 Partition 4 同じ”位の値”中の順番 0 1 0 0 2 1 0 3 0 0 1 0 1 0 1 1 0 0 0 0 1 2 1 3 44 45 17 30 46 31 5 47 32 6 33 18 19 48 7 49 50 8 20 34 51 52 21 53 スキャター先 22
23. Onesweep Radix Sort (1) 従来のRadix Sort実装はRadix DigitごとにVRAMにおおよそ3nのアクセス 新手法Onesweep: Decoupled Look-backに基づくRadix Sort • Radix Digitごとに一回のディスパッチでソートを完了 • パスごとにVRAMに2nアクセス Decoupled Look-back同様パーティションデスクリプターを持つ ただし”位の値”の種類分、つまり2^d個の「値+ステータス」 23
24. 24 Onesweep Radix Sort (2) • 全パスのグローバルヒストグラム計算、オフセット計算を先に行う (1nアクセス) • パーティションごと、”位の値”ごと パスの度に値が並び替わるので事前のパーティションごとのオフセット計算は不可能 • “位の値”ごとのヒストグラムを全パス分計算、それぞれスキャンして”位の値”ごとのオフセットを得る パーティションごとのオフセットはDecoupled Look-backで求められる 0 2 5 6 8 9 13 14 16 17 18 20 22 22 25 28 30 32 34 35 38 41 42 44 48 50 54 55 56 00: 01: 10: 11: (Per-Pass) Global Histogram パーティション数 (例: 7個) 2^d 2 3 1 2 1 4 1 2 1 1 2 2 0 3 3 2 2 2 1 3 3 1 2 4 2 4 1 1 14 11 16 15 16 11 19 10 … … … … 00 01 10 11 Pass 1 Pass 2 … 2^d 0 14 25 41 56 0 16 27 46 56 … … … … 56 00 01 10 11 パス数分の独立したスキャン (Per-Pass) Global Histograms (Per-Pass) Global Offset Arrays
25. 25 Onesweep Radix Sort (3) Pass 1 Pass 2 … 0 14 25 41 56 0 16 27 46 56 … … … … 56 00 01 10 11 (Per-Pass) Global Offset Arrays 71 231 5 18 51 162 32 127 ソート対象の値 LSD 1 1 1 0 1 0 0 1 1 1 0 1 1 1 0 1 1 1 1 0 0 0 0 1 0 0 0 0 0 0 0 1 18 96 26 101 73 27 68 47 0 0 0 1 1 1 0 1 1 0 1 0 0 1 0 1 0 0 0 1 0 0 1 1 0 0 1 0 1 1 0 1 19 124 193 82 67 83 49 7 1 0 1 0 1 1 1 1 1 0 0 1 1 1 0 1 0 1 0 0 0 0 0 1 0 1 0 0 0 0 0 0 Partition 3 Partition 2 Partition 4 同じ”位の値”中の順番 0 1 0 0 2 1 0 3 0 0 1 0 1 0 1 1 0 0 0 0 1 2 1 3 “位の値”のExclusive Prefix 44 45 17 30 46 31 5 47 32 6 33 18 19 48 7 49 50 8 20 34 51 52 21 53 スキャター先 00 01 10 11 5 3 5 3 00 01 10 11 6 4 7 7 00 01 10 11 8 6 9 9 • ”位の値”ごとにDecoupled Look-backを実行して (SIMD活用) パーティションごと、”位の値”ごとのExclusive Prefixを得る • “位の値”に対応するグローバルオフセット、Exclusive Prefix、パーティション中の順番 からスキャター先が求まる
26. 26 参考文献 • Parallel Prefix Sum (Scan) with CUDA • Single-pass Parallel Prefix Scan with Decoupled Look-back • Onesweep: A Faster Least Significant Digit Radix Sort for GPUs
27. 27 補足資料
28. 28 従来手法: Reduce-Then-Scan • VRAMへのアクセスが入力サイズnに対しておおよそ3n (2n reads, 1n writes) • パス数が入力サイズに応じて変わる 入力をブロックサイズ以下になるまで再帰的にブロックごとのリダクション 続いて逆向きにブロックごとのスキャンを繰り返す Input Buffer Dispatch 1: Block-wise Reduction Buffer 1 Buffer 2 Dispatch 3: Block-wise Scan Buffer 1 Input Buffer Dispatch 4: Block Exclusive Prefix + Block-wise Scan Dispatch 5: Block Exclusive Prefix + Block-wise Scan Dispatch 2: Block-wise Reduction Buffer 2
29. 29 従来手法: Scan-Then-Propagate • VRAMへのアクセスが入力サイズnに対しておおよそ4n (2n read, 2n writes) • パス数が入力サイズに応じて変わる 前段のリダクションステップだったところでスキャンも済ませておく 後段では読み出した値とブロックごとのオフセットを足すだけ Input Buffer Block-wise Scan Store Block Sum Buffer 1 Block-wise Scan Store Block Sum Buffer 2 Block-wise Scan Buffer 1 Input Buffer Block Exclusive Prefix + Scanned Block Block Exclusive Prefix + Scanned Block Dispatch 1 Dispatch 2 Dispatch 3 Dispatch 4 Dispatch 5