SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ali TehraniJamsaz\*, Quazi Ishtiaque Mahmud\*, Le Chen, Nesreen K. Ahmed2, Ali Jannesari
<sup>1</sup> Iowa State University, Ames, Iowa, USA
{tehrani, mahmud, lechen, jannesari}@iastate.edu
$^{2}$  Intel Labs, Santa Clara, CA, USA
nesreen.k.ahmed@intel.com
The remarkable growth and significant success of machine learning have expanded its applications into programming languages and program analysis. However, a key challenge in adopting the latest machine learning methods is the representation of programming languages, which directly impacts the ability of machine learning methods to reason about programs. The absence of numerical awareness, aggregate data structure information, and improper way of presenting variables in previous representation works have limited their performances. To overcome the limitations and challenges of current program representations, we propose a graph-based program representation called PERFOGRAPH. PERFOGRAPH can capture numerical information and the aggregate data structure by introducing new nodes and edges. Furthermore, we propose an adapted embedding method to incorporate numerical awareness. These enhancements make PERFOGRAPH a highly flexible and scalable representation that effectively captures programs' intricate dependencies and semantics. Consequently, it serves as a powerful tool for various applications such as program analysis, performance optimization, and parallelism discovery. Our experimental results demonstrate that PERFOGRAPH outperforms existing representations and sets new state-of-the-art results by reducing the error rate by  $7.4\%$  (AMD dataset) and  $10\%$  (NVIDIA dataset) in the well-known Device Mapping challenge. It also sets new state-of-the-art results in various performance optimization tasks like Parallelism Discovery and NUMA and Prefetchers Configuration prediction.
In recent years, the remarkable success of machine learning has led to transformative advancements across numerous fields, including compiler optimization and program analysis. The applications include compiler heuristics prediction, optimization decisions, parallelism detection, etc. [4, 15]. The training process generally involves feeding program data as input and transforming it into a representation suitable for machine learning models. The selection of program representation is crucial, as it can significantly impact the performance of the machine learning model [12]. With the development of graph neural networks (GNNs), an increasing number of graph representations of programs have been incorporated into GNN models for program analysis [3, 28, 10]. One of the pioneering efforts in developing a comprehensive graph representation for programs is PROGRAML [14]. PROGRAML incorporates control, data, and call dependencies as integral components of a program's representation. In contrast to prior sequential learning systems for code, PROGRAML
closely resembles the intermediate representations used by compilers, and the propagation of information through these graphs mimics the behavior of typical iterative data-flow analyses. Despite the success that PROGRAML has achieved, there are shortcomings in this current state-of-the-art program representation, especially for performance-oriented downstream tasks. These limitations stem from neglecting numerical values available at compile time and the inadequate representation of aggregate data types.
In this paper, we present PERFOGRAPH to address the limitations of the current state-of-the-art program representation. Additionally, we propose a novel way to embed numbers in programs in an elegant way so that our DL model will not face unknown numbers during inference time. Our experiments demonstrate that PERFOGRAPH sets new state-of-the-art results in numerous downstream tasks. For example, in the Device Mapping downstream task, PERFOGRAPH yield error rates as low as  $6\%$  and  $10\%$  depending on the target hardware. Moreover, PERFOGRAPH even outperforms the tools and models specially designed for specific tasks such as parallelism discovery.
Overall, the main contributions of this paper are:
- An enhanced compiler and language-agnostic program representation based on PROGRAML that represents programs as graphs.
- The proposed representation supports aggregate data types and provides numerical awareness, making it highly effective for performance optimization tasks.
- Evaluation of the proposed representation on common downstream tasks and exceeding the performance of PROGRAML.
- Quantification of the proposed approach on a new set of downstream tasks such as parallelism discovery and configuration of NUMA systems.
The rest of the paper is structured as follows: Section 2 presents the related works. In section 3, we provide a motivational example, showing the limitations of PROGRAML, the state-of-the-art program representation. This section is followed by section 4 where we present our proposed representation PERFOGRAPH along with the novel way of embedding numerical values. In section 5, experimental results on downstream tasks are provided, and finally, Section 6 concludes the paper and discusses some future works.
Machine learning has brought significant advancements in many fields, and program analysis and software engineering are no exceptions. However, Machine Learning (ML) and Deep Learning (DL) models can not directly process raw source code to reason about programs. Therefore, researchers have explored different approaches to represent applications in a format suitable for DL models. Generally, there are three types of commonly used program presentations: sequence of tokens, Abstract Syntax Tree (AST), and Intermediate Representation (IR).
Sequence of tokens: The initial attempts [16, 21, 27] represented source code as a sequence of tokens, such as identifiers, variable names, or operators. This approach intuitively treats programming languages similarly to natural languages. It allows for the utilization of advanced natural language process (NLP) techniques, such as large language models [19, 22, 23]. However, this token-based representation overlooks the inherent dependency information within the program's structure. It fails to capture the unique relationships and dependencies between different elements of the code, which can limit its effectiveness in tasks such as compiler optimization and code optimization.
AST: An AST represents the structure of a program by capturing its hierarchical organization. It is constructed based on the syntactic rules of the programming language and provides a high-level abstraction of the code. Previous works have leveraged ASTs as inputs to tree-based models for various code analysis tasks like software defect prediction [17] and code semantic study [9]. Moreover, there have been efforts to augment ASTs into graphs that incorporate program analysis flows such as control flow and data flow. These AST-based graph representations capture more comprehensive code dependency information and have shown superior results compared to traditional approaches in previous works [2, 3].
IR: IR is an intermediate step between the source code and the machine code generated by a compiler. Previous work [39] has utilized IR to train an encoding infrastructure for representing programs as a distributed embedding in continuous space. It augments the Symbolic encodings with the flow of information to capture the syntax as well as the semantics of the input programs. However, it
generates embedding at the program or function level and also requires a data-flow analysis type for generating the embedding. In contrast, our approach derives embedding from the representation and works at the more fine-grained instruction level. More recent works [6, 14, 8] have leveraged IR-based graph representation to better capture essential program information, such as control flow, data flow, and dependencies. However, despite their success, IR-based graph representations have certain limitations. For example, these representations may not be numeric-aware or may lack the ability to adequately represent aggregate data types. In this work, we propose PERFOGRAPH, a graph representation based on IR, to address these limitations.
As stated in the related work section, program representations based on the intermediate representation are very effective in enabling DL models to automate the process of various optimizations. One such representation is PROGRAML, whose performance surpasses other code representations, making it state-of-the-art in various optimizations and downstream tasks. However, despite its potential, it suffers from several limitations. To name a few: it is incapable of properly carrying out information regarding read and write operations to the memory location, has no support for aggregate data types, and discards numerical values. Listing 1 shows a code snippet where a 3-dimensional array is defined. Figure 1 shows the PROGRAML representation of this code snippet. For illustration purposes, instruction nodes and control flow edges are shown in blue, whereas red represents variable, constant nodes, and data flow edges. Green edges show the call graph. As it can be seen, PROGRAML fails to represent some critical information. For instance, code float arr [2] [3] [4] is converted to LLVM IR  $[2x[3x[4xfloat]]]$ , which is used to construct a node in PROGRAML. It eliminates the aggregate data structure information, like the array's dimension. Leaving it up to the DL model to infer the meaning behind the numbers in  $[2x[3x[4xfloat]]]$ . Moreover, in this representation, only the type of numbers (e.g., int8, float) are considered, and the actual values of the numbers are not given attention. The absence of numerical awareness limits the performance of PROGRAML in applications where numerical values play an important role. A numerically aware representation can help understand and optimize operations involving numeric data types, constants, and expressions. There are also some anomalies in the way temporary variables are depicted. For example, in 1, we see the fourth alloca node allocates memory for a variable, and two store instructions are applied on two separate nodes representing the variable. Thus, the information about the first store instruction is not carried out properly when the second store instruction happens. In the following section, we will see how PERFOGRAPH effectively addresses many limitations in the current state-of-the-art program representation. PERFOGRAPH uses PROGRAML representation as its initial graph and reconstructs the graphs by addressing the aforementioned limitations.
Listing 1:  $\mathbf{C} + +$  code example.
1 #include <stdio.h>
2
3 int main(int argc, char *argv[]) {
4     int arr_state = 0;
5         float arr[2][3][4]
6             = {{1.0f, 1.0f, 1.0f, 1.0f},
7                 {2.0f, 2.0f, 2.0f, 2.0f},
8                 {3.0f, 3.0f, 3.0f, 3.0f},
9                 {4.0f, 4.0f, 4.0f, 4.0f},
10                 {5.0f, 5.0f, 5.0f, 5.0f},
11                 {6.0f, 6.0f, 6.0f, 6.0f}
};
Figure 1: ProGRAML representation of Listing 1.
4 PERFOGRAPH: A fined-grained numerical aware graph representation
PERFOGRAPH is graph representation based on LLVM IR. It is built on top of PROGRAMML; however, it does not suffer from the limitations that the PROGRAMML has, helping DL models to reason over the complex structure of the programs and enabling them to make more accurate optimization decisions,
especially in terms of performance optimization. Figure 2 shows how various enhancements and improvements are applied to construct a more precise representation. Consider a simple code example of defining a variable and increasing it by one  $\{\mathrm{int} \quad \mathrm{i} = 0; \mathrm{i}++;\}$ . Figure 2a shows the PROGRAML representation of this code example. It has two store nodes, as one is responsible for storing 0 and the other one storing the incremented value of i. In the following subsection, we will explain how PERFOGRAPH is constructed by addressing the limitations shown in Figure 2a.
(a) Initial graph
(b) Local variables
(c) store nodes
Figure 2: PERFOGRAPH addresses the existing limitations in program representation.
(d) Numbers
4.1 Representing Local Identifiers and store instruction
Local Identifiers: Local identifiers' names are preceded by % in LLVM Intermediate representation. Memory allocation on the stack is done by alloca instruction. One of the limitations of the current state-of-the-art program representation, PROGRAML, is that it is unable to carry out information regarding the operations that happen to a memory location. For instance, in Figure 2a, the two store nodes represent storing values of 0 and 1 to variable i. However, as shown, each store instruction node is connected to a separate variable node, making it difficult for the graph neural network to reason over the operations that happen to a memory location. For the embedding vector of the second store node in 2a to represent the fact that some information regarding the variable i has changed, one has to increase the number of GNN layers to 3 to support up to 3 hops when propagating the messages in GNN. This can potentially limit the ability of the GNN model if there are a greater number of hops between the two store nodes shown in Figure 2a. To address this limitation, instead of having more than one variable node (oval-shape nodes) per identifier, PERFOGRAPH only considers one variable node in its graph representation. Any load or store instruction will refer to the same variable node. These changes are shown in Figure 2b. We see that the store nodes in Figure 2b access the same memory location, thus representing the fact that those store instructions are modifying the same memory location.
Store instruction: LLVM uses store instruction to write into memory. store instruction has two arguments: a value to store and the address to which it will store the value. PROGRAML differentiates between these two arguments by adding a position feature to the edges as shown in Figure 2a. However, since the store instruction modifies the contents at the corresponding memory address, we posit that it is better to reflect the fact the content of the identifier has changed. To present this information, PERFOGRAPH adds an extra edge from the store node to node representing the identifier whose value is modified by the store instruction. Figure 2c shows these changes in the graph constructed by PERFOGRAPH.
Numbers: Numbers can be a significant factor in optimization decisions. For example, they can show the loop bound, and different optimizations can be considered depending on the loop bound. PERFOGRAPH, unlike PROGRAML, not only considers the type of numbers such as i32, i64, float but also the actual values of the numbers. As illustrated in Figure 2d, numerical constant nodes have the actual value of the number in their feature set in addition to the type of the number. Even though numerical constant nodes have the value of the number as one of their features, there is a need to embed the numbers in a way that unknown numbers will not be seen in the inference. Unlike other tokens, numbers are harder to embed as an infinite amount of numbers exist, and to handle all ranges of numbers, we need to have a very large vocabulary set.
4.2 Numerical Awareness
However, with a very large vocabulary size, the DL models may still encounter numbers in the inference phase that they have not seen in the training phase. We propose a novel way of embedding numbers called Digit Embedding. Figure 3 shows our approach. To embed a number, we first break down the number to its digits; then, we consider a position for each one of the digits. The goal is to let DL models realize the place value of each digit. Then, each digit and its corresponding position are embedded and summed together. Therefore, we will have an embedding representing the information about the digits and their positions. For instance, in Figure 3, we embed each digit and its corresponding position with an output dimension of 3. Since the number has four digits, the results would be a vector/tensor of size  $4 \times 3$ . To make sure the Digit Embedding of numbers has the same length across numbers
with varying sizes of digits, we apply an aggregation function over the embedding dimension. Since the output embedding dimension is three in this example, we would have one vector of length three representing the number after aggregation. The aggregation function can be of any type (Max, Mean, etc.).
Figure 3: Overview of the digit embedding.
4.3 Aggregate Data Types
Aggregate data types, such as arrays and vectors, are an essential part of applications. They play an important role in many applications, such as matrix multiplications. Thus, presenting these data types helps the DL models better understand programs. Current LLVM IR-based program representations fail to present aggregate data types appropriately. For example, consider a three-dimensional integer array. In LLVM IR, this array is shown as  $3 \times [2 \times [3 \times i32]]$ . As can be seen, the length of the arrays and their data types are inferable. However, without proper representation, the DL model's capacities will be spent on learning these deterministic facts (i.e., the length of the arrays and their type). PERFOGRAPH considers aggregate data types as a new node type in its representations. Figure 4b shows how aggregate data types are supported by PERFOGRAPH. Unlike other LLVM IR-based representations, PERFOGRAPH supports multi-dimensional arrays and vectors. PERFOGRAPH creates a chain of nodes to present the different dimensions of the arrays. In Figure 4a, we see there is a node representing the three-dimensional array  $[3 \times [2 \times [3 \times i32]]]$ . PERFOGRAPH breaks down the corresponding node into three (since it is a three-dimensional array) white nodes as shown in Figure 4b. Then, each node has a context representing that specific dimension of the array. For example, the context for the third dimension is  $[3 \times i32]$ , whereas for the second dimension, the context is  $[2 \times [3 \times i32]]$ . For each aggregate type node, in addition to the context of the node,
we specifically add the length of the array and its type as additional features. For aggregate data types whose lengths are not known during compile time, we follow the LLVM conventions by considering the length of those data types as vscale. These enhancements will help the DL models to reason over the dimensions and types of aggregate data types. As a result, PERFOGRAPH will ultimately enable the DL models to have more accurate predictions for applications that deal with arrays and vectors.
(a) No support for aggregate data types.
(b) Break-down of aggregate data types.
Figure 4: PERFOGRAPH supports aggregate data types.
5 Experimental Results and Downstream Tasks
In this section, we evaluate PERFOGRAPH on six downstream tasks. For each downstream task, we will explain the task itself, the dataset, and the baselines.
In our experiments, we use DGL's [40] RGCN [37] implementation for PERFOGRAPH representation. The graphs from PERFOGRAPH are treated as heterogeneous and managed via the HeteroGraphConv module. We use a hardware setup of two 18-core Intel Skylake 6140 CPUs and two NVIDIA Tesla V100-32GB GPUs. The embedding space for numbers is generated by extracting digits and positions from a numeric token of an IR statement, then passed to a PyTorch [31] embedding layer for digit and position embeddings. These are combined for the final numeric token embedding. Non-numeric tokens directly go through the PyTorch embedding layer. Each PERFOGRAPH heterogeneous node converts to a 120-dimensional vector via this embedding. We use the Adam [26] Optimizer, relu [1] activation function, a learning rate of 0.01, and hidden_dim parameter of 60. Mean aggregation is applied to combine different node-type results before a linear classification layer, which outputs a probability distribution for each class. The class with the highest probability is the prediction.
5.2 Device Mapping
Problem Definition: We apply PERFOGRAPH to the challenging heterogeneous device mapping [13] problem. In this task, there are a number of OpenCL kernels that we need to predict which accelerator (CPU or GPU) yields higher performance. We compare PERFOGRAPH against DeepTune [13], Inst2Vec [6], and PROGRAML [14]. The results of the baselines are quoted from [14].
Dataset: For this task, we use the dataset published in [13]. In this dataset, there are 256 OpenCL kernels available, and 680 LLVM IR instances are extracted from them. There are
Figure 5: Performance comparison the device mapping task with state-of-the-art models [lower is better].
two types of GPUs: AMD and NVIDIA. For each of the GPUs, the runtimes of the kernels are recorded in the dataset. For AMD, 276 kernels show better performance in GPU, while 395 kernels show better performance in CPU. Whereas for NVIDIA, 385 kernels have better runtimes with GPU, and 286 kernels have better runtimes with CPU. We consider this as a binary CPU or GPU classification problem.
Results: As the dataset is small, we use the same 10-fold validation (with  $80\%$  training,  $10\%$  validation, and  $10\%$  testing) like PROGRAMML [14] and chose the model with the highest validation accuracy. The hand-crafted features of [20] are also used as graph-level features in our model to enhance the performance following the approach in [14]. Table 1 and 2 show the final precision, call, f1-score, and accuracy for AMD and NVIDIA devices. Figure 5 compares PERFOGRAPH with state-of-the-art models on the Device Mapping dataset. We can see that PERFOGRAPH sets new state-of-the-art results by achieving the lowest error rate among the baselines both for AMD and NVIDIA, indicating the effectiveness of PERFOGRAPH.
Table 1: PERFOGRAPH results for AMD devices.
Precision	Recall	F1-score	Accuracy
CPU	0.94	0.94	0.94	0.94
GPU	0.94	0.94	0.94
Table 2: PERFOGRAPH results for NVIDIA devices.
Precision	Recall	F1-score	Accuracy
CPU	0.87	0.90	0.89	0.90
GPU	0.92	0.89	0.90
5.3 Parallelism Discovery
Problem Definition: In this problem, given a sequential loop, we try to predict whether a loop can be executed in parallel. We treat this problem as a binary classification problem with two classes: Parallel and Non-Parallel.
Dataset: The OMP_Serial dataset [11] is used for this task. It contains around 6k compileable source C files with Parallel and Non-Parallel loops. The training dataset contains around 30k IR files. The OMP_Serial dataset has three test subsets to compare the performance with three traditional parallelism assistant tools: Pluto (4032 IR files), AutoPar (3356 IR files), and DiscoPoP (1226 IR files).
Results: We evaluate PERFOGRAPH on all three subsets and compare it with traditional rule-based tools: Pluto [7], AutoPar [33], DiscoPoP [29], and also Deep Learning based tools: Graph2Par [11], PROGRAML. Table 3 shows the results. The results of Pluto and Graph2par are reported from [11]. As PROGRAML does not have this downstream task in their paper, we used the PROGRAML representation in our pipeline to generate the results. Results show that traditional rule-based tools have the highest precision but the lowest accuracy because those tools are overly conservative while predicting parallel loops. So, they miss out on a lot of parallelism opportunities. PERFOGRAPH achieves considerably good precision scores across all the test subsets. In terms of accuracy, PERFOGRAPH surpasses the current state-of-the-art approaches by  $2\%$  in the Pluto and AutoPar subset. In the DiscoPoP subset, it achieves an impressive  $99\%$  accuracy and surpasses PROGRAML by  $9\%$ .
Table 3: Performance comparison of PERFOGRAPH on the OMP_SerI dataset.
Subset	Approach	Precision	Recall	F1-score	Accuracy
Pluto	Pluto	1	0.39	0.56	0.39
Graph2par	0.88	0.93	0.91	0.86
PROGRAML	0.88	0.88	0.87	0.89
PERFOGRAPH	0.91	0.90	0.89	0.91
autoPar	AutoPar	1	0.14	0.25	0.38
Graph2par	0.90	0.79	0.84	0.80
PROGRAML	0.92	0.69	0.67	0.84
PERFOGRAPH	0.85	0.91	0.85	0.86
DiscoPoP	DiscoPoP	1	0.54	0.70	0.63
Graph2par	0.90	0.79	0.84	0.81
PROGRAML	0.92	0.94	0.92	0.91
PERFOGRAPH	0.99	1	0.99	0.99
5.4 Parallel Pattern Detection