GPU Acceleration and Portability of the TRIMEG Code for Gyrokinetic Plasma Simulations using OpenMP

本論文は、ポータブルなOpenMP APIを用いたGPUアーキテクチャへのTRIMEGジャイロキネティック・プラズマシミュレーションコードの移植成功を提示するものであり、ハイブリッドMPI-OpenMP並列化およびイオン温度勾配モードシミュレーションを通じて、大幅な高速化と検証された物理的正確性を実証している。

要約

概要：宇宙規模の嵐を調理する

星の内部で起きている天気を予測しようとしている場面を想像してみてください。現実の世界では、太陽や核融合炉の中に直接温度計を突っ込むことはできません。あまりにも熱く、混沌としすぎているからです。その代わりに、科学者たちはスーパーコンピュータを使って、プラズマ（超高温の電気を帯びたガス）の「仮想シミュレーション」を実行します。

TRIMEGコードは、このプラズマをシミュレートするための、非常に高度で洗練された「レシピ」です。これは、何十億もの微細な粒子（嵐の中の砂粒一つひとつのようなもの）の動きを追跡し、それらがどのように渦巻き、衝突し、乱流を生み出すのかを調べます。問題は、このレシピが非常に「重い」ことです。標準的なコンピュータ（CPU）でこれを動かすのは、スプーン一本で山を動かそうとするようなものです。時間がかかりすぎます。

目標： 著者であるジョルジオ・ダネリ（Giorgio Daneri）氏は、GPU（グラフィックス・プロセッシング・ユニット）を使用して、この処理を高速化したいと考えました。CPUを「非常に賢いが、一度に一つの野菜しか切れない熟練のシェフ」だとしましょう。一方、GPUは「1万人の副料理長（ソムリエ）がいるキッチン」のようなものです。全員が同時に野菜を切ることができます。この論文は、その一人の熟練シェフのレシピを、1万人の副料理長たちのために完璧に機能させ、かつ、2つの異なるブランドのキッチン（NVIDIAとAMD）の両方で動作するようにする方法について書かれています。

課題：「ユニバーサル・トランスレーター（万能翻訳機）」の問題

著者は、翻訳を行うためのツールとしてOpenMPを選びました。OpenMPを、「レシピのこの部分をGPUに渡して」とコンピュータに伝えるユニバーサル・トランスレーターだと考えてください。

しかし、著者は2つの大きな壁にぶつかりました。

1. 「コンパイラ」の不具合： コードを翻訳するソフトウェア（コンパイラ）が完璧ではありませんでした。それは、ユニバーサル・トランスレーターを使っているのに、時々「塩」や「熱」という言葉を忘れてしまうような状態でした。著者は、この翻訳機の癖に合わせてコードの一部を書き直さなければなりませんでした。例えば、コードには高度な「ポリモーフィズム（多態性）」（オブジェクトが形や正体を変化させられるという高度な仕組み）が使われていました。GPU用の翻訳機（コンパイラ）はこの「形を変える性質」を理解できなかったため、著者はこれらを機能させるために、形を固定された箱へと「平坦化」する必要がありました。
2. 「交通渋滞」： メインコンピュータ（CPU）とGPU（副料理長たち）の間でデータを移動させるのは低速です。もし、材料を渡すたびに作業を中断して行ったり来たりしていたら、副料理長たちは何もできずに立ち往生してしまいます。著者は、レシピの材料を最初に一度だけGPUに移動させ、何度も行ったり来たりさせないように、コードの構造を作り直しました。

解決策：キッチンの再構築

コードをNVIDIAとAMDの両方のGPUで動作させるために、著者はTRIMEGコードに対していくつかの「手術」を行いました。

