[论文解读] 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 中实现了 Spectral Neighbor Analysis Potential (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 组件）上分子动力学应用的性能可移植性与可扩展性挑战。
• 克服现有 SNAP 互作用势 GPU 实现中的性能瓶颈，包括内存访问效率低下和负载不平衡问题。
• 利用 Kokkos 编程模型，在 LAMMPS 中开发高性能、可移植的 SNAP 实现，以支持未来硬件演进。
• 探索并验证能最大化现代 GPU 上算术强度、最小化内存延迟的算法级与底层优化技术。
• 建立基于 TestSNAP 代理应用的快速原型设计流程，以在集成到生产代码前加速优化迭代周期。

提出的方法

• 使用 TestSNAP 代理应用快速原型化并测试优化策略，再集成到 LAMMPS 中。
• 采用 Kokkos 实现 SNAP 势，后端为 CUDA，以确保在 CPU 和 GPU 架构间的性能可移植性。
• 通过 3D MDRangePolicy 映射到数组-结构-数组（AoSoA）数据布局，实现分层并行，其中最内层维度固定为 32（CUDA warp 大小），以实现内存访问合并。
• 通过重构数据布局优化内存访问，消除负载不平衡，提升 L1 缓存重用率，尤其针对 compute_Y 和 compute_U 内核。
• 采用内核分裂而非内核融合，以提高占用率并减少分支发散，尤其在涉及原子操作的 compute_Y 中效果显著。
• 引入共享内存缓存中间结果，并将复杂数据结构（如 Ylist）拆分为实部和虚部分量，以增强内存合并并减少原子冲突。

实验结果

研究问题

• RQ1 算法重构与底层优化如何提升现代 GPU 架构上 SNAP 势的性能？
• RQ2 数据布局（如 AoSoA）在实现 GPU 上的内存访问合并与负载均衡中起到何种作用？
• RQ3 在高占用率和原子操作较多的 GPU 内核中，内核分裂是否优于内核融合？
• RQ4 性能可移植框架（如 Kokkos）在异构架构间能否实现可移植的优化策略？
• RQ5 架构特定优化（如共享内存、 warp 级内存合并）在实现 GPU 计算饱和方面起到多大作用？

主要发现

• 最终优化实现相较于 NVIDIA V100-16GB 上的基线 GPU 实现，在 2J8 基准测试中实现了 19.6× 加速，在 2J14 基准测试中实现了 21.7× 加速。
• AoSoA 数据布局消除了 warp 内部的负载不平衡，使 2J8 和 2J14 问题规模均获得 1.4× 加速，得益于完美内存合并与均衡的工作分配。
• 内存使用显著降低：2J8 基准仅使用 0.1 GB GPU 内存，2J14 基准使用 0.9 GB，尽管问题规模更大。
• 性能提升以牺牲 CPU 性能为代价，因此 Kokkos 实现中需引入分支代码路径，CPU 路径仅使用至第五节的优化策略。
• 基础伴随导数重构是实现后续所有优化的关键，同时提升了数值稳定性和代码清晰度。
• 所有优化均已成功回溯至公开的 LAMMPS 发布版本中，未来工作计划基于相似的 AoSoA 原理，为 CPU SIMD 扩展提供支持。