OpenMP GPU Acceleration and Portability of TRIMEG-C1 for Electromagnetic Gyrokinetic Simulations in Tokamak Plasmas

技術要約：TRIMEG-C1におけるOpenMPによるGPU加速とポータビリティ

問題提起
TRIMEG-C1は、トカマクプラズマにおける電磁的不安定性をシミュレーションするために設計された、三角形メッシュベースのジャイロキネティック・パーティクル・イン・セル（PIC）ソルバーである。本コードは、$C^0$手法と比較して低いグリッド解像度で正確な物理現象を捉えるため、高次の$C^1$有限要素法を利用している。しかし、現代のスーパーコンピュータがヘテロジニアスなアーキテクチャに依存する中で、本コードは、微小な電子スケールや輸送タイムスケールを伴うマルチスケールシミュレーションへのスケーリングという課題に直面している。PICコードは高い並列性を持つため、GPUオフローディングは大幅な加速の可能性を提供するが、既存のアプローチはポータビリティに欠けることが多い。例えば、GEMやORB5で使用されているOpenACCはNVIDIAアーキテクチャでは強力にサポートされているが、HPCクラスター（例：Max Planck Computing and Data FacilityのViperクラスター）で普及が進んでいるAMD GPUに対する堅牢なサポートに欠けている。一方で、CUDAやHIPは、大幅なコードの書き換えや手動でのカーネル変換を必要とする。本研究が取り組む主要な問題は、既存のFortranコードの構造を損なうことなく、また大規模なコードの重複を必要とせずに、NVIDIAおよびAMDの両方のGPUアーキテクチャに対して、ポータブルで加速された実装を実現することである。

手法
著者らは、加速戦略としてOpenMPオフローディングフレームワークを選択した。この選択は、単一のFortranコードベース内でNVIDIAとAMDの両方のGPUをサポートし、構造的な変更やコードの分岐を最小限に抑えられるという理由に基づいている。移植プロセスは、最も計算負荷の高いコンポーネント、すなわち粒子プッシャー（particle pusher）カーネルと、グリッド・トゥ・パーティクル（G2P）操作（プルバック処理、密度計算、分布関数補間を含む）に焦点を当てた。

主な手法ステップは以下の通りである：

1. カーネルの適応: 独立してマーカー粒子を進化させる粒子プッシャーを、単一の大きなOpenMP target領域内にラップし、並列性を最大化した。これには、外部のBスプライン補間ライブラリを含むすべての呼び出し関数がデバイスから呼び出し可能であることを保証するためのリファクタリングが必要であった。
2. コンパイラの制限に対する回避策: amdflangおよびnvfortranコンパイラにおけるOpenMPオフローディングの制限を克服するために、多大な努力が払われた。具体的な課題は以下の通りである：
   • ポリモーフィズムと型結合プロシージャ: いずれのコンパイラも、デバイスコードにおけるこれらの機能を完全にはサポートしていない。著者らは、必要に応じてclassではなく非ポリモーフィックなtype宣言を使用するようにコードを修正し、Bスプラインライブラリにおける型結合プロシージャのバインディングを回避した。
   • メモリ管理: C++のstd::vectorを模倣する動的配列を扱うため、複雑な実行時ポインタマッピングの問題や、AMD MI300A APUでのOut-Of-Memory（OOM）エラーを避けるために、Array-of-Structures（AoS）データレイアウトを1次元配列へとフラット化（flattening）した。
   • レースコンディション: Bスプラインライブラリにおける共有一時配列の使用に起因するレースコンディションは、ローカルに宣言された自動配列への切り替えによって解決された。
   • 性能最適化: NVIDIA GPUにおいて、カーネルは一時配列の実行時割り当てにより、当初性能低下に見舞われた。これは、特定の関数に対してコンパイル時固定サイズ配列への切り替えを行うことで緩和された。この変更はターゲットアーキテクチャに基づいて条件的に適用された。
3. ハイブリッド並列化: 本研究では、ハイブル・MPI-OpenMPオフローディング戦略を評価し、リソースのオーバーサブスクリプション（GPUあたりのMPIプロセス数の増加）下での性能を評価することで、ノード割り当てが制限される大規模シミュレーションにおける実現可能性を検討した。

主な貢献

