A Unified Heterogeneous Implementation of Numerical Atomic Orbitals-Based Real-Time TDDFT within the ABACUS Package

ABACUS パッケージにおいて、数値原子軌道に基づくリアルタイム TDDFT の計算を高速化するための統一異種コンピューティングフレームワークを開発し、単一 GPU や複数 GPU 環境での高性能かつスケーラブルな実装を確立しました。

要約

この論文は、**「物質の超高速な電子の動きを、スーパーコンピューター（特に GPU）を使って、これまでよりもはるかに速く、正確にシミュレーションできる新しい仕組みを作った」**という内容です。

専門用語を避け、日常の例え話を使って解説します。

1. 何をしたの？（物語の舞台）

物質の中を走る電子は、光（レーザーなど）を当てると、一瞬で（フェムト秒という単位で）動き出します。この「電子の暴れっぷり」を計算して、新しい太陽電池や超高速な電子デバイスを設計しようとしています。

これまでは、この計算をするには**「CPU（普通のパソコンの頭脳）」を使っていましたが、計算量が膨大すぎて、「1 週間かかる計算を、1 日で終わらせたい」**という欲求がありました。

そこで、この研究チームは、**「GPU（ゲームや AI で使われる、並列処理が得意な計算機）」**をフル活用できる新しいシステムを「ABACUS」というソフトウェアの中に作りました。

2. 3 つの階層（新しい工場の仕組み）

このシステムは、まるで**「現代的な工場のライン」**のように 3 つの層に分かれています。

• ① 利用者層（お客様）：
  • 研究者が「この分子の構造を計算して」と入力する場所です。
  • 例え： 料理人が「今日の献立（入力データ）」を注文するカウンター。
• ② 開発者層（料理人）：
  • 物理の法則（電子がどう動くか）を計算するロジックです。
  • 例え： 料理人が「まず野菜を切り、次に炒め、最後に味付けをする」という**レシピ（アルゴリズム）**に従って作業する場所。
• ③ 基盤層（厨房と道具）：
  • ここが今回の最大の特徴です。CPU でも GPU でも、同じように動けるようにした「万能な道具箱」です。
  • 例え： 料理人が使う包丁やフライパンが、**「どんな種類のキッチン（CPU や GPU）でも、同じように使えるように設計された」**という状態です。料理人は「包丁の持ち方（メモリの管理）」を気にせず、ただ「切る（計算する）」ことだけに集中できます。

3. 最大の難所を解決した（「速度の壁」を突破）

この計算には、**「速度ゲージ（Velocity Gauge）」**という、物理的に正確だが計算が非常に重い方法があります。

• 問題点： これまで、この方法を使うと、計算の**「90% の時間」が「電子の位置に合わせた特殊な計算（球面上の積分）」**に使われてしまい、ボトルネック（渋滞）になっていました。
  • 例え： 高速道路を走る車（計算）が、**「信号待ち（特殊な計算）」**でほとんど止まってしまうような状態です。
• 解決策： 研究チームは、この「信号待ち」を GPU の得意技で処理する**「専用レーン」**を作りました。
  • 結果： これにより、その部分の計算速度が**「12 倍」**に向上しました。
  • 効果： 結果として、「速度ゲージ」という正確な方法が、もはや「遅い」という欠点を持たなくなりました。研究者は、正確さを犠牲にせず、好きな方法を選べるようになりました。

4. どれくらい速くなったの？（成果）

• 1 枚の GPU vs 56 コアの CPU：
  • 従来の「56 個の頭脳（CPU）」がフル稼働しても、**「1 枚の GPU（A800）」の方が「3〜4 倍速い」**ことがわかりました。
  • 特定の重い計算部分では、**「12 倍」**の差がつきました。
• 大規模な計算：
  • 原子が 1000 個以上ある大きな分子でも、40 枚の GPUを並列に使って効率的に動かすことができました。

5. まとめ（この研究の意義）

