Memory-aware SIMD-enabled permutationally invariant polynomial code-generation procedure: Application to the MB-pol water potential

Permutationally invariant polynomials have been used to develop a variety of highly accurate potential energy surfaces. Particularly when combined with the many-body expansion, such potentials present a promising avenue for reducing the computational cost of simulations with coupled-cluster accuracy. In practice, performing molecular dynamics simulations with many-body polynomial potential energy functions requires evaluating the polynomial and its gradients thousands of times at each simulation step. If the cost of that evaluation can be reduced, larger systems and longer time scales will be available for study. Each individual gradient calculation requires evaluating tens of thousands of floating-point operations, making polynomial evaluation an excellent candidate for vectorization using single-instruction multiple-data (SIMD) instructions available on modern central processing units. However, depending on how the polynomial is factored and optimized, each evaluation may require a substantial amount of memory, producing a memory bottleneck that prevents vectorization from achieving the nominal SIMD speedup. A new memory-aware procedure for generating highly optimized polynomial routines is presented and applied to the MB-pol potential energy function as a prototypical example. The algorithm reduces the memory consumption of the routines by identifying and extracting duplicate subexpressions and by iteratively reordering the program’s statements to minimize the size of the working set. These memory optimizations allow vector instructions to be efficiently exploited, producing speedups in polynomial evaluation of nearly 4× and 8× when the AVX2 and AVX-512 vector instruction sets are employed, respectively. Using this approach, the performance of polynomial-based many-body potential energy functions can be significantly improved.