• マップの平坦化： コードは粒子の位置を見つけるために複雑な「マップ」を使用していました。このマップは、まるで散らかった書類整理棚のようなものでした。著者は、GPUが迷うことなく即座に読み取れるよう、これを一つの直線的なリストへと平坦化しました。
• 「レース」の修正： 何千人もの副料理長が同時に同じホワイトボードに書き込もうとすると、お互いの文字を塗りつぶしてしまうことがあります（「レース条件」）。著者は、コードの中でこのようなことが起きている箇所を見つけ出し、全員が自分のレーンに書き込めるように修正しました。
• 「ワンサイズ・フィット・オール（汎用性）」の妥協： NVIDIAとAMDという2つのGPUブランドは、少し異なる言語を話すため、著者は両方で動作する単一のコード版を作成しなければなりませんでした。たとえそれが、片方のハードウェアにとって絶対的な最速ではないとしても、両方で機能する特定のメモリ割り当てを使用するなどの「回避策」を用いることになりました。

結果：うまくいったのか？

著者は、新しいGPUバージョンを、古いCPUバージョンと比較するために、2つの有名な「テストケース」（新しい車の標準的な運転テストのようなもの）を使用しました。

1. サイクロン・ケース： プラズマ乱流の簡略化されたシミュレーション。
2. TCV-X21 ケース： プラズマの端の部分を含む、より複雑で現実的なシミュレーション。

判定：

• 速度： GPUバージョンは大幅に高速化されました。あるテストでは、単一のマシン上でCPUバージョンの約30倍高速でした。
• 精度： GPUによる結果は、CPUの結果とほぼ完全に一致しました。「天候パターン」（エネルギーの成長や乱流構造）は同じでした。
• 移植性： コードは、完全に書き直すことなく、NVIDIAとAMDの両方のハードウェアで正常に動作しました。

注意点（限界）

著者は、以下の限界についても正直に述べています。

• 「翻訳機」はまだ完璧ではない： これらのGPU用のコンパイラ（コードを機械語に変換するソフトウェア）は、まだ発展途上です。時として、CPUとはわずかに異なる計算結果を生成することがあり、これが時間の経過とともに小さな誤差を生む可能性があります。
• ハードウェアのミスマッチ： CPUコアが多く、GPUが1つしかないコンピュータの場合、一度に多くのタスクをGPUに与えすぎると、GPUが処理しきれなくなることがあります。最高の成果を得るためには、「シェフ（MPIプロセス）」の数と「副料理長（GPUスレッド）」の数のバランスを取る必要があることを著者は発見しました。
• 「魔法の弾丸」はない： 粒子を動かす部分のコードは劇的なスピードアップを実現しましたが、その他の部分（磁場方程式を解く部分など）は、まだその部分をGPUに移行するためのツールが整っていないため、依然としてCPUで動作しています。

まとめ

要するに、この論文はエンジニアリングの創意工夫の物語です。著者は、重くて遅く、複雑なシミュレーションコードを取り上げ、現代の強力なグラフィックスカード上で動作するようにすることに成功しました。ソフトウェアのバグやコンパイラの制限という地雷原を切り抜け、2種類の異なるハードウェアで動作するバージョンを作り上げ、精度を損なうことなく、より速く核融合プラズマをシミュレートできることを証明しました。これは、核融合エネルギー研究をより効率的にするための重要な一歩ですが、完全に自動化された完璧な翻訳への道のりは、まだ終わっていません。

技術概要

技術要約：OpenMPを用いたジャイロキネティック・プラズマシミュレーションのためのTRIMEGコードのGPU加速とポータビリティ