この研究は、単に「計算を速くした」だけでなく、**「将来のコンピューター（NVIDIA だけでなく、AMD や中国製の DCU など）が変わっても、同じコードがそのまま動くように」という、「ハードウェアに依存しない万能な土台」**を作った点が画期的です。

一言で言うと：

「電子の超高速な動きをシミュレーションする『料理』を、これまで『手作業（CPU）』でやっていたのを、**『最新鋭のロボットキッチン（GPU）』で自動化し、さらに『最も時間のかかる工程』を劇的に短縮して、『どんなキッチンでも使える万能レシピ』**を完成させた」

これにより、将来の新材料開発や、光と物質の相互作用の解明が、これまで想像もできなかったスピードで進むことが期待されます。

技術概要

以下は、提示された論文「A Unified Heterogeneous Implementation of Numerical Atomic Orbitals-Based Real-Time TDDFT within the ABACUS Package」の技術的な要約です。

1. 研究の背景と課題 (Problem)

実時間時間依存密度汎関数理論（RT-TDDFT）は、励起状態の電子ダイナミクスや光 - 物質相互作用をシミュレーションする上で不可欠な手法ですが、計算コストが非常に高く、大規模システムや長時間シミュレーションには高性能計算（HPC）リソースが求められます。特に、数値原子軌道（NAO）に基づく RT-TDDFT 実装において、以下の課題が存在していました。

• GPU 加速の不足: 既存の GPU 対応コードの多くは基底状態の DFT や線形応答 TDDFT に焦点が当てられており、リアルタイムダイナミクス、特に局所基底（NAO）を用いた手法の GPU 加速は未開発でした。
• 速度ゲージ（Velocity Gauge）の計算ボトルネック: 周期系における光 - 物質相互作用を記述する際、速度ゲージは物理的に適切ですが、非局所擬ポテンシャルに位置依存の位相因子（$e^{-i\mathbf{A}(t)\cdot\mathbf{r}}$）が現れます。NAO 基底では、この項を解析的に評価できず、実空間グリッド上で数値積分を行う必要があり、計算コストが膨大になるという深刻なボトルネックとなっていました。
• ハードウェア依存性と保守性: GPU 向けコードの開発には明示的なメモリ管理やアーキテクチャ固有のチューニングが必要であり、コードの可搬性や長期維持が困難でした。また、NVIDIA GPU、AMD GPU、Hygon DCU などの多様なアクセラレータに対応する統一フレームワークが不足していました。

2. 提案手法とアーキテクチャ (Methodology)

著者らは、オープンソース第一原理計算パッケージ「ABACUS」内に、ハードウェアに依存しない統合的な異種計算フレームワークを構築しました。このフレームワークは、物理アルゴリズムとデバイス実装を分離する 3 層構造を採用しています。

• 統一データコンテナ（Tensor）:
  • Hamiltonian 行列や密度行列などの多次元配列を管理するための統一コンテナ「Tensor」を導入しました。
  • メモリ管理（CPU/GPU/DCU への自動割り当て）やデータレイアウトをカプセル化し、ソースコードレベルでのハードウェア特化を不要にしました。これにより、単一のコードベースで複数のアーキテクチャ上で効率的に実行可能になりました。
• 統一線形代数演算子:
  • 行列積（GEMM）、LU 分解（getrf）、線形方程式求解（getrs）などの標準的な線形代数ルーチンに対して、ハードウェア非依存のポリモーフィックインターフェースを提供しました。
  • CPU では BLAS/LAPACK、GPU/DCU では cuBLAS/cuSOLVER や対応するライブラリを自動的に呼び出す仕組みを構築しました。
• 統一グリッド積分インターフェース:
  • 実空間グリッド上の物理量（電荷密度、ポテンシャル、力など）の計算を抽象化しました。
  • 球面グリッド積分の GPU 加速: 速度ゲージにおける非局所擬ポテンシャルの位相因子積分を、原子中心の球面グリッド（Lebedev-Laikov 角積分と Gauss-Legendre 径方向積分）上で GPU キーネルとして最適化しました。これにより、従来の CPU 実装に比べて劇的な高速化を実現しました。

3. 主要な貢献 (Key Contributions)

