JavaScriptで2,000万個の粒子をシミュレーションする

 (dgerrells.com)
4 ポイント 投稿者 GN⁺ 2024-07-09 | 1件のコメント | WhatsAppで共有
  • CPUのみを使うJavaScript粒子シミュレーションから出発し、スマートフォンで100万個の粒子を60fpsで動かす目標を達成、最終サンプルは2,000万個の粒子まで拡張された
  • 性能の要は、オブジェクト配列を避けてTypedArrayとSharedArrayBufferで粒子データを連続メモリに置き、Web WorkerでCPUコアごとに処理を分割した点にあった
  • 実際のボトルネックは粒子位置の計算よりも描画とメモリアクセスに強く現れ、ピクセルバッファへのランダム書き込みがCPUキャッシュミスを引き起こした
  • Workerごとのピクセルバッファ、postMessage同期、ダブルバッファリング、ピクセルごとの粒子数グリッド、GPUテクスチャレンダリングを経て、ボトルネックを段階的に削減した
  • 最終構成ではCPUがシミュレーションし、GPUが画面解像度サイズのグリッドテクスチャをレンダリングし、M1 Macで2,000万個を約20fpsで処理した

目標と基本アプローチ

  • 課題は、WebAssemblyなしで純粋なJavaScriptとCPUだけを使い、スマートフォン上で100万個の粒子を60fpsでシミュレーションすることだった
  • 単純なJavaScriptオブジェクト配列には、モバイルの単一コア性能とメモリ配置の面で限界があると見て、データを連続配列として扱う方式を選んだ
  • 大量計算では、CPUキャッシュにできるだけ多くのデータを載せられる密な連続メモリ配列が重要になる
  • JavaScriptの通常の配列やオブジェクトは、低水準言語の配列やstructのように密に配置されるとは保証しにくい
  • TypedArrayは固定長の連続バイト配列を作れるため、JavaScriptでも比較的低水準のメモリ管理を可能にする

最初の実装: SharedArrayBufferとWeb Worker

  • 粒子データはSharedArrayBuffer上にFloat32Arrayビューを被せて保存した
    • 粒子1つはx、y、dx、dyの4つの32ビット浮動小数点値で表現した
    • stride = 4、byte_stride = stride * 4で平面バッファをインデックスした
  • シミュレーションは各フレームごとにx += dx、y += dyを行う単純な位置更新から始めた
  • マルチスレッド化はWeb WorkerとSharedArrayBufferで構成した
    • メインスレッドはWorkerごとに粒子チャンクを分割して渡した
    • Workerは自分に割り当てられたチャンクだけを読み書きし、同じメモリ位置に複数スレッドが書き込む状況を避けた
  • ブラウザでSharedArrayBufferを使うには次のヘッダーが必要だった
    • Cross-Origin-Embedder-Policy: require-corp
    • Cross-Origin-Opener-Policy: same-origin
  • 初期同期には別のシグナル用SharedArrayBufferを使った
    • メインスレッドはWorkerに実行シグナルを書き込む
    • WorkerはsetInterval(..., 1)でシグナルを確認し、作業が終わると準備状態を記録する
  • レンダリングはImageDataを使い、CPUベースのラスタライザのように処理した
    • 粒子1つを画面上の1ピクセルとして描く
    • 画面外の粒子はスキップして、配列範囲外アクセスを避けた
    • putImageDataでCanvasに結果を描画した
  • 初期プロファイリングではWorkerの処理は数ms程度で、時間の大半はメインスレッドのレンダリングに使われていた
  • メモリは起動時に一度だけ確保し、シミュレーションループ内で新しいオブジェクトを作らなかったため、ガベージ生成はほとんどなかった