問題提起
プラズマ物理学のシミュレーション、特にトカマク型核融合装置における不安定性や乱流の研究に使用されるジャイロキネティック・モデルは、計算集約的である。TRIMEGコードは、非構造化三角形メッシュ上でC1有限要素法を利用する高精度な粒子・格子法（PIC）ソルバーであるが、現実的なシミュレーションに必要となる膨大な数の粒子（しばしば$10^7$から$10^8$）によって、実行時間の面で大きな課題に直面している。本コードはすでにマルチノード並列化のためにMPIを採用しているが、粒子のプッシング（particle pushing）および格子・粒子間操作（G2P）が主要なボトルネックとなっており、全実行時間の最大80%を占めている。課題は、これらの特定の「ホットスポット」をグラフィックス処理装置（GPU）を用いて加速させつつ、異なるハードウェア・アーキテクチャ（具体的にはNVIDIAおよびAMD）間でのポータビリティを維持し、かつ、ポリモーフィズムや派生型を含むコードの複雑なオブジェクト指向構造を維持することにある。

手法
本研究では、OpenMP オフローディングAPI（バージョン4.0以降）を用いて、TRIMEGコードをGPUアーキテクチャへ移植することに焦点を当てている。手法は以下の通りである：

1. ターゲットの選定： 高い算術強度を持ち、粒子間の依存関係がないことから、粒子プッシャーのカーネルおよび関連するG2P操作（プルバック、密度計算、および分布関数補間）を、オフローディングの主要なターゲットとして特定した。
2. ポータビリティのためのコード再構築： amdflang（AMD）およびnvfortran（NVIDIA）の両コンパイラの制限を克服するために、大幅な再構築が必要となった。主な課題は以下の通りである：
   • ポリモーフィズム： 両コンパイラとも、GPUターゲット領域内でのclass()派生型および型結合プロシージャの扱いに苦慮した。解決策として、可能な限り非ポリモーフィックなtype()宣言を使用するようにコードをリファクタリングし、粒子クラスと電場クラス間の循環依存関係に対して、ベースクラス/拡張クラスの階層構造と、関数本体を複製するためのFortranのINCLUDEディレクティブを用いた回避策を実装した。
   • 動的配列： コードではC++のvectorを模倣したカスタムライブラリを使用して動的配列を利用していた。GPUカーネルは動的なメモリ確保や複雑なポインタの誘導を容易に扱うことができないため、バウンディングボックスとメッシュ三角形の間のマッピング構造を、配列の構造体から1次元配列へと「平坦化」し、効率的なメモリ転送を容易にした。
   • メモリ管理： ランタイムのレイテンシを最小限に抑えるため、初期化フェーズにおける予備的なメモリ確保を実装した。AMDプラットフォームでは、可能な限りユニファイド共有メモリ（USM）を活用し、USMをサポートしていないNVIDIAプラットフォームでは、明示的なenter data、update、およびexit dataディレクティブを使用した。
   • 数値的一貫性： GPUの結果がCPUのリファレンスと一致することを確認するため、浮動小数点演算を変更するコンパイラ最適化（Fused-Multiply-Add命令など）を無効化した（AMDでは-ffp-contract=off、NVIDIAでは-Mnofma）。外部Bスプラインライブラリにおけるレースコンディションは、共有オブジェクトメンバからローカルに宣言された自動配列に切り替えることで解決した。
3. 性能評価： 実装のテストは、Viperクラスター（AMD MI300A）、Raven（NVIDIA A100）、およびPitagora（NVIDIA H100）クラスターで実施した。性能評価は以下を通じて行った：
   • カーネル・プロファイリング： rocprof-computeおよびnsysを使用して、リソース占有率、メモリ帯域幅、および命令ミックスを分析した。
   • スケーラビリティ研究： ハイブリッドMPI-OpenMPオフローディングの効率を評価するために、強力スケーリング（strong scaling）テストを実施し、特にGPUを複数のMPIプロセスでオーバーサブスクライブ（oversubscribing）することの影響を検討した。
   • グリッドサイズの探索： ハードウェア利用率を最大化するために、OpenMPチーム数およびチームあたりのスレッド数をチューニングした。

主な貢献