• NAO ベースの RT-TDDFT の完全な異種実装: ABACUS において、数値原子軌道に基づく RT-TDDFT の全プロセス（波動関数の時間発展、ハミルトニアン構築、力評価、エレンフェスト動力学）を GPU 上で実行可能にしました。
• 速度ゲージのボトルネック解消: 位置依存位相因子を含む非局所擬ポテンシャルの積分を GPU 上で効率的に処理する専用カーネルを開発し、速度ゲージ使用時の計算コストを劇的に削減しました。
• ポータブルで保守可能なフレームワーク: 「Tensor」抽象化レイヤーにより、物理アルゴリズムとハードウェア実装を分離し、NVIDIA、AMD、Hygon などの異なるアクセラレータへの対応を容易にしました。
• 大規模マルチ GPU 実装: 分散環境での波動関数伝播を効率的に行うため、QR 分解に基づく並列ソルバ戦略を採用し、多数の GPU での強スケーリングを実現しました。

4. 結果と性能評価 (Results)

• 物理的妥当性の検証:
  • 有機分子（アントラセン）、ナノクラスター（(CdSe)6）、1 次元水素鎖、2 次元 h-BN、3 次元バルクシリコンなど、多様な次元性と化学結合を持つ系で光学特性（吸収スペクトル、誘電関数）を計算しました。
  • 長さゲージ、速度ゲージ、ハイブリッドゲージの結果が互いに整合しており、既存の平面波基底や他の NAO コード（SIESTA, CP2K, Octopus 等）のベンチマークと高い一致を示しました。
• 単一 GPU での性能向上:
  • バルクシリコン（最大 1200 原子）のシミュレーションにおいて、単一の NVIDIA A800 GPU は、56 コアの Intel Ice Lake CPU ノード（フル稼働）と比較して、3〜4 倍のウォールクロック時間短縮を達成しました。
  • 特に、速度ゲージ特有の球面積分カーネルは、最適化された CPU ベースラインと比較して12 倍以上の高速化を達成し、速度ゲージのアルゴリズム的ペナルティをほぼ排除しました。
  • 波動関数伝播（$O(N^3)$）の計算時間は、GPU 使用により CPU の最良構成に対して約 6〜7 倍、単一 MPI 構成に対して 12 倍以上高速化されました。
• マルチ GPU 強スケーリング:
  • 最大 40 GPU（10 ノード、1728 原子）でのテストにおいて、約 76% の並列効率を維持しました。
  • 波動関数伝播モジュールは良好なスケーリングを示しましたが、線形スケーリングのグリッド積分モジュールは通信オーバーヘッドによりある程度飽和しました。しかし、全体計算時間におけるグリッド積分の割合が GPU 加速により相対的に減少しているため、全体のスケーリング性能には大きな影響を与えませんでした。

5. 意義と展望 (Significance)

この研究は、局所基底（NAO）の計算効率と GPU 加速の高性能さを融合させた、大規模かつ長時間の第一原理電子ダイナミクスシミュレーションのための高品質でポータブルなプラットフォームを確立しました。

• 科学的意義: 速度ゲージの計算コスト問題を解決したことで、周期系における非平衡電子ダイナミクス（高調波発生、光誘起相転移など）を、より物理的に正確なゲージ選択でシミュレーションできるようになりました。
• 技術的意義: ハードウェアに依存しない抽象化レイヤーの設計は、将来の新しいアクセラレータ（DCU など）への対応や、AI 科学（AI for Science）との統合（自動微分など）への道を開くものとして重要です。
• 将来展望: 今後の課題として、cuBLASMp/cuSOLVERMp への依存を解消し、AMD や Hygon などの非 NVIDIA ハードウェアでも分散マルチ GPU 計算が可能なポータブルな分散ライブラリへの対応、および Tensor コアを活用した混合精度演算の導入が挙げられています。

総じて、この論文は、超高速電子ダイナミクスの研究において、大規模複雑材料のシミュレーションを現実的な計算時間で可能にする画期的な進展を提供しています。