Exascale-ready neural network interatomic potentials with CabanaMD

Computational predictions for materials require sufficiently accurate physics models which can be simulated in a reasonable amount of time. In the drive towards exascale computing, new hardware and software technologies are enabling more complex, accurate, and expensive models, but only with rethinking of algorithms, communication patterns, and data layouts. We exemplify this trend with a re-implementation of the Behler-style neural network potential (NNP), for classical molecular dynamics (MD) with near-quantum level accuracy. We use the Co-design center for Particle Applications (CoPA) Cabana particle library which, i) is built on Kokkos for on-node parallelism on various hardware, ii) provides performant particle-centric functionality, including MPI communication and neighbor lists, and iii) enables optimization of data structure for a given architecture through arrays-of-structs-of-arrays (AoSoA), intermediate between AoS and SoA. The NNP is added to the CabanaMD proxy app, where we show performance portability, including many-core CPU and GPU.