https://medium.com/graalvm/machine-learning-driven-static-profiling-for-native-image-d7fc13bb04e2 Open in app Sign up Sign in [ ] Write Sign up Sign in [1] Machine Learning-Driven Static Profiling for Native Image Milan Cugurovic graalvm Milan Cugurovic * Follow Published in graalvm * 11 min read * Just now -- Listen Share Machine learning (ML) enables the training of models that, based on the characteristics of a program, can accurately predict its execution. In this blog we will explain how we used ML to develop the static profiler GraalSP (a static profiler that predicts profiles in Native Image), integrated it into Oracle GraalVM Native Image, and achieved a 7.5% improvement in runtime performance! In the first part of this blog, we describe static profilers and compare them to their dynamic counterparts, which we covered in a previous blog. In that blog we discussed program profiles and how they "guide" optimizations. Using a sorting example we will delve into the design and development of GraalSP. At the end we'll also present the results of integrating ML into Native Image and discuss the deployment of the ML models into production. Dynamic and Static Profilers In AOT compilation, dynamic profilers work by building an instrumented image, collecting profiles, and then building an optimized image. They collect high-quality profiles but also have a few drawbacks. First, they complicate the optimization process by requiring two builds and a profile-collection run. This process is usually time and memory consuming, placing an extra burden on programmers and the machines used for optimization. Also, finding appropriate workloads for profile collection can be challenging. Fortunately, there is a much cheaper alternative where there's no need to go through the run-build-run compilation cycle or hunt for suitable inputs for profile collection. It's called "static profiling" or "static profile prediction". A static profiler is a profiler that doesn't collect a profile during program execution -- instead, it predicts a program's profile. A state-of-the-art static profiler takes advantage of a machine learning model to predict a profile based on a set of static features that characterize a program. Figure 1 illustrates the pipeline of a Profile-Guided Optimization (PGO) build driven by a static profiler relying on an ML model. A static profiler extracts features (we will get back to them later!) that characterize a program and then perform ML model inference to predict a profile. This predicted profile is then used just like one obtained by dynamic profiling: the compiler utilizes the information about program "execution" to create an optimized executable. This way, an ML model enables the compiler to utilize a profile without needing to run the program, reducing the usability burden and the challenge of finding suitable workloads for profile collection. Figure 1: Pipeline of a PGO build driven by an ML-based static profiler. Introducing the Graal IR It is important to note that a static profiler does not operate on a program's source code but use the compiler's intermediate representation (IR) instead. A compiler translates a program's source code into IR, a form that is more manageable and suitable for various optimizations, such as eliminating redundant code and duplicating frequent loops. In this context, Native Image uses Graal Intermediate Representation (Graal IR) to represent a program as a graph. Graal IR is a high-level IR that captures a program's structure and includes additional information that enables the compiler to perform advanced analyses and optimizations, leading to better-optimized programs. Graal IR represents a program as a graph consisting of nodes that correspond to the control flow and nodes that correspond to the data flow in the program. For example, consider the simple for loop. Graal IR translates the for loop into the IR graph shown below (Figure 2), with control flow edges colored red and data flow edges colored blue. The condition i < n is parsed into an IF node that evaluates the value of the counter i. The loop body corresponds to the true branch of the IF control split node in the IR. Therefore, executing the loop body corresponds to executing the true branch of the corresponding IF node. Figure 2: Graal IR graph corresponding to a simple for loop. You can learn more about Graal IR in one of our previous blog posts and in the paper Graal IR: An Extensible Declarative Intermediate Representation. GraalSP: The Static Profiler for Native Image Now that we are familiar with the pipeline of static profilers and Graal IR, let's discuss the ML in Native Image. We developed GraalSP: a precise, highly efficient, lightweight, polyglot, and robust static profiler as a part of the Native Image tool. Let's delve into GraalSP using an example -- we'll use the heap sort function to demonstrate how (static) profiling can help us improve performance. Running Example: Heap Sort In the code block below we show the source code of the heapSort function from the Java Development Kit (JDK). This sorting implementation uses the pushDown method to push elements down the heap. /** * Sorts the specified range of the array using heap sort. * * @param a the array to be sorted * @param low the index of the first element, inclusive, to be sorted * @param high the index of the last element, exclusive, to be sorted */ private static void heapSort(int[] a, int low, int high) { for (int k = (low + high) >>> 1; k > low; ) { pushDown(a, --k, a[k], low, high); } while (--high > low) { int max = a[low]; pushDown(a, low, a[high], low, high); a[high] = max; } } Impact of the Program Profiles on Performance Function inlining optimization determines whether to inline function calls based on the probability of executing loop bodies in for and while loops. More precisely, function inlining is a very complex optimization, and decisions about inlining a pushDown invocation aren't made solely based on the probabilities of loop body execution. However, these probabilities play a significant role in determining whether or not to inline calls. Therefore, optimization of the heapSort function highly depends on program profiles that contain information about the execution probabilities of the for and while loops. Consider sorting an array of 10 million integer values using the heapSort function from the JDK. Function inlining optimization can inline calls to the pushDown method, reducing the overhead of function calls during the sorting process. When the pushDown method calls are inlined, the average sorting time is 1.80 seconds. Conversely, if these calls are not inlined, the average sorting time increases to 2.18 seconds. As inlining calls to the pushDown method can speed up program execution by more than 20%, the execution probabilities of the loop bodies are very important. Goal: Accurately Predict Program Profiles Our goal is to predict the execution probabilities of the bodies of the for and while loop. Let's focus on, for example, for loop. Predicting the probability of executing the loop body translates to predicting the probability of executing the true branch of the corresponding IF node in the Graal IR; so our goal is to predict the execution probability of the true branch of the IF node corresponding to the for loop. To do so, first, we need to take a look at the Graal IR and "define" branches of the IF node. ML Features that Characterize Code In Figure 3, we illustrate the IF node in the Graal IR that corresponds to the for loop (named "1. If"), along with its true and false branches defined by the blocks of a control flow graph (CFG). If you're wondering about the CFG, it's essentially composed of blocks, where each block contains nodes from the IR graph. The CFG consolidates both control flow and data flow from the Graal IR, providing a clear representation of the program's execution order and enhancing comprehension of its flow. Therefore, we utilized the CFG to characterize the branches of IF statements. Figure 3: IF node that corresponds to the for and loop of the heapSort function. We "define" branches of the IF node in terms of the nodes in the Graal IR and blocks of the CFG built on top of the Graal IR. The true branch of the "1. If" node consists of blocks B2-B8, while the false branch consists of blocks B9-B24. Once we define blocks corresponding to the branches of the IF node, to fully characterize the IF node we also extract features from the CFG block that hosts the IF node (block B1 in our example) as well as blocks that point to that block in the CFG (block B0). The features we extract include, for example, the estimated assembly size of a branch, the estimated number of CPU cycles the computer will perform to execute instructions from the branch, and the nesting loop depth of an IF node. The margin of this blog is too narrow to discuss all the features. For those interested in more details, our paper GraalSP: Polyglot, efficient, and robust machine learning-based static profiler provides more details. Here, it is important to note that the feature characterizing the IF node, for which we aim to predict the probability of executing its true branch, is represented by a vector of floating point numbers. Machine Learning for Profile Prediction After extracting the feature vector, we use supervised learning techniques to train the ML model to predict profiles. We use the XGBoost ensemble of decision tree (DT) models for regression to predict profiles. A DT model uses a tree-like structure consisting of nodes, branches, and leaves to model the data. Each node evaluates a feature and directs decisions down the tree, while leaf nodes predict profiles. On the left side of Figure 4, we illustrate a shallow DT model that predicts the branch execution probability of a true branch based on the loop depth of an IF node, as well as the assembly size and number of CPU cycles of a branch. For example, if an IF node is at a depth of two or more and the estimated CPU cycles for a true branch are less than five, the DT model predicts the execution probability of that branch to be 0.15. The main benefits of DTs are their interpretability, speed, and ease of use. To improve prediction performance, reduce overfitting, and increase accuracy, we utilize the XGBoost ensemble, combining DTs as weak learners. On the right side of Figure 4, we illustrate an ensemble consisting of 1,500 decision trees. Each tree in the ensemble predicts the execution probability for a target branch, and the ensemble aggregates all predictions (for example, by averaging them) to determine the outcome. Usage of the ensembles enables us to achieve precise predictions, while usage of the decision trees enables us to achieve a highly efficient static profiler. Also, by using lightweight decision trees, we end up with a small model of only ~250 kilobytes. Figure 4: DT and XGBoost ML models. By defining the static profiler on top of the IR of the Graal compiler, we've created a polyglot static profiler capable of predicting profiles for all programs compiled to Java bytecode (for example, Java, Scala, and so on!). Furthermore, we designed and developed two profile prediction heuristics that handle scenarios where input data deviates from the real-world Java and Scala programs. Of course, these heuristics can't predict profiles better than ML models, but they can help prevent the model from making mistakes. For instance, one of the heuristics ensures that the probability of executing the body of a loop is not less than 0.2. This way, the ML model won't make cardinal mistakes on frequent loops. By slightly increasing the binary size of the generated programs, these heuristics have enabled us to create a robust static profiler that effectively handles outliers. The impact of a static profiler on program performance can vary depending on the program's code. To evaluate the impact, we use programs from test suites Renaissance, DaCapo, and DaCapo con Scala. We integrated GraalSP into Native Image and achieved a 7.46% speedup in execution time (geometric mean) compared to the default configuration. The default image build assumes a uniform distribution of execution probabilities over the branches of a control split. Figure 5 reports the geomean speedup of programs aggregated according to the test suite. Figure 5: Runtime speedup of GraalSP across the test suites Renaissance, DaCapo, and DaCapo con Scala. Runtime improvement comes at the cost of an average increase in the size of generated programs of only 3.9%. The increased size of generated binary files results from duplication-based compiler optimizations. The static profiler generates profiles for the entire codebase, even though usually only a small section of the code is executed. Duplication-based optimizations duplicate code even in sections that won't be executed, thus slightly increasing the size of the compiled programs. It is important to emphasize that the performance of programs optimized using PGO with dynamically collected profiles is much higher than those optimized using profiles predicted by GraalSP. In our experiments, PGO with dynamically collected profiles produces programs that are 33% faster and 15% smaller. As expected, the quality of the dynamically collected profiles is superior because they lack the inherent errors found in ML models. Additionally, dynamically collected profiles include information that GraalSP currently does not predict: method call profiles, reached types in virtual calls, details about monitor locking and unlocking, and so on. Of course, as there is no such thing as a free lunch, these performance improvements comes with the hurdles of dynamic profiling. Model Inference: Deployment to Production Since the releases of GraalVM for JDK 17 and GraalVM for JDK 20 (version 23.0), GraalSP has been enabled by default in Oracle GraalVM ! During the default build, you'll find out that Native Image optimizes your program using PGO with ML-inferred profiles. Figure 6 illustrates the integration of GraalSP and dynamic profiling in Native Image. When users enable dynamic profiling, Native Image instruments the code and runs the instrumented executable to collect profiles. Otherwise, if users do not enable the dynamic profiler, Native Image will run the GraalSP to predict profiles and optimize programs based on the predicted values. Figure 6: Integration of the GraalSP and dynamic profiler in the Native Image. GraalSP uses the ONNX Java Runtime, developed by Microsoft for cross-platform accelerated ML. This ensures compatibility with various architectures, including Windows amd64, Linux amd64 and aarch64, and Darwin amd64 and aarch64, providing flexibility and usability of GraalSP across different platforms. Also, by using the ONNX runtime, we implemented static profiling with a compilation time overhead of only 10%. Conclusion To summarize, GraalSP is enabled by default in Oracle GraalVM Native Image and offers the benefits of PGO without requiring profile collection. GraalSP characterizes programs in terms of Graal IR and uses the XGBoost ensemble to predict profiles. Native Image then uses these predicted profiles to perform PGO and create optimized programs. This represents the first successful application of ML in Oracle GraalVM. However, at Oracle Labs, we are preparing some exciting updates, improvements, and ML novelties to be announced very soon. Stay tuned! For more details about the ML magic inside Oracle GraalVM, please refer to our paper GraalSP: Polyglot, efficient, and robust machine learning-based static profiler. Graalvm Native Image Machine Learning Xgboost Onnx Graalvm -- -- graalvm graalvm Published in graalvm 2.6K Followers *Last published just now GraalVM team blog - https://www.graalvm.org Milan Cugurovic Milan Cugurovic Follow Written by Milan Cugurovic 0 Followers *1 Following Senior Researcher at Oracle Labs. Ph.D. Candidate and Teaching Assistant at the Faculty of Mathematics, University of Belgrade. ML and compilers. Follow No responses yet Help Status About Careers Press Blog Privacy Terms Text to speech Teams