ParaGraph: Weighted Graph Representation for Performance Optimization of HPC Kernels

Recently, with breakthroughs in Machine Learning (ML) and specifically Deep Learning (DL), researchers have been applying data-driven approaches to various software engineering tasks, and challenges, ranging from code comment generation [3] to compiler optimizations [4]. However, the source code can not simply be fed to DL models as DL models do not accept strings as input. Therefore, applications need to be presented in a way usable by the models. Some works have presented programs as a sequence of tokens [5, 6]. While this kind of representation has worked well for some downstream tasks, presenting programs as a sequence of tokens discards critical syntactical and semantic characteristics of programs which are essential for deep learning models to reason.

Recently, to better present programs’ semantics and syntactic information, graph-based representations have been proposed. Allamanis et al. [7] proposed enhanced ASTs and added new edges to present control flow information. They applied the augmented AST to two tasks: variable misuse detection and variable name suggestion. Inspired by their work, we also augment the AST by adding new edges. However, we incorporate additional edges to better represent $loops$ and $if$ conditions.

Cummins et al. in [4] proposed a lower-level graph representation based on LLVM intermediate representation for solving compiler optimization tasks.These representations are effective for downstream tasks, such as simple algorithm classification. To the best of our knowledge, there does not exist a representation tailored toward representing the characteristics of parallel loops and if-statements.

When different architecture supports either a different native language or various optimization techniques for the same language, program portability becomes a serious problem. Regrettably, programming languages and compilers are still far from being able to handle portability on their own. Program portability should be the primary priority for developers who have to deal with multiple device memory accesses on the same data or someone whose programs must work effectively on systems with various node architectures, such as many-core vs. GPU-accelerated nodes.

Utilizing a directive-based programming paradigm, such as OpenMP, the de-facto standard for parallel programming in C/C++ and Fortran, is one approach to ensuring portability across several architectures. OpenMP intends to move to extremely heterogeneous architectures [8] and has supported GPU offloading from specification 4.0. Unfortunately, even with OpenMP, optimizing large-scale applications remains a challenging task.

Since GPUs are very parallel machines, programmers should take full advantage of this fact to get the most performance possible. The “omp parallel for” pragma alone will only parallelize a code for CPUs; it won’t offload the computation to a GPU. We expect a high level of coarse grain parallelism on a platform like GPU. The amount of parallelism that a GPU can use is constrained by this design. Figure 1 illustrates how OpenMP teams and distribute directives create additional levels of parallelism. At the start of a target region, only one team and one member thread are active. The teams distribute directive distributes the full loop iteration space among all available teams. Additionally, in the case of nested parallel loops, we can utilize the parallel for directive on them to further distribute the iterations of the nested loops among threads within a team. We utilize the combined directive “teams distribute parallel for” to distribute the iteration space of one loop among teams and threads inside a team when there is just one level of parallel loops or when the outer loop has sufficient parallelism. teams are used to group threads, and distribute enables a team group to be scheduled to run in a loop. Teams resemble CUDA threadblocks on NVIDIA GPUs, as there is no synchronization primitive to function as a barrier between threads from different teams. And the threads within each team are parallelized typically using parallel for.

There are several frameworks being worked on right now that will automatically assist application developers in handling severe heterogeneity. These frameworks necessitate analytical cost models, which will help them select the kind of optimization the application requires. Hand-tuned cost functions are extensively employed currently; however, calculating optimization costs requires a deeper understanding of the underlying hardware. Despite its efficacy, manually building a cost model for a single architecture can take months. Since cost functions are crucial and manual tuning is time-consuming, compiler engineers are investigating Machine, and Deep Learning approaches as a means of automating this process.