• 初のクロスベンダー・ポート： 本研究は、単一のコードベースを用いて、OpenMPオフローディングを介して複雑なオブジェクト指向Fortranジャイロキネティック・コードをNVIDIAおよびAMDの両方のGPUに移植するという先駆的な取り組みを提示している。
• コンパイラの回避策： 本論文は、ポリモーフィズム、動的配列、およびプロシージャポインタに関するコンパイラの制限を回避するための、具体的かつ非自明な回避策を記録している。これは、nvfortranおよびamdflangにおけるこれらの機能に関する包括的なドキュメントの欠如を浮き彫りにしている。
• ハイブリッド並列化の分析： ハイブリッドMPI-OpenMPオフローディングのトレードオフに関する詳細な分析を提供し、GPU加速が粒子プッシャーに対して有効である一方で、元のコードにOpenMPマルチスレッディングが存在しないことが、高コア数だがGPUリソースが限定的なノード上でのスケーラビリティを制限する可能性があることを示した。
• 数値的検証： エネルギー成長率と2Dモード構造をCPUの結果と比較する厳格な検証プロセスを含んでおり、コンパイラ固有の浮動小数点処理による軽微な数値的偏差はあるものの、GPU版が高度な忠実度で物理現象を再現することを確認した。

結果

• スピードアップ： $32 \times 10^6$個の電子を用いた現実的なワークロードにおいて、GPU実装は、TOKクラスター上のGCCコンパイルされたCPUバージョンと比較して、AMD Viperノードで約14.8倍、NVIDIA Pitagoraノードで約29.6倍のスピードアップを達成した。
• カーネル効率： 粒子プッシャーのカーネルが実行時間の大部分を占めた。AMD MI300Aでのプロファイリングでは、高い算術強度（80%以上のL1/L2キャッシュヒット率）を示したが、メモリアクセスのコアレセ（coalesced）率は18%に留まった。
• スケーラビリティの制限： 強力スケーリングテストにより、GPU加速された部分は良好にスケールする一方で、アプリケーション全体のスピードアップは、非加速部分（例：PETScを用いた電場ソルバー）やGPUのオーバーサブスクライブによるオーバーヘッドによって制約を受けることが明らかになった。NVIDIA Pitagoraクラスターでは、OpenMPによるマルチGPUサポートがテストされたコンパイラバージョン（nvfortran 24.9）において機能していないことが判明し、ノード上の利用可能なすべてのGPUを同時に活用することが制限された。
• 正確性： Cycloneケース（ITGモード）およびTCV-X21ケース（非線形ITG不安定性）のシミュレーションにより、GPU版がCPU版で見られるエネルギー成長率およびモード構造を正しく再現していることが確認された。差異は、アルゴリズムのエラーではなく、乱数生成器の初期化およびコンパイラ固有の浮動小数点演算の差異に起因するものである。

意義と主張
本論文は、OpenMPオフローディングが異なるHPCアーキテクチャ間でのポータビリティに向けた有望な道筋を提供する一方で、複雑なレガシーコードにとって「シームレスな」解決策ではないことを主張している。本研究は、動作する高性能なGPU版を実現するためには、高度なFortran機能に対するコンパイラサポートの制限を回避するために、広範なコンパイラ探索と大幅なコード再構築が必要であることを示している。

著者らは、このポータビリティの成功は、単なるプログラミング・パラダイムではなく、特定のコンパイラバージョンに強く依存していることを強調している。彼らは、TRIMEGのGPU実装が、ジャイロキネティック・シミュレーションのための機能的かつ正確なツールであり、最も計算コストの高い部分に対して大幅なスピードアップを提供できることを結論づけている。しかし、ハードウェアの全潜在能力（特にマルチGPUノード）は、現在の未成熟なコンパイラによるマルチデバイス・オフローディングのサポートおよび、基礎となるCPUコードにおけるOpenMPマルチスレッディングの欠如によって阻害されていることを控えめに述べている。本研究は、ヘテロジニアスなアーキテクチャへ複雑なFortranコードを移植しようとする人々にとって、実践的なガイドおよび「代替ドキュメント」としての役割を果たすものである。