Automated optimization of a molecular simulation program

16 min read1 day ago--Note: this post is a simplified overview of this paper 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 MBXThe 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.htmlThere 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 sizeRelationship 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 sizeDiagram showing how SIMD increases a program’s memory demandLivesizeOur 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] unaliveLivesize 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 removalThe initial polynomial code was generated by a scientific computation package called Maple, 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 algorithmTo 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 sizeA 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.ResultsThe 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 sizeA comparison of the program’s livesize graph after redundant subexpression removal and statement reordering (to be discussed later)This also means that this