Recently, a lot of research has been done on how to use learning-based techniques in compilation as well. Early work exploiting ML in compilers primarily explored its use to help optimize sequential programs. However, its application to the task of optimizing parallel programs has attracted significant attention in the past decade due to the prevalence of multi-core platforms and, more recently, heterogeneous systems [9, 10]. Mirka et.al. [11] devised a decision-tree-based approach to predict the scheduling policy for an OpenMP parallel region. Also, Denoyelle et.al. [12] uses machine learning techniques to optimize OpenMP programs for scheduling policies and the number of threads. Alcaraz et.al. [13] uses neural networks and performance counters to predict the number of threads. Dutta et.al. [14] uses an LLVM-IR based graph representation and performance counter to train a GNN model to predict number of thread, chunk size and scheduling policy for OpenMP loops. Tree and graph-based features have also been used by Malik et.al. [15], who present a unique graph-based approach for feature representation. Learning-based techniques were used to build classifiers to determine whether to offload OpenCL code [16] and to select a clock frequency at which the processor should operate [17]. A high level of accuracy was reported; however, the benefits could not be quantified as the work did not attempt to generate modified code. They also explored regression techniques to build curve fitting models to search for the sweet spot for work partitioning among processors [18] or a trade-off of energy and performance [19].

The results of prior efforts from applying ML on compiler optimizations are encouraging. However, new feature engineering practices need to be explored that can help learning-based methods to learn more about a code and its computational needs. In COMPOFF [2], the authors provide a proof of concept for using an Artificial Neural Network model to make better decisions in offloading a kernel to a GPU using OpenMP. They utilize a fully-connected feed-forward network, also referred to as multi-layer perceptrons (MLPs), which are effectively stacked layers of linear regression.

OpenMP Advisor [20] is a first-of-its-kind compiler tool that enables OpenMP code offloading to a GPU via Machine Learning. The Advisor is divided into three major modules: Kernel Analysis, Cost Model, and Code Transformation. The Kernel Analysis module recommends various variants for a given application, and the Code Transformation module generates codes for those variants. In our work, we use the code transformation module of OpenMP Advisor to generate various kernels for training our model. Through the use of COMPOFF, a machine learning-based cost model, this tool can identify the kernels that are most suitable for offloading by predicting their runtime. But COMPOFF has some limitations. It requires figuring out how many operations are contained within a kernel, which is a challenging task in and of itself. The Advisor’s current functionality is only restricted to GPUs, though it has the potential to be expanded to support additional OpenMP-capable hardware. One technique for expanding to other devices is to develop a model that works outside GPUs as well. Extending COMPOFF to work beyond GPUs is a daunting task, and there is a need for a new cost model that can bridge this gap.

A recent tool called BLISS [21] allows for auto-tuning parallel applications without the use of instrumentation or domain-specific knowledge. This auto-tuner develops a variety of lightweight models with the aid of Bayesian optimization that can compete with the output quality of a complex, heavyweight Machine Learning (ML) model. In the same context, Bayesian optimization is frequently used by other autotuners, including BOHB [22], HiPerBOt [23], GPTune [24] and ytopt [25]. Unfortunately, due to their expensive evaluation functions, tuning large-scale HPC systems, which are becoming increasingly heterogeneous, complex, and expensive to evaluate, is still quite challenging. Nonetheless, most of the autotuners mentioned above are online, meaning that even in the inference time, they execute applications iteratively to tune them. Therefore, in terms of complex heterogeneous large-scale HPC systems, these tuners will be very costly to apply. In contrast, $ParaGraph$ is an offline approach and does not require the execution of applications in the inference phase. COMPOFF  [2] provides a novel portable cost model that statically calculates the cost of OpenMP offloading on various architectures. COMPOFF is also an offline approach. In this work, we compare $ParaGraph$ to COMPOFF since both can predict the runtime of a kernel on a GPU statically.

$ParaGraph$ can be used by existing GNN models, which are a type of neural network that can operate over graphs, for downstream tasks. Especially, those tasks that are related to performance optimizations. $ParaGraph$ can be exploited both as a heterogeneous graph or as a homogeneous graph. For simplicity, in this study, we treat $ParaGraph$ as a homogeneous graph. Typically a graph is shown as Equation 1, where $V$ is a set of nodes and $E$ contains an adjacency matrix. Elements of $E$ are either 1 or 0 to show whether an edge exists between two nodes.

We extend the AST for our new representation by including new edges like NextToken, NextSib, and other types that were explained earlier. We formally define $ParaGraph$ as follows:

Where $V$ and $E$ are previously defined in Equation 1 and $T\in\mathbb{Z}^{+}$ represents the type of the edges (such as NextToken, NextSib etc). $W\in\mathbb{Z}^{+}$ presents the weight which is zero for any edge type other than Child.