2番目の実装: 入力と重力風の相互作用

  • 入力データは既存のシミュレーション状態用SharedArrayBufferに追加した
    • デルタ時間、マウスx/y、タッチ状態、画面の幅/高さを保存した
    • タッチ状態はbooleanに近い値だが、4バイトを使った
  • 画面を押すと粒子がタッチ位置へ引き寄せられる効果を追加した
    • 実際の重力式である1 / r²の形を参考にしつつ、正確な物理シミュレーションよりも面白い相互作用を狙って調整した
    • 摩擦に相当する減衰を加え、粒子が無限に加速して画面外へ飛んでいくのを防いだ
  • レンダリング色は粒子位置に応じて変わるように修正した
    • x / width、y / heightの値を使ってRGBチャンネル値を調整した
    • 粒子ごとに色を保存すると粒子サイズが大きくなるため、データ量を減らす目的で位置ベースの色を使った
  • 200万個の粒子まで試せるようになり、URLのcountクエリパラメータで粒子数を変更できるようにした
  • 1,000万個の粒子ではフレームレートが落ち、プロファイリングの結果、WorkerよりもImageDataへのピクセル書き込みが遅い部分だと分かった
  • RGB全体を書く代わりに青とアルファチャンネルだけを書くよう変えると、環境によってはレンダリング時間が約80〜90msから16〜18msまで短縮された
    • この差はCPUによって異なり、M1 MacBook AirよりデスクトップRyzenでより大きく現れた

3番目の実装: Workerでピクセルバッファを書く

  • より多くのCPUコアを活用するため、Workerがシミュレーションだけでなくピクセル描画も担当するように変えた
  • 複数のWorkerが同じピクセルバッファに同時書き込みしないよう、Workerごとのピクセルバッファを用意した
    • CPU_CORES * window.innerWidth * window.innerHeight * 3サイズのSharedArrayBufferを使用した
    • 各Workerは画面全体のRGBバッファを1つ持つ
  • メインスレッドはWorkerごとのピクセルバッファをすべて合算したあと、ImageDataに書き込んだ
  • このバージョンでは画面のちらつきが発生した
    • Workerが自分のピクセルバッファをfill(0)で消去している間に、メインスレッドが同じバッファを読んでいた
    • 以前のレンダリングでも古いデータを読む可能性はあったが、フレーム間の差が目立たず、今回はWorkerの消去処理が画面の一部のちらつきとして現れた
  • 解決候補はAtomicsベースの同期、postMessageでWorker完了を待つ方式、ダブルバッファリングだった

4番目の実装: postMessageベースの同期

  • ちらつきをなくすため、シグナル配列をやめてpostMessageでWorker完了を通知するようにした
  • メインスレッドはactiveWorkersの値を管理した
    • Workerは終わるたびにメッセージを送り、activeWorkersを減らす
    • すべてのWorkerが終わるとレンダリングし、次のrequestAnimationFrameを要求する
  • この方式でちらつきは解消した
  • フレーム時間は大きく改善せず、むしろ悪化することもあった
    • メインスレッドがWorker完了を待っている間、何もできないためだ
    • 200万個の粒子では、Workerがメインスレッドのレンダリングを待たなければ約7ms削減できる場面があった
  • このボトルネックを減らすため、グラフィックスドライバで一般的なダブルバッファリングを適用した

5番目の実装: ダブルバッファリングとキャッシュのボトルネック

  • ダブルバッファリングは2つのピクセルバッファを作り、フレームごとにアクティブバッファを切り替える方式だ
    • メインスレッドが片方のバッファを描画している間、Workerはもう片方のバッファに次フレームを準備する
  • 適用後、Workerはメインスレッドのレンダリングを待たなくなり、以前は50msを超えていた一部のフレームが16msにより近づいた
  • 引き続き遅かったのは、ピクセルバッファのRGBチャンネルに値を加算する3行のコードだった
  • 原因は、データ自体は連続して保存されていても、アクセス順序が連続していない点にあった
    • 粒子のx、yから計算したpixelIndexは粒子配列の順序と整列していない
    • ループごとにピクセルバッファの任意の位置へ飛ぶため、ランダムアクセスに近いパターンになる
  • このアクセスパターンは、粒子データとピクセルデータがどちらもキャッシュに収まらないとキャッシュミスを引き起こす
  • iPhoneを基準にした計算例は次のとおり
    • 公式: screen width * screen height * bytes per pixel + bytes per particle * total particles / number of workers
    • 2532 * 1170 * 3 + 16 * 2,000,000 / 5
    • 結果は15MBを超え、L1キャッシュに入れるには大きすぎる
  • このバージョンではM1チップで、200万個の粒子シミュレーションは約4ms、ピクセルデータアクセスは約3.5ms、ピクセルバッファ書き込みは約7msかかった
  • スマートフォンでは100万個の粒子を60fpsで維持できたが、20本超のスレッドを持つデスクトップでは、メインスレッドが20個超のピクセルバッファを毎フレーム合算するのに30msを費やした

