[논문 리뷰] Rapid Exploration of Optimization Strategies on Advanced Architectures using TestSNAP and LAMMPS

Rahulkumar Gayatri, Stan Moore외 6명|arXiv (Cornell University)|2020. 11. 25.

Machine Learning in Materials Science참고 문헌 12인용 수 4

한 줄 요약

이 논문은 LAMMPS 내에서 스펙트럼 이웃 분석 포텐셜(SNAP)의 성능을 22배 향상시켰으며, Kokkos를 활용한 알고리즘 재설계와 GPU 전용 최적화(예: 계층적 병렬성, AoSoA 데이터 레이아웃, 커널 분할)를 통해 초고성능 아키텍처에서 높은 계산 포화도와 성능 이식 가능성을 달성하였다.

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 구성 요소를 가진 이종 초고성능 아키텍처에서 분자의 동역학 응용 프로그램의 성능 이식 가능성과 확장성 문제를 해결한다.
• 기존 GPU 구현에서 발생하는 메모리 접근 비효율성과 로드 불균형으로 인한 성능 저하 문제를 극복한다.
• 향후 하드웨어 진화를 고려하여 Kokkos 프로그래밍 모델을 사용해 LAMMPS 내에서 SNAP의 고성능 이식 가능 구현을 개발한다.
• 현대 GPU에서 산술 강도를 극대화하고 메모리 지연을 최소화하기 위한 알고리즘적 및 저수준 최적화를 탐색하고 검증한다.
• 생산 코드에 통합하기 전 최적화 주기를 가속화하기 위해 TestSNAP 프oxy 애플리케이션을 활용한 빠른 프로토타이핑 파이프라인을 수립한다.

제안 방법

• LAMMPS에 통합하기 전 최적화 전략을 신속하게 프로토타이프 및 테스트하기 위해 TestSNAP 프록시 애플리케이션을 활용하였다.
• CPU 및 GPU 아키텍처 간 성능 이식 가능성을 확보하기 위해 CUDA 백엔드를 사용해 Kokkos를 활용해 SNAP 포텐셜을 재구현하였다.
• 3D MDRangePolicy를 사용해 배열의 배열의 구조(AoSoA) 데이터 레이아웃에 계층적 병렬성을 적용하였으며, 내부 차원은 항상 32(CUDA 워프 크기)로 고정하여 연속된 메모리 접근을 보장하였다.
• 메모리 접근을 최적화하기 위해 데이터 레이아웃을 재구성하여 로드 불균형을 제거하고 L1 캐시 재사용을 향상시켰으며, 특히 compute_Y 및 compute_U 커널에서 효과를 보였다.
• 특히 원자 연산이 포함된 compute_Y 커널에서 점유율을 향상시키고 분열을 줄이기 위해 커널 분할(kernel fission)을 사용하여 커널 융합(kernel fusion)보다 성능을 향상시켰다.
• 중간 결과를 공유 메모리 캐시에 저장하고 복잡한 데이터 구조(예: Ylist)를 실수 및 허수 성분으로 분리하여 코ales싱을 향상시키고 원자 연산 충돌을 줄였다.

실험 결과

연구 질문

• RQ1 알고리즘 재구성 및 저수준 최적화는 현대 GPU 아키텍처에서 SNAP 포텐셜의 성능을 어떻게 향상시킬 수 있는가?
• RQ2 데이터 레이아웃(예: AoSoA)은 GPU에서 연속된 메모리 접근과 로드 균형을 달성하는 데 어떤 역할을 하는가?
• RQ3 고점유율 및 원자 연산을 포함한 GPU 커널에서 커널 분할이 커널 융합보다 성능 면에서 뛰어날 수 있는가?
• RQ4 성능 이식 가능한 프레임워크(예: Kokkos)는 이종 아키텍처 간 최적화 전략을 얼마나 효과적으로 지원할 수 있는가?
• RQ5 아키텍처 특화 최적화(예: 공유 메모리, 워프 수준의 코ales싱)는 GPU에서 계산 포화도를 달성하는 데 어떤 기여를 하는가?

주요 결과

• 최종 최적화된 구현은 NVIDIA V100-16GB에서 기준 GPU 구현 대비 2J8 벤치마크에서 19.6×, 2J14 벤치마크에서 21.7×의 성능 향상을 달성하였다.
• AoSoA 데이터 레이아웃은 워프 내 로드 불균형을 완전히 제거하여, 2J8 및 2J14 문제 크기에서 각각 1.4×의 성능 향상을 가져왔으며, 이는 완벽한 코ales싱과 균형 잡힌 작업 분배 덕분이었다.
• 메모리 사용량이 크게 감소했다: 2J8 벤치마크는 GPU 메모리 0.1 GB만 사용했고, 2J14 벤치마크는 더 큰 문제 크기임에도 불구하고 0.9 GB만 사용하였다.
• 성능 향상의 대가로 CPU 성능이 저하되었으며, 이에 따라 Kokkos 구현에서 CPU 경로는 섹션 V까지의 최적화만 적용된 별도의 코드 경로가 필요로 하였다.
• 기본적인 애드조인트 재구성(의도된 재구성)은 이후 모든 최적화를 가능하게 하고 수치 안정성 및 코드 명확성을 향상시키는 데 필수적이었다.
• 모든 최적화는 공개된 LAMMPS 릴리스에 성공적으로 후방 포팅(backported)되었으며, 향후 CPU SIMD 확장은 유사한 AoSoA 원칙을 기반으로 계획되어 있다.