These additional edges and attributes are all added to the graphs statically. Therefore one obvious question could be why we do not let the model reason about these additional attributes by itself rather than presenting them in the graph? On one hand, adding these attributes to the graphs is not costly; on the other hand, we do not want to waste the resources and the learning capabilities of the model to realize these obvious facts. Therefore by augmenting the ASTs and including the attributes mentioned earlier, we can train GNN models more effectively. Later on, our ablation study further improves that the model benefits from the attributes and information presented in $ParaGraph$.

We adapt Relational Graph Attention Networks (RGAT) [27] and use $ParaGraph$ representation as input to train a model for predicting runtime of applications on different accelerators. In RGAT, attention logits are computed per each edge type. In the result section, we will see that with $ParaGraph$, the model has a very small error in its predictions and outperforms the current state-of-the-art approach.

Figure 3 shows the overall workflow of our GNN-based pipeline for runtime prediction. The first step is preparing different variants of an application for each accelerator used in this study. This step is further explained in detail in the next section. Then, using Clang, ASTs are produced and a series of augmentations, as discussed, are applied to them to construct $ParaGraph$. In order to train a model to predict the runtime of an application, we need to create a dataset with a list of applications and their variant accompanied by their runtime. Therefore, we execute each one of the variants on the specified accelerator and measure the runtime. Lastly, the $ParaGraph$ representation of variants and their runtime are used to train the GNN model. Along with $ParaGraph$s, our feature set also includes the number of teams and threads used for executing an application. These features are then fed into the GNN-based model for predicting the runtime.

There are three parts to our data collection: Code Variant Generation, Graph Generation, and Runtime Collection.

Once we have enough data on both the Summit and Corona clusters, we begin training our model. We implemented a GNN-based neural network using RGAT [27] as the convolution layers. Our model is implemented using Pytorch-Geometric library with Mean Squared Error as the loss function, and Adam [33] as the optimizer. To embed the graph, the model uses three graph convolution layers based on RGAT, followed by two fully connected layers with ReLu activation function. As mentioned, the number of teams and threads are considered as two additional features. Another fully connected layer is used to embed these two features. Finally, the embedding of the graph and the two features are concatenated together and passed through the last fully connected layer for runtime prediction.

The edge weights and two additional features are normalized using MinMaxScaler. The dataset is split into train-validation sets using a 9:1 ratio.

To evaluate the performance of $ParaGraph$, we use $RMSE$ which is Root Mean Square Error (Equation 3).

Where $x_{i}$ stands for the runtime of a data point in microseconds, $\hat{x}_{i}$ is the predicted runtime by $ParaGraph$, and $N$ is the total number of samples.

Since the range of runtime differs across platforms, the normalized version of $RMSE$ is also considered. Normalized $RMSE$ is calculated by dividing the $RMSE$ by the distance between the minimum and maximum runtime. We also use relative error (i.e., absolute error divided by the range of runtime) to report the error rate.

Table III shows the experimental results for each accelerator. We have NVIDIA V100, AMD MI50, and IBM POWER9 with 22 cores and AMD EPYC 7401 with 24 cores. As shown, the $RMSE$ values range from 280 (ms) to 4325 (ms). The reason why we have different $RMSE$ values, such as 4325 (ms) for POWER9 accelerator, lies in the fact that the runtime dispersion differs across the accelerators. Standard Deviation in Table II gives us some insights on how dispersed the collected data are. Moving on to Normalized-RMSE, which is independent of the range of runtime, we see $ParaGraph$ has relatively the same error across accelerators; thus, it can be applied to different accelerators.

To further analyze our results, we have calculated a relative error per bins of 10 seconds. Figure 4 shows 11 bins for all four accelerators. Each bin has a 10-second range except the last bin. The figure shows that the relative error is small across different bins and accelerators (less than %10), meaning that our $ParaGraph$ model has stable behavior for different ranges of runtimes.

Moreover, we analyze how stable the model is during the training process in Figure 5. The figure shows validation $RMSE$ for four accelerators. In the first few epochs, the $ParaGraph$ model is not very stable, resulting in fluctuations in the $RMSE$; however, as the model is trained further, it is able to pick and learn the features from the representation and reduce $RMSE$ value per each epoch and ultimately converges.