6番目の実装: 開始位置へ戻る力

  • 新しい効果は、粒子が開始位置から離れるほど強く引き戻されるようにすることだった
  • そのために粒子へ開始位置sy、sxではなく開始位置sx、syの2値を追加した
    • 粒子1つはx、y、dx、dy、sx、syの6つのfloatを持つ
    • particleStride = 6に増えた
  • Workerコードには力計算関数が追加され、ガベージ生成を避けるためにキャッシュオブジェクトを使った
  • 結果としてゼリーや流体のような視覚効果が生まれ、粒子はしわくちゃの紙や布のような渦を見せた
  • この過程で、ピクセルバッファをWorkerごとに保持せず、ピクセルごとの粒子数だけを数えるグリッドに変えるアイデアが生まれた

7番目の実装: ピクセルごとの粒子数グリッド

  • RGB値を保存する代わりに、各ピクセル位置にある粒子数だけを保存すれば、メモリを1/3に減らせる
  • 色は粒子ごとのデータではなくx、y位置から計算するため、ピクセルごとのカウントだけでレンダリングできる
  • すべてのWorkerが1つの粒子数グリッドを共有するように変更した
    • Workerごとのバッファを合算するメインスレッドの処理が不要になった
    • Worker数が増えるとレンダリング合算コストが増える問題を解消した
  • 新しいキャッシュ計算式は次のとおり
    • screen width * screen height + bytes per particle * total particles / number of workers
    • iPhone基準で2532 * 1170 + 16 * 2,000,000 / 5
    • 約9.3MBとなり、以前より約30%減ったが、それでもL1キャッシュには大きい
  • レンダリングでは全ピクセルを走査し、そのピクセルの粒子数を読み、RGB値を計算したうえで、グリッドをfill(0)で初期化する
  • Workerはシミュレーション後、画面内にある粒子のpCountIndexを計算し、activeGrid[pCountIndex]++を実行する
  • この変更はWorker数が多いときに特に有効で、デスクトップでは24スレッドすべてを使えるようになった
  • それでもM1 MacBook AirがRyzen 9 3900Xデスクトップより速い場合があった
    • Ryzen 9 3900XのL3キャッシュは64MBだが、L1は64KB
    • 第1世代M1のL1キャッシュは320KBで、そのうち128KBがデータ、192KBが命令用
    • より大きなL1データキャッシュがメモリ待ち時間の削減に有利に働いたと見られる

CPU専用最適化の限界

  • 複数の最適化の結果、最初のマルチスレッド版に比べて約2倍の高速化を得た
  • JavaScriptだけでスマートフォン上のCPUベース100万個粒子を60fpsで処理するという目標は達成した
  • 全体の結果は特別に印象的というほどではなく、コンパイル言語なら10倍速いと期待している
  • tight loopでSIMD命令を使えれば、さらに高速化の余地がある
  • 次の段階では、シミュレーションはCPUに残したまま、レンダリングだけを従来のGPUパイプラインへ移して比較した