• ポータブルなGPU実装: 本論文は、OpenMPを用いてTRIMEG-C1をNVIDIAおよびAMDの両方のGPUアーキテクチャへ移植することに成功した事例を提示しており、多様なHPCシステム（Viper, Raven, Pitagoraなど）での実行を可能にした。
• コンパイラ固有の課題の解決: ポリモーフィズム、動的メモリ割り当て、および外部ライブラリの統合に関するFortranコンパイラの制限に対する、具体的な回避策の詳細を記述している。
• 性能ベンチマーク: 複数のクラスター（AMD MI300A, NVIDIA A100, NVIDIA H100）および問題サイズにわたる包括的な性能評価を実施し、個々のカーネルの加速率と強いスケーラビリティの両方を分析した。
• 検証: GPU実装の正当性は、標準的なベンチマークケースを用いてCPUの結果と比較することで厳密に検証され、移植プロセスが物理的な不正確さを導入していないことが確認された。

結果

• 加速率: フルノードにおける$32 \times 10^6$個のマーカー粒子という現実的な問題サイズにおいて、粒子プッシャーカーネルは、最速のCPUのみの実装（AMD EPYC 9754）と比較して、AMD MI300A (Viper) および NVIDIA H100 (Pitagora) の両方で約9倍の加速を達成した。
• オーバーサブスクリプション: 研究の結果、GPUリソースをオーバーサブスクライブ（ノードあたりのMPIプロセス数を増やす）すると、シリアライゼーションとコンテキストスイッチングにより個々のカーネルの加速率が低下するものの、ハイブリッドアプローチは依然として実行可能であることが分かった。ただし、全体のコードの加速率はワークロードのサイズに大きく依存する。小さなワークロード（$10^6$粒子）の場合、データ転送オーバーヘッドと算術強度の低さにより、完全に活用されたCPUノードと比較した場合の利点は著しく減少する。
• スケーラビリティ: 強スケーリング分析により、GPU加速コードは複数ノードにわたって良好にスケールすることが示されたが、粒子・トゥ・グリッド（P2G）フェーズが支配的になる場合、そのフェーズがCPUのみのノードの高いコア数から大きな恩恵を受けるため、全体的な性能向上は限定的となる。
• 正当性: CycloneベースケースおよびTCV-X21ケースの両方を用いたイオン温度勾配（ITG）モードシミュレーションの比較において、CPU版とGPU版の間で良好な一致が見られた。エネルギー成長率は、Cycloneケースでは2%未満、TCV-X21ケースでは一般に10%未満（高いトロイダルモード数において偏差が大きくなる）であった。二次元のモード構造も保持されていた。
• プロファイリング: AMD GPUにおけるプロファイリングでは、OpenMP APIが共有メモリを明示的に管理できないことに起因する、アクティブな演算ユニット（CU）を制限する高いレジスタ圧迫と、ローカルデータシェア（LDS）の利用率の低さが明らかになった。

意義と主張
本論文は、OpenMPを用いたTRIMEG-C1のGPUアーキテクチャへの移植成功が、将来の大規模な物理研究のためのポータブルな計算基盤を提供すると主張している。著者らは、コンパイラの制限を回避するために大幅なコード再構築が必要であったことを強調しつつも、結果として得られたコードは、異なるベンダーのハードウェアに対して個別のコードベースを維持することなく、多様なプラットフォーム上で電磁的不安定性のシミュレーションを可能にするものであるとしている。

本研究は、粒子プッシャーがGPU加速の主要な候補であり、粒子数が多い場合に大幅な加速をもたらすことを確立した。しかし、著者らは、現在のコードの全体的な性能は、GPU加速が容易ではない、あるいはCPUの並列化からより多くの恩恵を受ける他のフェーズ（場の方程式の求解やP2G操作など）によってボトルネックとなっていることを謙虚に述べている。したがって、GPU加速の全潜在能力は、プッシャーが実行時間の大部分を占める大規模な粒子数を持つシミュレーションにおいて実現される。

著者らが概説する今後の課題には、PETScを介したGPU対応ソルバーの統合、オフローディングの効率向上と移植作業の軽減のためのTRIMEGのC++ポート、および粒子・トゥ・グリッド操作へのGPUオフローディングの実装が含まれる。現在のGPU加速されたTRIMEG-C1は、マルチスケール乱流研究および現実的なダイバート型トカマク構成におけるプラズマ輸送の研究に向けた踏み台として位置付けられている。