[論文レビュー] Rapid Exploration of Optimization Strategies on Advanced Architectures using TestSNAP and LAMMPS

Rahulkumar Gayatri, Stan Moore+6 more|arXiv (Cornell University)|Nov 25, 2020

Machine Learning in Materials Science参考文献 12被引用数 4

ひとこと要約

この論文は、KokkosとGPU固有の最適化（階層的並列化、AoSoAデータレイアウト、カーネル分離）を用いたアルゴリズムの再設計により、LAMMPSにおけるスペクトル近接分析ポテンシャル（SNAP）で22倍の性能向上を達成した。これにより、エクサスケールアーキテクチャを横断した高い計算飽和度とパフォーマンスポータビリティが実現された。

ABSTRACT

The exascale race is at an end with the announcement of the Aurora and Frontier machines. This next generation of supercomputers utilize diverse hardware architectures to achieve their compute performance, providing an added onus on the performance portability of applications. An expanding fragmentation of programming models would provide a compounding optimization challenge were it not for the evolution of performance-portable frameworks, providing unified models for mapping abstract hierarchies of parallelism to diverse architectures. A solution to this challenge is the evolution of performance-portable frameworks, providing unified models for mapping abstract hierarchies of parallelism to diverse architectures. Kokkos is one such performance portable programming model for C++ applications, providing back-end implementations for each major HPC platform. Even with a performance portable framework, restructuring algorithms to expose higher degrees of parallelism is non-trivial. The Spectral Neighbor Analysis Potential (SNAP) is a machine-learned inter-atomic potential utilized in cutting-edge molecular dynamics simulations. Previous implementations of the SNAP calculation showed a downward trend in their performance relative to peak on newer-generation CPUs and low performance on GPUs. In this paper we describe the restructuring and optimization of SNAP as implemented in the Kokkos CUDA backend of the LAMMPS molecular dynamics package, benchmarked on NVIDIA GPUs. We identify novel patterns of hierarchical parallelism, facilitating a minimization of memory access overheads and pushing the implementation into a compute-saturated regime. Our implementation via Kokkos enables recompile-and-run efficiency on upcoming architectures. We find a $\sim$22x time-to-solution improvement relative to an existing implementation as measured on an NVIDIA Tesla V100-16GB for an important benchmark.

研究の動機と目的

• CPUとGPUを組み合わせた多様なコンponentを有する異種のエクサスケールアーキテクチャにおいて、分子動力学アプリケーションのパフォーマンスポータビリティとスケーラビリティを向上させること。
• メモリアクセスの非効率性と負荷の不均衡に起因する、従来のSNAP相互作用ポテンシャルのGPU実装におけるパフォーマンスボトルネックを克服すること。
• 将来のハードウェア進化に対応できる、Kokkosプログラミングモデルを用いた高性能でポータブルなSNAPの実装を開発すること。
• 現代のGPUにおける算術強度の最大化とメモリ遅延の最小化を実現するためのアルゴリズム的および低レベルの最適化を調査・検証すること。
• LAMMPSに統合する前の最適化サイクルを加速するために、TestSNAPプロキシアプリケーションを用いた迅速なプロトタイピングパイプラインを構築すること。

提案手法

• LAMMPSに統合する前の最適化戦略の迅速なプロトタイピングとテストのために、TestSNAPプロキシアプリケーションを採用した。
• CPUおよびGPUアーキテクチャを横断したパフォーマンスポータビリティを確保するため、CUDAバックエンドを用いてKokkosを用いてSNAPポテンシャルを再実装した。
• 3次元MDRangePolicyを用いた階層的並列化を実装し、配列の配列の構造（AoSoA）データレイアウトにマッピングした。内側の次元は32（CUDAワープサイズ）に固定し、メモリアクセスのコalescingを実現した。
• データレイアウトの再構築により、特にcompute_Yおよびcompute_Uカーネルにおいて負荷の不均衡を解消し、L1キャッシュの再利用を向上させた。
• 特にアトミック演算が関与するcompute_Yにおいて、占有率の向上と分岐の低減を図るため、カーネル融合ではなくカーネル分離を採用した。
• 中間結果の共有メモリキャッシュ化と、Ylistなどの複雑なデータ構造を実部と虚部に分割することで、コalescingの向上とアトミック競合の低減を実現した。

実験結果

リサーチクエスチョン

• RQ1 アルゴリズムの再構築と低レベルの最適化は、現代のGPUアーキテクチャにおけるSNAPポテンシャルのパフォーマンスをどのように向上させるか？
• RQ2 データレイアウト（例：AoSoA）は、GPUにおけるコalescedメモリアクセスと負荷のバランスをどのように実現するか？
• RQ3 高占有率とアトミック演算を含むGPUカーネルにおいて、カーネル分離はカーネル融合を上回る性能を発揮できるか？
• RQ4 パフォーマンスポータブルフレームワーク（例：Kokkos）は、異種アーキテクチャに跨る最適化戦略のポータビリティをどの程度実現できるか？
• RQ5 アーキテクチャ固有の最適化（例：共有メモリ、ワープレベルのコalescing）は、GPUにおける計算飽和度の達成にどの程度寄与するか？

主な発見

• 最適化された最終実装は、NVIDIA V100-16GB上でベースラインGPU実装と比較して、2J8ベンチマークで19.6倍、2J14ベンチマークで21.7倍の高速化を達成した。
• AoSoAデータレイアウトにより、ワープ内での負荷の不均衡が解消され、2J8および2J14の両問題サイズで1.4倍の高速化が得られた。これは、完全なコalescingと均等な作業分配のおかげである。
• メモリ使用量が顕著に削減された：2J8ベンチマークではGPUメモリがわずか0.1 GB、2J14ベンチマークでは0.9 GBで、問題サイズが大きくなっても抑えられた。
• パフォーマンス向上の代償としてCPUパフォーマンスが低下したため、Kokkos実装ではCPUパスとGPUパスを分離する必要があり、CPUパスはセクションVまでの最適化のみを適用した。
• 基本的な随伴因子再構築（adjoint refactorization）は、以降のすべての最適化を可能にし、数値安定性とコードの明確さを向上させるために不可欠であった。
• すべての最適化は、公開LAMMPSリリースにバックポートされた。今後の作業として、同様のAoSoA原則を用いたCPU SIMD拡張が計画されている。