2,000万個の粒子とGPUレンダリング

  • 最初のGPUレンダリング比較は、threejsのGPU instancingで粒子ごとにplane/quadを描く方式だった
  • この方式は予想より性能が低かった
    • threejsのinstancingでは粒子位置ごとに行列を更新する必要があり、この作業はCPUの単一スレッドで行われる
    • 粒子の変換データを毎フレームGPUへアップロードする必要があり、この転送が遅い
  • M1チップで400万個の粒子をシミュレーションした場合、フレーム時間の80%がGPUへのデータ送信に使われた
  • デスクトップではデータ転送の比率は30%だったが、それでも最大のボトルネックだった
  • 代替案は、前述の粒子数グリッドをテクスチャとして使い、全画面quadをレンダリングする方式だった
    • GPUへ送るデータ量を粒子数ではなく画面解像度に固定できる
    • 大規模シミュレーションでは転送量が減る一方、小規模シミュレーションではより多くのデータを送る可能性があるというトレードオフになる
  • threejsとカスタムピクセルシェーダで素早く実装した結果、メインスレッドのレンダリングはシミュレーション規模に関係なく数ms程度になった
  • この構成ではWorkerは可能な限り速く実行され、シミュレーションは大規模でも小規模でもCPUコア数に直接縛られる
  • 各Workerは時間の約50%を数値計算に、約50%を粒子グリッドのカウント更新に使っていた
  • 最終版の性能は次のとおり
    • バッテリー駆動中のM1 Macで2,000万個の粒子を約20fpsで動かした
    • デスクトップでは同程度のfpsで約3,000万個の粒子を処理した
    • 32コアCPUを持つ友人のテストでは約4,000万個の粒子に達した
  • 関連デモ:

まとめ

  • CPUとGPUは数値計算を非常に高速に処理できる
  • 遅いのはデータの移動であり、とくにランダムメモリアクセスが問題になる
  • 高速なコードを書くには、ハードウェアがデータをどうキャッシュし、どう移動させるかを理解する必要がある
  • SharedArrayBufferとWeb Workerは、JavaScriptで複数コアを活用する実験に有用だった
  • WebGPUがさらに広く普及すれば、compute shaderも試してみる予定だ

関連記事