Lastly, we calculate the average relative error per application to see if the $ParaGraph$ model is able to have prediction with low error for all types of applications. Figure 6 shows the error rate of each application. As can be seen, the model is indeed capable of having accurate predictions for a wide variety of applications resulting in a low error rate. Therefore, the model is not biased toward one application. On the AMD MI50 GPU, the Laplace data was corrupted during collection. Consequently, neither this study nor the training process includes that data.

$ParaGraph$ representation, as explained in Section III, is built by applying a series of major augmentations on top of AST. In this section, we quantify the impact of these augmentations. First, we consider the AST itself without additional edges and weights, we call it Raw AST. Then, we add additional edges and edge types to the AST and name it Augmented AST. Lastly, we have $ParaGraph$ that contains both additional edges and also edge weights. Table IV shows the results of the ablation study. We see that Raw AST results in the highest error for all four accelerators. Adding new edges and introducing new types of edges (Augmented AST) improve the prediction to some extent. For example, the $RMSE$ of V100 drops from 2114 (ms) to 786 (ms) with the addition of these new edges.

One of the key characteristics of our proposed program representation is the edge weights. Edge weights convey essential information about how often different regions of the AST execute. Therefore, we see quite a good improvement in $RMSE$ when edge weights are added. For instance, the $RMSE$ for V100 is further improved to 280 (ms).

We analyze the addition of edges and their weights further to see how the training process of the $ParaGraph$ model is affected by these augmentations.

Figure 7 shows how the model behaves during training. This figure depicts the $RMSE$ value per each epoch for Raw AST, Augmented AST, and $ParaGraph$ on the MI50 accelerator.

Using only the Raw AST without any augmentations, which means having only one edge type, the model is able to learn some characteristics of the applications and reduce the $RMSE$ per epoch; however, this reduction in $RMSE$ is not significant. Augmented AST contains 8 different types of edges. We see the addition of these edges destabilized the training process of the model. In the first few epochs, the model is challenged to learn different relations between the nodes. However, eventually, after several epochs, the prediction of the model is stabilized and it achieves $RMSE$ of 1177 (ms). Once edges of the Augmented AST are augmented with weights, thus $ParaGraph$ is constructed, we see further improvements in the model’s predictions. Although the validation $RMSE$ has some fluctuation in the initial epochs, it ultimately converges with a considerably smaller error.

To the best of our knowledge, $COMPOFF$ [2] is the only state-of-the-art OpenMP GPU offloading cost model. We further compare the results from $ParaGraph$ with those of $COMPOFF$. As mentioned in Section II-D OpenMP Advisor uses $COMPOFF$ for predicting runtime for OpenMP kernels.

While OpenMP Advisor eventually needs a cost model that can forecast for all possible underlying architectures, $COMPOFF$ is currently only suitable for GPU execution. In contrast, $ParaGraph$ can model the application runtime regardless of the underlying architecture. In our experiments, $ParaGraph$ was used to predict runtime on both CPUs and GPUs. Therefore, we compare $ParaGraph$ and $COMPOFF$ using data points collected from the NVIDIA V100 GPU only.

For both $ParaGraph$ and $COMPOFF$, Figure 9 demonstrates a strong correlation between the actual and predicted data. The results for $COMPOFF$ are represented by blue dots, while those for $ParaGraph$ are represented by orange dots. Despite the fact that both perform admirably, $ParaGraph$ demonstrates a much stronger correlation between the predicted and actual runtime. In Figure 8, $COMPOFF$ (blue) demonstrates a slightly higher error rate for smaller runtime kernels, but as the runtime increases, this error rate decreases (or disappears entirely). However, for all kernels, the error rate is significantly lower for $ParaGraph$ (orange). Similar results were also observed on the AMD GPU. We were unable to compare the outcomes of our prediction on CPUs because $COMPOFF$ does not support CPUs. This is one significant advantage $ParaGraph$ has over $COMPOFF$.

$ParaGraph$ is able to model the execution of applications statically. Our experimental results show the effectiveness of $ParaGraph$ in predicting applications’ runtime across different accelerators. On the other hand, there are applications whose runtime behavior is not stable. That is, they show different behaviors across various executions. These applications can not be modeled statically. As a result, predicting their runtime using only static features is a challenge. One way of tackling this challenge could be developing a hybrid approach using static and dynamic features (for example, Performance Counters) of applications. However, collecting dynamic features of applications requires the execution of applications which will potentially increase the cost of prediction.