WebGPU で MLS-MPM を実装し，ブラウザ上で 10 万粒子規模のリアルタイムシミュレーションを実現する
以前に書いた WebGPU で実装したリアルタイム 3D 流体シミュレーションの紹介 の続編として，MLS-MPM という手法による流体シミュレーションを WebGPU で実装して公開したところ，ありがたいことに予想を遥かに超える反響をいただきました．
本記事では，このシミュレーションに関する以下の話題について書きます．
- なぜ MLS-MPM をシミュレーションに用いたのか
- MLS-MPM のパフォーマンス
- WebGPU による MLS-MPM の実装
また，流体のレンダリングについては，以前の記事である WebGPU で実装したリアルタイム 3D 流体シミュレーションの紹介 に詳しく書いたので，そちらをご覧ください．
以下がデモとリポジトリへのリンクです．
- Demo :
- Repository :
なぜ MLS-MPM をシミュレーションに用いたのか
まずは，今回のシミュレーションに MLS-MPM を用いた理由や背景を説明します．なお，MLS-MPM の理論の詳しい説明については，そもそも筆者自身が完全に理解していないこともあり，他の文献に譲ろうと思います．参考文献 の [1]~[4] は，MPM や MLS-MPM の理論を理解するのに非常に参考になる資料です．
以前書いた WebGPU で実装したリアルタイム 3D 流体シミュレーションの紹介では，シミュレーション手法として SPH 法（Smoothed Particle Hydrodynamics）を実装し，パフォーマンスのボトルネックである近傍探索を高速化しました．しかしながら，依然として近傍探索のコストは大きく，さらに大規模なシミュレーションを行うためには近傍探索を必要としない手法を実装する必要がありました．
そこで白羽の矢が立ったのが，2018 年に Hu らによって提案された MLS-MPM（Moving Least Squares Material Point Method） という新しい手法です．この手法は，流体に限らず，様々な特性を持つ物質を統一的に扱える MPM (Material Point Method) というシミュレーション手法を高速化したものです．流体シミュレーションには，粒子どうしの相互作用を直接計算してシミュレーションを行う粒子法と，空間を格子(グリッド)状に分割して，そのグリッドの上で計算を行う格子法という 2 つの枠組みがあるのですが，MPM や MLS-MPM は粒子とグリッドを両方用いるハイブリッドな手法です．具体的には，粒子からグリッドへと情報を受け渡す Particle to Grid（P2G） ステップと，グリッドから粒子へと情報を書き戻す Grid to Particle (G2P) ステップを繰り返して，粒子の座標を更新していきます．
MPM や MLS-MPM によるシミュレーションの大きなメリットは，近傍探索が必要ないことです．必要なのは，P2G ステップで各粒子の周辺 3x3x3=27 マスへ情報を書きだすことと，G2P ステップで同じく周辺 27 マスから粒子へ情報を書き戻すこと（3 次元の場合）のみであり，粒子どうしの相互作用を SPH 法のように直接計算する必要はありません．それゆえ，近傍探索がボトルネックだった以前までのシミュレーションをさらに高速化できると考え，MLS-MPM を実装することにしました．
シミュレーションのパフォーマンス
実装面について説明する前に，MLS-MPM を実装することで，以前の記事で実装した SPH 法と比べてどれくらいの高速化が達成されたのかを実際に見てみます．以下は，ノートパソコン向けの AMD Ryzen 7 5825U という CPU の統合グラフィックス上で，MLS-MPM, SPH のそれぞれで 10 万粒子のシミュレーションを実行したときの様子です．
MLS-MPM による 10 万粒子のシミュレーション
SPH による 10 万粒子のシミュレーション
MLS-MPM のほうがが圧倒的にスムーズです．統合グラフィックス & ブラウザ上という環境で，この規模の粒子をリアルタイムシミュレーションできるのは素晴らしいですね．
もっと強い GPU を用いると，さらに大規模なシミュレーションも可能です．以下は，RTX 3060 Laptop という GPU を搭載したノートパソコン上で，50 万粒子をシミュレーションしたときの様子です．
MLS-MPM による 50 万粒子のシミュレーション（RTX 3060 Laptop 上）
ブラウザ上でここまでの規模のシミュレーションが可能だとは，まったく予想していませんでした．SPH 法を用いて，ブラウザ上でこの規模のシミュレーションを実現するのは不可能に近いと予想しています．
WebGPU による MLS-MPM の実装
ここでは，MLS-MPM の実装面について説明します．今回の実装は，nialltl さんによる mpm guide というチュートリアルを大いに参考にしています．
今回実装した MLS-MPM の 1 シミュレーションステップは，以下の計算からなります．
- グリッドの初期化
- P2G ステップ（1 回目）
- 粒子の質量をグリッドに分配します．
- P2G ステップ（2 回目）
- 1 回目の P2G ステップでグリッドに分配した質量から，粒子の体積を計算します．
- 粒子の運動量をグリッドに分配します．
- グリッドの速度の更新・境界条件の適用
- G2P ステップ
- グリッドから粒子へ情報を書き戻して粒子の速度を計算し，粒子の座標を更新します．
ここでは，特に実装の際のポイントとなった，P2G ステップの WebGPU による実装について説明します．
atomicAddによる P2G
mpm guide の MLS-MPM は，比較的素直に WebGPU で実装できたのですが，P2G ステップの実装には少し工夫が必要でした．
P2G ステップを GPU 上で実行する場合，各粒子のデータをその周辺のグリッドに対して並列に足し合わせる操作が必要になります．この操作は，WebGPU におけるアトミック加算のための命令であるatomicAddで実現できます．atomicAddにより，データ競合を起こすことなく，粒子のデータをグリッドへ並列に分配することが可能です．
しかし，WebGPU におけるatomicAddは 32bit 整数向けのものしか存在せず，浮動小数点数として保持されているデータに対してそのままatomicAddを用いることができません．粒子やグリッドの物理量は浮動小数点数として保持されているので，これは困ります．代替案として，いわゆる compare and swap に相当するatomicCompareExchangeWeakという命令を使うという選択肢もあったのですが，セルへデータをを分配するたびにいちいちループで待たなければならないため，性能への悪影響を懸念しました．（ちなみに，WebGPU-Fluid-Simulationというリポジトリの MLS-MPM の実装では，P2G をこの方法で実装しているようです．デモが手元の環境で動かないので性能は分からず．）
この問題を解決するために採用したのが，固定小数点を用いてデータを保持するという方法です．つまり，物理量 atomicAddを使うことができるようになります．たとえば，粒子の運動量をグリッドのセルへ分配するのは，以下のようなコードで行われます．
atomicAdd(&cells[cell_index].vx, encodeFixedPoint(momentum.x));
atomicAdd(&cells[cell_index].vy, encodeFixedPoint(momentum.y));
atomicAdd(&cells[cell_index].vz, encodeFixedPoint(momentum.z));
momentum.x（y,z も同様）は浮動小数点として保持されており，encodeFixedPointではその値に cells[cell_index].vx（vy, vz も同様）は整数として保持されているため，atomicAddを用いることができます．cells[cell_index].vxから物理量を復元する際には，
このテクニックは，Chris Lewin による A Position Based Material Point Method (SIGGRAPH 2024) という論文の reference implementation のリポジトリである pb-mpm で発見したものです．よく知られているテクニックだったりするんでしょうか？
まとめ & 所感
以下がまとめです．
- MLS-MPM（Moving Least Squares Material Point Method）を用いた流体シミュレーションを，WebGPU を用いて実装した．
- SPH 法でパフォーマンスのボトルネックとなっていた近傍探索を行う必要がなくなり，CPU の統合グラフィックス上で 10 万粒子規模のシミュレーションが可能になった．
- 固定小数点数を用いることにより，P2G ステージを
atomicAddを用いて実装することができた．
筆者がはじめてブラウザ上での流体シミュレーションに触れたのは，david.li の fluid という作品でした．このシミュレーションに触れてから，10 万粒子規模のブラウザ上流体シミュレーションを自分で実装するのは長い間の目標だったので，MLS-MPM を用いてそれを実現できたのは嬉しいです．ブラウザで Unreal Engine レベルの流体シミュレーションを！... とまではさすがにいきませんが，ブラウザ上でもそれなりに大規模な流体シミュレーションができることを知れたのは大きな収穫でした．
次の方向性としては，以下のようなことを漠然と考えています．
- 水しぶきのシミュレーション機能を追加して，以下の動画のような波を再現してみる． 水しぶきや泡といった効果を実装する方法について包括的に論じた Unified Spray, Foam and Bubbles for Particle-Based Fluids という有名な論文があるのですが，この論文の手法は近傍探索を必要とするため，ブラウザ上でのリアルタイムシミュレーションを行うにはもっと高速な方法を見つける必要がありそうです．もし実現できれば，かなり見ごたえのあるシミュレーションになるだろうと予想しています．
- SIMD を用いて，CPU 上での流体シミュレーションの高速化を行う．つい先日，John Lin という方の "This Real-Time 3D Fluid Sim Runs on One CPU Core." という以下の動画を見つけたのですが，この動画ではなんと 1 CPU 上で非常に大規模な流体シミュレーションを実現しています（おそらく数十万粒子規模？）．概要欄によれば，MLS-MPM を AVX2 を用いて高速化した結果，この規模のシミュレーションが可能になったそうです（ちなみに，MLS-MPM の論文の著者である Yuanming Hu がコメント欄に出没しています）．筆者は，highload.fun という高速化コンペティションサイトで AVX2 を用いた異常高速化を行うのが趣味なので，そこで培った技術を以下のような大規模流体シミュレーションの実装に活かしたいと考えています．
参考文献
- [1] : A Moving Least Squares Material Point Method with Displacement Discontinuity and Two-Way Rigid Body Coupling (SIGGRAPH 2018)
- Hu らによる MLS-MPM の元論文です．
- 論文の reference implementation が置かれたリポジトリである taichi_mpm には，たった 88 行からなる MLS-MPM の実装があります．
- [2] : The Material Point Method for Simulating Continuum Materials
- 2016 年の SIGGRAPH の Course にて公開された，MPM に関する講義資料です．MPM の理論に関する資料の中では，最も体系的に書かれたものだと思います．
- [3] : mpm guide
- nialltl さんという方によって書かれた，MLS-MPM をインクリメンタルに実装していく，実践的な記事です．今回のシミュレーションの実装は，この記事に基づいています．
- [4] : 解説: Jelly
- ハイクオリティな物理シミュレーションを数多く実装されている，saharan さんの記事です．MLS-MPM のみならず，SPH や PIC/FLIP といった流体シミュレーションの諸手法とそれらのつながりについても書かれており，流体シミュレーション全体を俯瞰するのに非常に参考になる記事です．
- [5] : pb-mpm
- Chris Lewin による A Position Based Material Point Method (SIGGRAPH 2024) の reference implementation が置いてあるリポジトリです．P2G のために固定小数点を用いるアイデアはここで学びました．
- [6] : Water3D
- saharan さんによる，MLS-MPM を用いた流体シミュレーションです．このデモを触ってみることを通して，MLS-MPM による大規模な流体シミュレーションの可能性を知ることができました．[4] には WebGL による MLS-MPM の実装についても書いてあり，併せて読むのがおすすめです．
Discussion