1件のコメント

 
GN⁺ 2024-07-09
Hacker News のコメント

  • ブログ内にシミュレーションを直接埋め込んでくれるとよかった
    実際のシミュレーションへのリンクを見つけるまで、かなり上下にスクロールしなければならなかった
    古いスマートフォンでもそこそこ動く値を選ぶか、フレームレートに応じて調整するようにすることもできるし、あるいは記事の上部にリンクをいくつか置くだけでもよい
    https://ciechanow.ski/ は、記事内にシミュレーションを入れるとどれほど素晴らしくなるかを示す世界最高レベルの例
    昔はどのウェブサイトも、こうした素晴らしいインタラクティブ要素を埋め込んでいたものだった
    それに、WebGPU なしでも GPU 上で粒子シミュレーションを動かせると思う
    例: https://news.ycombinator.com/item?id=19963640

    • これまで見た解説系ブログの中でも最高レベルの例の一つ
      情報共有の標準がこうなってほしい。共有する価値があるなら、理解しやすくする価値もある
    • あの流体シミュレーションは本当にすごい
    • 埋め込みが良いという点には同意
      各バージョンを埋め込もうかとも考えたが、ワーカーが同時に多すぎる数動くのではないかと心配だった
      記事の最後に最終版を埋め込むよう更新する予定
      あのブログは例ごとの完成度がものすごく高く、本当に良い
      追記: 埋め込み版を追加しようとしたが、必要なヘッダーが他の埋め込みとうまく合わず、古いバージョンはいまだに codesandbox に縛られている
    • 「最終アプリを触ってみたい場合は末尾へ飛んでください」と案内すればいい

  • モバイルデモ [0] を直接触ってみたら、かなり驚いた
    [0] https://dgerrells.com/sabby

    • 複数の指でも動作する
      純粋な JavaScript でここまでできるとは驚きで、指が増えるほど遅延が大きくなるのを見ると、スレッド間でやり取りされるデータが増えているようだ
    • 妙に有機的に見えて、原始スープの中の脂質構造のよう
    • 私のスマートフォンでは Firefox のほうが Chrome よりよく動いて満足

  • 純粋な好奇心だが、各粒子を x, y, dx, dy という4つの32ビット浮動小数点数で表す代わりに、このデータを1つの JS 数値にエンコードできるのか気になる
    JS の安全な整数範囲は 2^53 - 1 = 9,007,199,254,740,991 で、記事で使っている Float32Array の範囲は -3.4e38 から 3.4e38 まで
    例えば画面位置が 1000x1000 のキャンバスなら 0〜1,000,000 個の数値で表現でき、そこにサブピクセル 10 段階を加えても 100,000,000 個なので JS の範囲に十分収まる
    速度 dx, dy も (Math.random()*2-1)*10 で -10 から +10 の間のランダムな小数として計算しているが、小数1桁に制限して十分なら -100〜+100 の整数範囲に変えられ、取り得る値は 10,000 個で足りる
    2つを合わせると 10,000 * 100,000,000 = 1,000,000,000,000、つまり1兆個の数値で粒子を表現でき、JS の MAX_SAFE_INTEGER 内に収まる
    なので単一粒子のデータを1つの MAX_SAFE_INTEGER や Float32Array の要素1つに収められそうで、そうすればストライドが不要になり、データの一貫性もより確実にできるかもしれない
    ただし、エンコード/デコードのコストがメモリ節約を上回り、性能上まったく価値がない可能性もある

  • 本当に素晴らしいけれど、この部分が気になる
    Javascript does support an Atomics API but it uses promises which are gross. Eww sick.
    waitAsync[1]を除けば、Atomics APIは Promise を使っていないように見える
    以前 Atomics を使ったことがあるが、非同期／Promise のコードを扱う必要はなかった
    内部的に Promise を使っているのか、それとも見落としている部分があるのか気になる
    [1] https://developer.mozilla.org/en-US/docs/Web/JavaScript/Refe...

  • ほぼ 2,000万個の粒子を扱うシミュレーションもある
    実行してプログラミングできる: https://tinlizzie.org/~ohshima/shadama2/
    https://tinlizzie.org/~ohshima/shadama2/live2017/
    https://www.youtube.com/watch?v=tqBO_5GMfWw

  • 動画は素晴らしいが、try it out here の codesandbox リンクは MacOS Chrome デスクトップでは動かない
    Uncaught ReferenceError: SharedArrayBuffer is not defined や ERR_BLOCKED_BY_RESPONSE.NotSameOriginAfterDefaultedToSameOriginByCoep のような CORS エラーが出る

    • codesandbox のインラインエディタがヘッダー設定を妨げるため、プレビューを 専用タブで開く必要がある
      プライバシー重視のブラウザを使っているとブロックされる場合もある
      今後は埋め込みサンプルも含めてみる
      最終版はこちらで試せる: https://dgerrells.com/sabby

  • 素晴らしい
    コードを理解していじるのに1時間使い、ここに リアルタイム実装を作っておいた: https://particules.kapochamo.com/index.html

    • 最初に開いたときは驚いた
      本当にきれいで、JS でこれほど多くの粒子をシミュレートしているとは信じがたい
    • いいね
      そこまで深くコードを掘り下げたのが印象的だ

  • chrome://tracing がさらに多くの洞察を与えてくれるか確認してみる価値がある
    ここで知った: https://youtu.be/easvMCCBFkQ?t=114

  • これを、遅い React アプリを作っている UI チームに早く送りたい
    JS は本当に速く、特にうまく書けば ものすごく速い

    • 問題は、慣用的な JS と高速な JS が実際にはほぼ正反対だという点
      高速な JS は出来の悪い C 方言のように見え、ガベージコレクション対象のオブジェクト割り当てをしてはいけないのに、この言語には割り当てなしで使える良い代替手段が少ない
      複雑な JS アプリで、すべてのメモリ割り当てを巨大な ArrayBuffer 内のポインタとして扱う人はほとんどいないし、それなら WebAssembly に移るほうが簡単だ
    • JavaScript は、各種エンジン最適化に投入された 人間の時間が、おそらく最も多い言語だろう
      残念なのは JS だけを信じて進めるわけにはいかず、DOM 操作を大量に混ぜる必要がある点で、私たちが作る UI でたいてい遅い部分はその DOM 操作だ