Automated optimization of a molecular simulation program

# Automated optimization of a molecular simulation program Source: [https://shishir-iyer.medium.com/automated-optimization-of-a-molecular-simulation-program-4a28a2bc05ad](https://shishir-iyer.medium.com/automated-optimization-of-a-molecular-simulation-program-4a28a2bc05ad) [](https://shishir-iyer.medium.com/?source=post_page---byline--4a28a2bc05ad-----------------------------------------) *Note: this post is a simplified overview of*[*this paper*](https://doi.org/10.1063/5.0350433)*I worked on with the Paesani Research Group at UC San Diego\.* Physics and molecular simulation programs often perform repetitive computations on vast quantities of data\. This allows them to benefit from parallelization, especially since these computations tend to be independent of each other\. MBX is a molecular simulation software package I worked on with the Paesani Research Group at UC San Diego, and it has many opportunities for optimization via parallel computation\. In particular, a section of the code referred to as the three\-body water polynomial performs numerous floating point operations for all the water molecules in the system; previously, these computations were all done sequentially\. Because water interactions contribute significantly to molecular simulation, parallelizing this part of the code had the potential for pretty big performance improvements\. However, the performance of a program is ultimately based on multiple different factors and is often a combination of memory throughput and compute speed\. Obtaining the full benefits of parallelism requires careful consideration of how the program uses memory so that the memory subsystem can keep up with the increased compute speed\. And because the polynomial evaluations are quite lengthy, we explored automated methods of code analysis and restructuring for these optimizations\. ## An overview of vectorization and its application in MBX The primary method of parallelization we explored for this task was compiler vectorization\. Most modern CPUs support vectorized instructions that allow for performing multiple computations at once\. For instance, the loop below: ``` for(int i = 0; i < 8; i++) { a[i] = b[i] + c[i];} ``` gets compiled into a vectorized instruction that does something like the following: ``` a[0:8] = b[0:8] + c[0:8] ``` The CPU has eight “lanes” with which to do these computations, so all eight addition operations can be done using a single instruction\. These instructions are thus often referred to as SIMD \(Single Instruction Multiple Data\) because they operate on eight sets of operands at once\. This allows for a speedup by up to a factor of 8\. Either 4\-wide or 8\-wide vectorization is possible depending on the instruction set used to compile the program\. Visualization of vectorization\. Source:[https://lappweb\.in2p3\.fr/~paubert/ASTERICS\_HPC/6\-6\-1\-985\.html](https://lappweb.in2p3.fr/~paubert/ASTERICS_HPC/6-6-1-985.html)There also can’t be memory dependencies between loop iterations as that would mess up the parallelism\. Therefore, the compiler can’t vectorize the code below\. ``` for(int i = 0; i < 8; i++) { // Assume a has at least 9 elements a[i + 1] = a[i] + b[i];} ``` The previously mentioned polynomial computation in MBX looks like the following, where`t`is an array used to store intermediates\. ``` t[1] = 5;t[2] = t[1] + 7;t[3] = t[1] + t[2];t[4] = 2*t[2] + t[3];t[5] = t[1] + t[3] + t[4];// ... ``` By batching the evaluations and looping through them, we can make use of compiler vectorization to evaluate the batches in parallel\. ``` // 8 wide vectorizationfor(int i = 0; i < 8; i++) { t[1*8 + i] = 5; t[2*8 + i] = t[1*8 + i] + 7; t[3*8 + i] = t[1*8 + i] + t[2*8 + i]; t[4*8 + i] = 2*t[2*8 + i] + t[3*8 + i]; t[5*8 + i] = t[1*8 + i] + t[3*8 + i] + t[4*8 + i]; // ...} ``` This initial vectorization step was a simple refactor and could easily be done with a short Python script\. We also restructured the data being passed into the procedure to ensure that corresponding elements for each batch were stored contiguously as shown below\. Press enter or click to view image in full size Relationship between original and vectorized input orderHowever, when running the new vectorized code, its speedup was still well below the theoretical speedup of 8x\. Because the code was now doing eight evaluations at once, eight times as many intermediates were required, which substantially increased the memory requirements of the program\. The original program used 32,609 intermediates for a single polynomial evaluation— being 64\-bit double precision floating point values, this equated to around 260 KB of memory\. Doing eight evaluations at once with SIMD vectorization makes the memory demand balloon all the way to nearly 2087 KB, which is way larger than the L1 and L2 caches on the CPU\. This results in a significant amount of data being read from either the slower L3 cache or even slower main memory\. Ultimately, the bulk of our work \(and thus the paper and this post\) was focused on code refactor procedures to optimize the memory usage\. Press enter or click to view image in full size Diagram showing how SIMD increases a program’s memory demand## Livesize Our primary method for keeping track of the program’s memory usage is something we’ve called the “livesize\.” At any given point in the program, some intermediates are “live” \(i\.e\. their value will be referenced at a later point in the program\), and others become “unalive” \(as their value will not be used again\)\. The livesize is thus the number of live intermediates at that point in the program\. Here’s how the livesize changes throughout the example program below: ``` // Assume t[0], t[1], t[2] previously definedt[3] = 2 * t[0]; // t[3] alivet[4] = t[0] + t[1]; // t[4] alive; t[0], t[1] unalivet[5] = t[4] + t[2]; // t[5] alive, t[2] unalivet[6] = t[4] + t[5]; // t[6] alive; t[4], t[5] unalivet[7] = t[3] + t[6]; // t[7] alive; t[3], t[6] unalive ``` Livesize throughout the example programThe maximum livesize throughout the program serves as the lower bound for how much memory we need; the array of intermediates needs to be large enough to store all the live intermediates at any given point\. Therefore, most of our memory improvements fall into two categories: bringing the number of intermediates used in the program down to the maximum livesize, and reducing the maximum livesize itself\. Overall, this latter category proved to require the most complex analysis of the program\. ## Redundant subexpression removal The initial polynomial code was generated by a scientific computation package called[Maple](https://www.maplesoft.com/), which does perform some basic optimizations on its generated routines\. While this includes removal of some duplicate subexpressions, it isn’t an exhaustive process and many still remain; for instance, the expression`t\[21\] \+ t\[22\] \+ t\[24\] \+ t\[26\]`appears at least 10 times throughout the entire file\. In this example, reassigning this expression to a new intermediate could potentially free up all four of the intermediates above, provided they weren’t used elsewhere\. This would then decrease the livesize and the amount of memory used\. Replacing repeated computations would also have the side effect of reducing the number of floating point operations the program does\. In the context of a particular subexpression, we thus defined one of its intermediates as “replaceable” if assigning the subexpression to a new intermediate actually made the intermediate go unalive earlier\. We checked this by determining whether the intermediate was used anywhere after the last occurrence of the subexpression\. If not, it satisfied this condition and would be considered replaceable\. Of course, assigning the subexpression to the new intermediate would come at the cost of introducing an additional intermediate\. Therefore, this would only be done if the subexpression had at least two replaceable intermediates \(in which case you would save more intermediates by assigning the subexpression\)\. ### The algorithm To start out, we processed each computation into a parse tree to more effectively keep track of the various operations and operands\. Python provides the`ast`module in the standard library to generate an abstract syntax tree for us, but this data structure depends on the order of the operands\. Because the operands in the redundant subexpressions sometimes showed up in varying orders throughout the file, we instead created our own parse tree data structure which kept track of children in a sorted list rather than organize the nodes in a binary tree like the original AST did\. Press enter or click to view image in full size A depiction of the parse tree data structure constructed from an example expressionThe flattening and sorting operation meant that any equivalent computations with a different order of operands would also get converted into the exact same parse tree\. Therefore, two expressions would share a common subexpression if their parse trees shared a common subtree; performing the replacement then just meant removing the subtree and replacing it with a new node\. This check was easily done by recursively traversing the tree and comparing the children\. We then checked every single pair of parse trees present in the file to see if any had a computation in common, and selected the common subtrees we found as candidates for future replacement\. Once we determined all the candidate subexpressions, we checked each line for replacement candidates\. If there were multiple candidates on the same line, we started replacement with the shared computation that had the largest number of replaceable indices\. A single pass of this algorithm eliminates most of the redundancies, but it actually introduces new ones because some of the subexpressions themselves have common subexpressions between them\. Therefore, we ran two passes of the algorithm to get rid of a majority of the redundant subexpressions\. ### Results The redundant subexpression removal was able to significantly reduce the livesize, with over a 20% reduction in maximum livesize from the first pass alone\. The second pass produces a significant yet smaller reduction\. When testing this algorithm we did try running a third pass, but this resulted in an even smaller livesize reduction, indicating that repeatedly running the algorithm produces diminishing returns\. Press enter or click to view image in full size A comparison of the program’s livesize graph after redundant subexpression removal and statement reordering \(to be discussed later\)This also means that this algorithm does not actually arrive at the definitive minimum livesize\. This is to be expected; ultimately, the algorithm is a simple greedy algorithm, which often doesn’t have the best track record with this sort of optimization problem\. Trying to minimize the number of floating point operations required for the polynomial calculation is also likely an NP\-hard problem, so any algorithm that actually solved this problem would be even slower than the one we came up with\. Ultimately, the solution we came up with is “good enough” for our purposes, and even running more passes wouldn’t really be worth it for the small improvement in results we would see\. By far the most significant improvement is actually seen from reordering statements, which will be discussed in the following section\. ## Reordering statements ### Polynomial evaluation as a DAG A computer program is generally thought of as a linear sequence of statements\. However, it often helps to view it as a graph, where each statement comprises a node in the graph\. As a statement will have various inputs and produce an output used by future statements, these dependencies will form the edges of the graph\. Specifically, we can draw an edge from Statement A to Statement B if the result of A is used as an input to B\. It’s impossible in such a structure to have a circular dependency between statements, so any program’s graph representation would be a directed acyclic graph \(DAG\)\. Therefore, an ordering of statements for the program can be thought of as a topological ordering of the program graph\. A topological ordering is an ordering of nodes in which edges can only point forward; if there is an edge from node A to node B, B must appear after A in the ordering\. As a topological ordering respects the dependencies of the DAG by definition, any ordering of statements in topological order will always yield the same result\. Different topological orderings can result in vastly different livesizes, though, so we wanted to determine which ordering minimizes the livesize\. In a DAG, this problem can actually be expressed as a concept known as the[vertex separation number](https://en.wikipedia.org/wiki/Pathwidth#Vertex_separation_number), or pathwidth\. For a given topological ordering, the vertex separation number of the graph is the smallest number*s*such that for any vertex*v*in the graph, at most*s*vertices are connected to*v*or a later vertex in the ordering\. For a given*v,*the vertices that are either attached to it or vertices after it represent statements that either depend on the statement corresponding to*v*or statements that show up after it, meaning the corresponding intermediates stay live\. Therefore, the vertex separation number and maximum livesize of a topological ordering are equivalent\. This is illustrated in the example below\. In the first ordering, the vertex separation number is 3 because for*a*₃,*a*₁ and*a*₂ are directly connected to it, and*a*₀ is connected to*a*₄, which shows up after*a*₃ in the ordering\. But moving*a*₀ further up after*a*₃ \(as is done in the second ordering\) eliminates this extra dependency, bringing the vertex separation down to 2\. How topological ordering impacts the livesize on an example DAGFinding the minimum vertex separation number over all topological orderings of a graph \(and by extension, finding the topological ordering that minimizes the vertex separation\) is also an NP\-hard problem, so no efficient algorithm currently exists for determining the optimal solution\. Therefore, we once again turn to a not\-quite\-optimal greedy algorithm to arrive at a satisfactory solution\. ### Algorithm Phase 1: Clustering The clustering phase starts with an initial topological ordering and divides it into slices that form the clusters\. We then consider what happens if we move one of the clusters of vertices to a different position in the ordering, making any other necessary adjustments to other vertices to preserve a valid topological ordering\. A cluster move can then be “scored” based on how much it reduces the mean livesize of the program\. For each cluster, we use a[ternary search](https://en.wikipedia.org/wiki/Ternary_search)to determine the optimal index to move it to that minimizes the mean livesize\. Ternary search is generally used to find the maximum or minimum value of a “unimodal” function — a function that only strictly increases until a global maximum, then only strictly decreases \(and vice versa\)\. It turns out that the score of a cluster is roughly a[convex function](https://en.wikipedia.org/wiki/Convex_function)of the index, which makes ternary search viable here \(though not always perfect\)\. We then look at each of the clusters in increasing order of their best scores, since the most optimal scores are the most negative values\. For each of these clusters, we examine each of the vertices in the cluster and remove it if doing so improves the score\. Similarly, we look at the predecessor and successor vertices — ones that have edges going into or leading out of the cluster — and add them to the cluster if doing so improves the score\. We continue this iterative process until we’ve either reached a maximum number of iterations or an iteration doesn’t modify the cluster\. Finally, we recompute the cluster score and move the cluster if the score is negative \(i\.e\. moving the cluster decreases the mean livesize\)\. Once we’ve considered all the clusters, the clustering phase is complete\. ### Algorithm Phase 2: Minimization The minimization phase considers the topological ordering at the vertex level rather than a cluster level, and tries moving each vertex to its optimal location for minimizing the mean livesize\. Since we’re only considering a single vertex at a time, we can compute the changes in livesize with each move instead of recalculating the livesize each time, which reduces the optimal index computation for each node down to linear time and eliminates the need for ternary search\. After moving all the vertices to their optimal location, we check if the decrease in mean livesize is greater than some predetermined cutoff, and repeat the process again if so\. Otherwise, the minimization phase is considered to have converged, and the algorithm goes back to the clustering phase\. These two phases alternate for a predetermined number of iterations\. ### Results As shown in the livesize graphs in the previous section on eliminating redundant subexpressions, it was this step of reordering statements that provided the greatest reduction in maximum livesize, reducing it by over 80% from the program even with most redundant subexpressions removed\. Like with the redundant subexpression algorithm, we also see a significant diminishing returns effect here as the number of iterations increases; the decrease in livesize becomes pretty insignificant after only around six iterations\. The cluster phases also display much larger reductions in livesize than the minimization phases\. This makes sense as the cluster phases are more “coarse\-grained” and move more vertices at once, albeit not quite as optimally\. Press enter or click to view image in full size Livesize as a function of the number of iterations of the reordering algorithm\. Filled circles represent cluster phases while open circles represent minimization phases## Reassigning intermediate indices The improvements discussed so far have greatly reduced the livesize, but by this point the program still uses many more intermediates than the maximum livesize\. This is because many of these are only used a few times in the program and then never again, so they stick around for a long time even after they’re no longer live\. Reassigning new expressions to unalive intermediates instead of creating new ones would resolve this issue\. ### The algorithm The reassignment algorithm ultimately ended up being pretty simple\. We first make a pass through the file, keeping track of every intermediate’s last usage\. We also maintain a list for each line of the file of all intermediates whose last usage was on a given line\. Finally, we maintain a list of all intermediates which are eligible for reuse \(no longer live\)\. Then, we make one last pass through the file, checking the following at every single line: 1. Do any intermediates go unalive on this line? If so, add them to the free list, using the list of last usages by line\. 2. Are there any assignments to an intermediate on this line? If so, replace that intermediate with one from the free list, if possible\. Additionally, maintain a list of all intermediates and their replacements if applicable\. 3. Replace any other intermediates on the line that have replacements in the list mentioned above\. ### Discussion Measuring the exact performance improvement of this step is rather difficult and subjective, since depending on when this algorithm is run the number of intermediates can be many times the livesize\. However, this algorithm will always bring the number of intermediates down to the maximum livesize of the program, so it is optimal in the sense that the program literally cannot use less memory than that\. As a reminder, this algorithm on its own doesn’t actually change the livesize and instead reduces the number of intermediates to be strictly equal to the maximum livesize\. This makes sense as we now use intermediates based on whether they’re live or not\. ## Overall Results ### Livesize As a baseline, the original program used 32,609 intermediates for a single polynomial evaluation\. Our optimization procedures were able to reduce this number all the way down to 1,457\. Now, even eight polynomial evaluations with SIMD use only 11,656 intermediates, which is still only around a third of the original memory requirement for a single evaluation\. We were able to obtain significant improvements both from decreasing the livesize of the program as well as reassigning intermediates to bring the usage down to the maximum livesize\. The maximum livesize of the original program was still only 9400, and only 22,000 of the intermediates were actually assigned or referenced anywhere \(this is probably a quirk of how Maple generates the programs\)\. Therefore, reassigning intermediates already decreased memory usage by over 50%\. Then, as discussed previously, multiple iterations of redundant subexpression removal and statement reordering decreased the maximum livesize and number of intermediates used down to 1457, or a nearly 85% decrease\. If anything, this demonstrates the utility in even non\-optimal iterative processes and greedy algorithms\. Using theoretically optimal algorithms for this problem would have likely been significantly slower and more complex and in this scenario would have yielded diminishing returns; at this point, all the data for a set of polynomial evaluations can fit comfortably in the L1 cache anyway\. ### Runtime Runtime measurements were performed using up to 32 cores of an Intel Xeon Platinum 8380 processor, using varying simulation sizes and number of OpenMP threads\. We also test both 4\-wide \(AVX2\) and 8\-wide \(AVX512\) SIMD vectorization\. For the polynomial evaluation we see a speedup of up to 4\.15x with AVX2 and 8\.05x with AVX512; the greatest speedups are observed on larger simulations with fewer threads\. It may seem a bit strange that we see speedup beyond what the “theoretical” maximum would be \(4x and 8x for AVX2 and AVX512 respectively\), but this is likely due to the fact that the original program was also slowed down by cache misses\. The increased performance on larger box sizes is also likely due to the lessening contribution of various fixed overheads to the runtime\. Meanwhile, as the number of threads increases, the SIMD speedup contributes less while the OpenMP overhead of scheduling and load balancing stays fixed, leading to less of a speedup with more threads\. Press enter or click to view image in full size Runtime per simulation step as a function of \# threads on two simulation sizes\. MBX v1\.4 here represents our optimized version of the polynomial evaluationBecause the polynomial evaluation is only one component of MBX’s total runtime, speeding it up significantly will result in a smaller speedup overall if the rest of the program remains unchanged\. Indeed, we see only around a 1\.82x speedup with 8\-wide vectorization between MBX v1\.2 \(the original version\) and MBX v1\.4 \(the optimized version\) for a 256 water molecule box, and a 2\.46x speedup on a 2048 water molecule box\. At this point, the polynomial evaluation contributes a way smaller fraction of the total runtime\. This also explains why there’s not much of a difference between 4\-wide and 8\-wide vectorization; most of the speedup has already been realized by this point\. Press enter or click to view image in full size MBX’s runtime broken down by components for each versionWith the previously time\-consuming polynomial evaluation having now been optimized, we now turn our attention to the electrostatics computation, which now dominates the runtime\. We are currently working on a full rewrite of the electrostatics computation, particularly the[particle\-mesh Ewald](https://en.wikipedia.org/wiki/Ewald_summation#Particle_mesh_Ewald_(PME)_method)\(PME\) calculation\. The refactoring will hopefully allow us to more effectively leverage parallelization techniques like SIMD and multithreading\. ## Conclusion Through the use of various algorithms borrowed from compilers and graph theory, we were able to improve the memory usage of the polynomial evaluation to effectively leverage the speedup from SIMD vectorization\. These changes have currently not been ported to GPU implementations of MBX, and many similar memory considerations to what we’ve discussed here apply to GPU parallelism\. Therefore, we hope these algorithms will provide a good foundation for future GPU work\. [MBX](https://github.com/paesanilab/MBX)and the[optimization code](https://github.com/paesanilab/polynomial_optimization_tools)are both publicly available on GitHub\. ## Acknowledgements I would like to thank co\-author[Ethan Bull\-Vulpe](https://scholar.google.com/citations?hl=en&user=tavAa0AAAAAJ)for doing most of the work on the statement reordering process and helping out with the data collection, as well as[Francesco Paesani](https://scholar.google.com/citations?user=01BwypIAAAAJ&hl=en)for providing me with the opportunity to do this research and advising me during the process\.