LEARNING EXECUTION THROUGH NEURAL CODE FUSION
Authors: Zhan Shi
Kevin Swersky
Daniel Tarlow
Parthasarathy Ranganathan
Milad Hashemi
Venue:  (2020)
Keywords: code understanding
graph neural networks
learning program execution
execution traces
program performance

Abstract: As the performance of computer systems stagnates due to the end of Moore’s Law,
there is a need for new models that can understand and optimize the execution
of general purpose code. While there is a growing body of work on using Graph
Neural Networks (GNNs) to learn static representations of source code, these
representations do not understand how code executes at runtime. In this work, we
propose a new approach using GNNs to learn fused representations of general
source code and its execution. Our approach defines a multi-task GNN over
low-level representations of source code and program state (i.e., assembly code
and dynamic memory states), converting complex source code constructs and data
structures into a simpler, more uniform format. We show that this leads to improved
performance over similar methods that do not use execution and it opens the door
to applying GNN models to new tasks that would not be feasible from static code
alone. As an illustration of this, we apply the new model to challenging dynamic
tasks (branch prediction and prefetching) from the SPEC CPU benchmark suite,
outperforming the state-of-the-art by 26% and 45% respectively. Moreover, we
use the learned fused graph embeddings to demonstrate transfer learning with high
performance on an indirectly related algorithm classification task.

Published as a conference paper at ICLR 2020
LEARNING EXECUTION THROUGH NEURAL CODE FUSION
Zhan Shi"
The University of Texas at Austin zshil700 ishil7@utexas edu
Kevin Swersky, Daniel Tarlow, Parthasarathy Ranganathan, Milad Hashemi Google Research fkswersky, dtarlow, parthas, miladh]]@googee.co edu
ABSTRACT
As the performance of computer systems stagnates due to the end Moore's Law, there is need for new models that can understand and optimize the execution of general purpose code. While there a growing body of work on using Graph Neural Networks (GNNs) to learn static representations of source code, these representations do not understand how code executes at runtime. In this work, we propose new approach using GNNs learn fused representations of general source code and its execution. Our approach defines a multi-task GNN over low-level representations of source code and program state (i.c., assembly code and dynamic memory states), converting complex source code constructs and data structures into simpler, more uniform format. We show that this leads to improved performance over similar methods that do not use execution and it opens the door to applying GNN models to new tasks that would not be feasible from static code alone. As an illustration of this, we apply the new model to challenging dynamic tasks (branch prediction and prefetching) from the SPEC CPU benchmark suite, outperforming the tfat-oo-theeart by 26% and 45% respectively. Morcover, we use the learned fused graph embeddings to demonstrate transfer learning with high performance indirectly related algorithm classification task.
INTRODUCTION
Over the last 50 years, hardware improvements have led to exponential increases in software per- formance, driven by Moore's Law. The end of this exponential scaling has enormous ramifications for computing (Hennessy & Patterson 2019) since the demand for compute has simultaneously grown exponentially, relying on Moore's s Law to compensate (Ranganathan) (017). ) As the onus of performance optimization shifts to software, new models, representations, and methodologies for program understanding are needed to drive research and development in computer architectures, compilers, and to aid engineers in writing high performance code.
Deep learning has emerged as powerful framework for solving difficult prediction problems across many domains, including vision (Krizhessky et al. 2012), speech (Hinton et al. 2012), and text (Sutskever et al. 2014). Recent work has started to frame many canonical tasks in computer architecture analogous prediction problems, and have shown that deep learning has the potential to outperform traditional heuristics (Hashemi et al. 2018). In this work, we focus on two representative tasks: address prefetching (modeling floa-flow during execution) [Jouppi] 1990; Wenisch et al.) 2009 Hashemi et al. 2018) and branch prediction (modeling control-flor -foow during execution) (Jiménez & Lin 2001 | Seznec 2011 Smith 1981 Traditional models for solving these tasks memorize historical access patterns and branch history to make predictions about the future. However, this approach inherently limited as there are simple cases where history- methods cannot generalize
Work completed during an internship at Google.
As Moore's Law ends, prediction techniques in these fields have also stagnated. For example, the winner of the recent branch prediction championship increased precision by 3.7% (Dundas 2016).

---

Published as a conference paper at ICLR 2020
Instead, we argue that these tasks Stranch-predietiee and prefetching) jointly model the intermediate behavior of a program as it executes. During execution, there is a rich and informative set of features in intermediate memory states that models can learn to drive both prediction tasks. Additionally, since programs are highly structured objects, static program syntax can supplement dynamic information with additional context about the program 's execution.
(Section 4.6p.
We combine these two sources of information by learning a representation of a program from both its static syntax and its dynamic intermediate state during execution. This incorporates new set of previously unexplored features for prefetching and branch prediction, and we demonstrate that these can be leveraged obtain significant performance improvements. Inspired by recent work on learning representations of code (Allamanis et al. (017), our approach is distinguished by two aspects. First, instead of using high level source code, we construct a new graph representation of low-level assembly code and model it with a graph neural network. Assembly makes operations like register reads, memory accesses, and branch statements explicit, naturally allowing us to model multiple problems within a single, unified representation. Second, to model intermediate state, we propose a novel snapshot mechanism that feeds limited memory states into the graph (Section 3.2)
We call our approach neural code fusion (NCF). This same representation can easily be leveraged for bevy of other low-level optimizations (including: indirect branch prediction, value prediction, memory disambiguation) and opens up possibilities for multi-task learning that were not previously possible with traditional heuristics. NCF can also be used to generate useful representations of programs for indirectly related downstream tasks, and we demonstrate this transfer learning approach on an algorithm classification problem.
On the SPEC CPU2006 benchmarks (Sta 2006), NCF outperforms the state-of-the-ard in address and branch prediction by significant margin. Moreover, NCF is orthogonal to existing history- based methods, and could easily combine them with our learned representations to potentially boost accuracy further. To our knowledge, NCF is the first instance of a single model that can learn simultancously on dynamic control-flow and data-flow tasks, setting the stage for teaching neural network models to better understand how programs execute.
In summary, this paper makes the following contributions:
An extensible graph neural network based representation of code that fuses static code and dynamic execution information into one graph.
A binary representation for dynamic memory states that generalizes better than scalar or categorical representations.
The first unified representation for control-flow and data-flow during program execution.
State-of-the-ar performance in branch prediction (by 26%) and prefetching (by 45%).
We show that NCF representations pre-trained on branch prediction are useful for transfer learning, achieving competitive performance on an algorithm classification task.
BACKGROUND
In order to generate fused representation (Figure 1. combine three fundamental components. The representation itself builds on Graph Neural Networks (GNNs). Instead of directly representing source code, our static representation uses assembly code. To drive dynamic information through the GNN, we use binary memory snapshots. We start with background on these three components.
2.1
GATED GRAPH NEURAL NETWORKS
A generic graph neural network structure G = (V, FV) consist of a set of nodes V and K sets of directed edges E = E Eg where E,g C V is the set of directed edges of type k. Each node E V is annotated with initial node embedding denoted by Z. E RDD and associated with node state vector ir E for each step of propagation == =1,,..... I....,T.
Our work builds on specific GNN variant Gated Graph Neural Networks (GGNNs) (L al. 2015).
GGNNs propagate information in the graph through message passing. At each step of propagation, "messages" to each node computed
m
Σ E:(
S(h, Ox),
(1)
2

---

Published as a conference paper at ICLR 2020
Static/Dynamic Data Collection
Learning Control/Data Flow Behavior
Program preprocessing
Static graph
Gated Graph neural networks
Source code
hh
Assemblly code
Compieer
Dynamic snapshots
Linear Mapping
Linear Mapping
Control
Data
Figure 1: Overview of the fused static/dynamic graph representation.
where m is the zero vector if there are no edges of type k directed towards v. f is a linear layer with parameters 0% this model. but can be an arbitrary function. To update the state vector of node v, all nonzero incoming messages are aggregated as:
m g(( for k such that Bu. v) E Ex}).
(2)
Here is an aggregation function, for which we use element-wise summation. Finally, the next state vector is computed using a gated recurrent unit (GRU) (Chung et al. 2014):
h +++1
GRU K.)).
The propagation is initialized with h H. and repeated T times. The state vectors hi are considered as the final node embeddings. For each task, we mark specific node v" as the "task node". We feed its final state vector h... to a linear output layer to make final predictions.
2.2 PROGRAM REPRESENTATIONS
Here we give a brief review of how compilers and processors represent source code and program state, along with tools for extracting these representations from programs and their executions.
Dynamic Execution State. The dynamic state of a program is the set of values that change as program executes. This is defined by a fixed set of registers (referenced by names like rrdi and 8rax) and memory (which is much larger and indexed by an integer memory address). Values are moved from memory to registers via load instructions and from registers to memory via store instructions. Finally, the instruction pointer specifies which instruction should be executed next.
So, what the correct subset of dynamic state to feed into a model? In principle could include all registers and memory. However, this can be difficult to work with (memory is very large) and it is expensive access arbitrary memory at test time. Instead, we restrict dynamic state snapshot that only includes CPU general purpose registers and recently used memory states. These values are cheaply obtainable hardware through buffers that hold recently used data and software through dynamic instrumentation tools like Pin (see Tools section).
Assembly Code. Assembly code is compiled from source code and is specific to a particular processor architecture (such as x86). It is a sequence of instructions, some of which operate on register values, some of which move values between registers and memory (loads and sto res), and some of which conditionally branch or jump other locations in the program. A common way of organizing assembly code is in a control flow graph (CFG). Nodes of a CFG are basic blocks, which sequences of instructions without any control flow statements. Edges point from a source basic block to target basic block when it is possible for control to jump from the source bock to the target block. For x86 direct branches, there are only two possible target blocks for given source block, which we can refer to as the true block and false block. A benefit of assembly code in our context that it is typically less stylish and tighter to program semanties. For example, programs that are syntactically different but semantically equivalent tend to correspond to similar assembly (Figure (
While we only use assembly for static code, it is also possible to link assembly code to the source code it was generated from to gain additional information about high-level constructs like data structures.
3

---

Published as a conference paper at ICLR 2020
For loop
While loop
Assembly
Semanties
Font;
i<10:
[++)
while
(i<10)
4006c1: cmpl
S0x9, (xx6))p
compare and 9 (10-1)
{
4006c5: ig
4006c9
jump out if >= 10
itl;
4006cb: addl
(x4((bb))))
add to loop back
4006c7: jmp
4006c1
4006c9:
(a) Assembly example for vs. while.
if-else
Ternary
Assembly
Semantics
(a<bb
u=b
4004da:
mov
-0xe(%rbb))
fetch
4004dd:
cmp
-0x8(%irb)), %eax
compare and b jump b am # fetch
clse
4004c0:
jge
4004ea
b;
4004e2:
Oxet%atbp) %eax
4004e5:
mov
%eeux (%x4sr)))
# i # jump out fetch b
4004c8:
jmp
4004:0
4004aa:
mov 0x8(%tb)))%aas
4004ed:
mov
400410:
50x1.%6cax
(b) Assembly example 2: if-else vs. ternary.
Two examples where syntactically different but semantically equivalent source code is compiled to the same assembly. Corresponding sections of source/assembly are colored the same.
Figure 2:
Tasks. We test learned understanding of control-flow during execution using the branch prediction task. Branch prediction traditionally uses heuristics to predict which target basic block will be entered next. The instruction pointer determines which basic block is currently being executed, and the target output a boolean specifying either the true block or false block.
Branch prediction is difficult problem with large performance implications for small relative im- provements. Modern microprocessors execute hundreds of instructions speculatively, mispredicted branch means that the processor has to discard all work completed after that branch and re-execute.
Learned understanding of data-flow during execution tested using the prefetching task. Prefetching predicts the memory address that will be accessed in the next load operation. Since data access time the largest bottleneck in server applications, solving data prefetching has significant implications for scaling computer architectures (Hashemi al. 2018). Note that there is generally interleaving of branching and memory instructions, so predicting the next memory access may depend on an unknown branch decision, and vice versa.
Tools. Compilers convert source code into assembly code. We use gec. Creating a usable snapshot of the dynamic state program is nontrivial. Given the large size of memory, we need focus on memory locations that are relevant to the execution. These are obtained by monitoring the dynamic target memory addresses of load instructions that are executed. To obtain these snapshots, we instrument instructions during execution with tool called Pin (Luk et al. 2005).
3 MODEL
Dynamic snapshots are used as features to inform the GNN of the instruction-level dynamics during execution (Section' 3.2 which we show leads to model to learn the behavior of the application (Section
model the static assembly as GNN (Section 3.1).
3.1 GRAPH STRUCTURE
BigureB an example of our graph structure translating from lines of assembly to GNN. The graph consists of three major types of nodes: instruction nodes (in white), variable nodes (in yellow), and pseudo nodes (in grey).
Instruction nodes created from instructions to the backbone of the graph. Each instruction can have variable nodes or pseudo nodes as child nodes.
Variable nodes represent variables that use dynamic values, including registers and constants.
4

---

Published as a conference paper at ICLR 2020
Instead of connecting instructions nodes directly to their child variable nodes, Pseudo nodes represent the sub-operations inside an instruction. The value associated with a pseudo node is computed in a bottom-up manner by recursively executing the sub-operations of its child nodes. For example, in instruction O in Figure 3 a pseudo node is created to represent the source operand that loads data from memory which contain child constant 4x and child register & rrbx. There number of different pseudo node types listed in the appendix.
Three major types of edges are used to connect nodes in the graph: control-flow edges, parent edges and usage edges. Control-flow edges con- nect an instruction node to all potential subse- quent instruction nodes. For non-branch instrue- tions, the control-flow edge from an instruction node points to the next sequential instruction node in the program. For branch instructions, control-flow edges are used to connect to both the next instruction and the branch target. Parent edges are used to connect child variable nodes pseudo nodes their parent instruction nodes or pseudo nodes. Usage edges provide the graph with data flow information, connecting variable nodes with their last read or write. Given this static structure, Section 332 describes how the GNN is initialized and used.
PCO: $0x48( %rdi
PC1: cmp (%rdi, %rax, 1), Sesi PC2: jne PCO
$0x48
brC
%rrbx
dest
Airdi
offset
3.2 FUSED STAMIC
SiC
Krax
base
GATED GRAPH NEURAL NETWORKS
$0x1
Node initialization. Unlike previous ap- proaches code analysis where node embed- dings are initialized with the static text of source code, we fuse the static graph with dynamic snapshots by using dynamic state to initialize nodes in the graph.
dest
%erzi
ins
Instruction mode
Control fow edge
Variable
Parent
Poeudo node Usaga adga
Each variable node and pseudo node is initial- ized with a dynamic value from the memory snapshot. These values are converted into initial node embeddings via learned embedding layer. We find that the numerical format of the dynamic values are critical to allowing the model to understand the application. We consider three types of rep- resentations for data values: categorical, scalar and binary. Our results (Section (.6) show that binary has an inherent ability to generalize more efficiently than categorical scalar representations. The intuition behind why binary generalizes so well is that the representation is inherently hierarchical, which allows for stronger generalization to previously unseen bit patterns.
Figure 3: Graph structure on assembly code.
Lastly, instruction nodes are initialized with zero vectors as embeddings. Given the initial embeddings, the GNN runs for predefined number of propagation steps to obtain the final embeddings.
Defining tasks on the graph. Tasks are defined on nodes using masking. Similar to masking in RNNs to handle variable sequence lengths, masking in GNNs handles different numbers of task nodes. A node defined with a task has a mask value of and the ones without task are masked out using 0 during both forward and backward propagation.
Branch-prediction is defined on the branch instruction node. Since each branch can either be taken or taken, this is a binary decision. The final node embeddings are fed into linear layer generate a scalar output using sigmoid activation and a cross entropy loss.
Prefetching defined on the src pseudo node that represents memory load operation. The task is to predict the 64-bit target address of the next memory load from this node. A 64-bit output is generated by feeding the final node embeddings of the task node to a different linear layer. In this
"In x86 assembly, parentheses represent addressing memory
5

---

Published as a conference paper at ICLR 2020
case, the output layer is 64-dimensional to correspond to a 64-bit address. The loss is the summation of sigmoid cross entropy loss on all 64 bits $
Scaling to large programs. For large-scale programs, it is unrealistic to utilize the static graph built on the entire assembly file (the gee benchmark has 550K instructions). As in (Li et al. 2015), to handle large graphs, only nodes which are within 100 steps to the task node will affect the prediction.
EXPERIMENTS
4.1 DATA COLLECTION
Our model consists of two parts, the static assembly and dynamic snapshots. To collect static assembly we use gee to compile source code for each binary. This binary is then disassembled using the GNU binary utilities to obtain the assembly code.
The dynamic snapshots are captured for conditional branch and memory load instructions using the dynamic instrumentation tool Pin (Luk al. 2005). We run the benchmarks with the reference input set and use SimPoint (Hamerly et al.] (005) to generate a single representative sample of 100 million instructions for each benchmark. Our tool attaches to the running process, fast forwards to the region of interest and outputs values of general registers and related memory addresses into a file every time the target conditional branch instructions or memory load instructions are executed by the instrumented application. We use SPECint 2006 to evaluate our proposal. This is a standard benchmark suite commonly used to evaluate hardware and software system performance.
4.2 EXPERIMENTAL SETUP
We train the model on each benchmark independently. The first 70% of snapshots are used for training, and the last 30% for evaluation. Hyperparameters are reported in the appendix.
4.3 METRICS
To evaluate branch prediction we follow computer architecture research and use mispredictions per thousand instructions (MPKI) (Jiménez & Lin 2001] [ee al.) 1997) as metric. Prefetching is harder problem as the predictor needs to accurately predict all bits of a target memory address. A prediction with even 1 bit off, especially in the high bits, is an error at often distant memory locations. We evaluate prefetching using complete accuracy, defined as an accurate prediction in all bits.
4.4 MODEL COMPARISONS
We compare model to three branch predictors. The first is bimodal predictor that uses 2- saturating counter for each branch instruction to keep track of its branch history (Lee et al. [997). The second is widely used, state-of-the-art perceptron branch predictor (Jiménez & Lin (001) > that uses the perceptron learning algorithm on long sequential binary taken/not-taken branch histories (Jiménez) 2016). As a more powerful baseline, we implement offline non-linear multi-layer perceptron (MLP). The MLP has two hidden layers and each layer is of the same size as the input layer. A default SGD solver used for optimization. The results are shown in Figure We find that NCF reduces MPKI by 26% and 22% compared to the perceptron and MLP respectively. Note that some of the benchmarks (libquantum, perlbench) have zero MPKI.
Three baselines are used to evaluate our prefetching model in Figure 5 The first is a stride data prefetcher (Chen Baer 1995) that good at detecting regular patterns, such as array operations. The second is a state-of-the-art address correlation (AC) prefetcher that handles irregular patterns by learning temporal address correlation (Wenisch et al. 2009). LSTM-delta a learning-batt based prefetcher that captures correlation among deltas between addresses (Hashemi et al. 2018). Due to our binary representation, NCF achieves nearly 100% coverage of all addresses, unlike the 50-80% reported for the LSTM- -rreffetett of (Hashemi et al., 2018). Figure 5 shows that NCF achieves significantly higher performance than prior work by handling both regular and irregular patterns with its binary representation. In both Figures 4 and 5 the applications are sorted from most-challenging to least-challenging. We find that NCF particularly outperforms the traditional baselines on the most challenging datasets. The traditional baselines in both branch prediction and prefetching leverage long sequential features. Our NCF does not yet use sequential features or sequential snapshots, we leave this for future work.
OOrr framework supports multitasking in that it handles control-flow and data-flow tasks simultancously. However, in our ablation studies, we did not see significant evidence that these tasks currently help each other.

---

Published as a conference paper at ICLR 2020
25
Bimod
MPKI (lower is better)
Perceptron
MLP
15
NCF
5
RCc
mcf
sjeng
astar
bxip2
gobmk
omnetpp
xalan
h264
hmmer
libquant
perlbench
average
Figure Evaluation of the branch-prediction task (lower better).
Stride
Accuracy (higher is better)
0.9
AC
LSTM-delta
0.8
NCF
0.7
0.6
0.5
0.4
0.3
0.2
0.1
mcf
gcc
hmmer
xalan
astar
omnetip
bzlp2
h264ref
gobmk
sjeng
libquantum
peribench
average
Figure 5: Evaluation of the prefetching task (higher is better).
The effectiveness of the GNN depends on the input graph, and we perform ablation studies in the appendix (Section B.I.
4.5 ALGORITHM CLASSIHICATION
To test the model has learned about the behavior of the application, we test the NCF representation an algorithm classification dataset (Lili Mou 2016). We randomly select subset of 15 problems from this datase and generate inputs for each program. 50 programs are randomly selected from each class. These are split into 30 for training, 10 for validation (tuning the linear SVM described below) and 10 for testing.
We generate the graph for each program post-compilation and obtain memory snapshots via instrumentation tool. The representation is pre-trained on branch prediction and the resultant embed- dings are averaged to serve as the final embedding of the program. linear SVM is trained using the pre-trained embeddings to output predicted class.
2018 achieves 95.3% on This yields 96.0% test accuracy, where the saute-of-the-see (Ben-Nun al.
the same subset. In contrast to Ben-BBnnin Nun et al. (2018), 8), which pre-trains an LSTM on over 50M lines of L.LVM IR, our embeddings are trained on 203k lines of assembly from the algorithm classification dataset itself. This shows that branch prediction can be highly predictive of high-level program attributes, suggesting that it may be fruitful to use dynamic information to solve other static tasks.
4.6 GENERALIZATION TEST ON REPRESENTATIONS
Lastly, we test the effectiveness of binary representations of memory state. There are three major options for representing dynamic state: categorical, real-valued scalar. and binary. State-of-the-
We use subset because the programs had to be modified (by adding appropriate headers, fixing bugs) to compile and in order to retrieve the assembly code and dynamic states.
7

---

Published as a conference paper at ICLR 2020
art data prefetchers tend to categorical representations. Recent advances in representing and
manipulating numbers for neural arithmetic logic units use scalar representations (Trsk et al. 2018).
We evaluate the generalization ability of these representations using a simple loop. We replace the constant 10 total iterations of the loop in Figure [(a] with a variable k. The control-flow of the loop decides
Yes
Completely correct
stay in jump out of the loop by comparing variable and k. The branch will be not taken for the first k times but will be taken at the kth time, Since tradi- tional state-of-the-art branch predictors depend
Categorical
Scalar Binary
on memorizing past branch history, they will always mispredict the final branch (as has al- ways been taken). Our proposal is able to make the correct prediction at the kth time.
No
10
20
30
60
70
80
k values
The challenge for our model is that the value k can change during program execution, and the model needs to generalize unseen values. We use this example to test the three representations and create testing 40 with step size of 3 (I, 4, 7, 37). We feed all three representations to MLP predictors that have one hidden layer of the same size of each input representation (160 for categorical, 2 for scalar and 14 for binary). The results are shown in Figure b
Figure 6: Generalization ability of representations. "Yes" means predictions on all branches of the loop with & value are correct.
using & values from 1 80. The training set only contains h values from to
The categorical representation can only correctly predict training samples, missing every two out of three k values, where scalar and binary representations are both able to generalize across a continuous range, filling the "holes" between training samples. The binary representation generalizes to larger range than scalar representation, as long as the bits have been seen and toggled in the training set. Since binary is inherently hierarchical (the range increases exponentially with the number of bits), this advantage is greater in real world 64-bit machine.
5 RELATED WORK
5.1 LEARNING FROM SOURCE CODE & EXECUTION BEHAVIOR
There is significant body of work on learning for code, and we refer the reader to Allamanii et al. (2018) for survey. We focus the most relevant methods here. Li al. (2015) use GNNs to represent the state of heap memory for a program verification application, Allamanis et al. (2017)
learn to represent source code with GNNs.
Similar to us, Ben-Nun et al. (0018): Mendis et al. (2018) learn representations of code from low-level syntax, the LLVM intermediate representation (IR) assmebly, but do not dynamic information. We use assembly code instead of IR to maintain a 1:1 mapping between dynamic state and the static backbone of the graph (since instructions are atomic when executed). Prior work that builds graphs purely based on static source code disregard the instruction-level dynamics that are created during program execution, as single static piece of code can execute in different ways depending on the provided inputs.
Wang et al. (2017) embed the sequences of values that variables take on during the execution of program as dynamic program embedding. The code is not otherwise used. The states are relatively simple (variables can take on relatively few possible values) in contrast our dynamic states that are "from the wild." Cummiins et al.) (2017) embeds code and optionally allows flat vector of auxiliary features that can depend on dynamic information. Abstract program execution can also used as a basis for learning program representations (DeFreezz et al. 2018 Henkel et al. 2018). However, neither uses concrete program state.
5.2
USING PROGRAM STATE TO GUIDE PROGRAM SYNTHESIS
There are several works that learn from program state to aid program synthesis (Balog et al. 2016
Parisotto al. 2016 Devlin et al. 2017: Zohar & Wolf 2018 Chen et al., 2019 Vijayakumar ct al. 2018] Menon et al. 2013). particular, Balog et al. (2016) use neural networks to learn a mapping from list-of-integer-valus input-output examples to the set of primitives needed. All of
8

---

Published as a conference paper at ICLR 2020
these operate programs in relatively simple Domain Specific Languages and are learning mappings from program state to code, rather than learning joint embeddings code and program state.
5.3 DYNAMIC PREDICTION TASKS
Branch prediction and prefetching are heavily studied in the computer architecture domain. High- performance modern microprocessors commonly include perceptron (Jiménez & Lin 2001) or table-based branch predictors that memorize commonly taken paths through code (Seznec 2011).
While there has been significant amount of work around correlation prefetching in academia (Wenisch et al. | CChanne & Reeves 1995; Roth et al. 1998), modern processors only commonly implement simple stream prefetchers (Chen & Baer, 1995 Jouppi 1990, Gindele [977). Recent work has related prefetching to natural language models and shown that LSTMs achieve high accuracy (Hashemi al. 2018). However, their categorical representation covers only a limited portion of the access patterns while the binary representation described here is more general.
CONCLUSION
We develop novel graph neural network that uses both static and dynamic features learn rich representation for code. Since the representation is based on a relational network, it easy to envision extensions that include high-level source code into the model to add new prediction tasks. Instead of focusing on rrrlwaee-realizaable systems with real-time performance, our primary focus in this paper develop representations that explore the limits of predictive accuracy for these problems with extremely powerful models, so that the improvements can be be eventually be distilled. This is common machine learning research, where typically the limits of performance for a given approach are reached, and then distilled into a performant system, e.g. (Van Den Oord et al. 2016 Oord et al.) 2017). However, benefits can still be derived by using the model to affect program behavior through compilation hints (Chilimbi & Hirzel 2002: Jagannathan & Wright 1996. Wolf et al. 1996). making this exploration immediately practical. We argue that fusing both static and dynamic features into one representation is an exciting direction to enable further progress in neural program understanding.
REFERENCES
Miltiadis Allamanis, Mare Brockschmidt, and Mahmoud Khademi. Learning to represent programs with graphs. arXiv preprint arXiv: 2017.
Miltiadis Allamanis, Earl Barr, Premkumar Devanbu, and Charles Sutton. A survey of machine learning for big code and naturalness. ACM Computing Surveys (CSUR), 5 51()):81, 2018.
Matej Balog, Alexander L Gaunt, Mare Brockschmidt, Sebastian Nowozin, and Daniel Tarlow. Deepcoder: Learning to write programs. arXiv preprint arXiv:/61/ 1.1199991 2016.
Tal Ben-Nun, Alice Shoshana Jakobovits, and Torsten Hoefler. Neural code comprehension: A learnable representation of code semantics. arXiv preprint arXiv:1806. 07336, 2018.
M. J. Charney and A. P. Reeves, Generalized correlation-basel hardware prefetching. Technical Report EE-CEG-99 9,-1- Comell Univ., 1995.
Tien-Fu Chen and Lean-Loup Baer. Effective hardware-basee data prefetching for high-performance processors. IEEE transactions on computers, 1995.
Xinyun Chen, Chang Liu, and Dawn Song. Execution-guided neural program synthesis. 2019.
Trishul M Chilimbi and Martin Hirzel. Dynamic hot data stream prefetching for general-purposs. programs. In ACM SIGPLAN Notices, 2002.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014.
Chris Cummins, Pavlos Petoumenos, Zheng Wang, and Hugh Leather. End-to-end deep learning of optimization heuristics. In PACT, 2017.
Daniel DeFreez, Aditya V Thakur, and Cindy Lubio-Gonzálee Path-based function embedding and its application to specification mining. arXiv preprint a0X 2018.

---

Published as a conference paper at ICLR 2020
Jacob Devlin, Jonathan Uesato, Surya Bhupatiraju, Rishabh Singh, Abdel-rahman Mohamed, and Pushmeet Kohli. Robustfill: Neural program learning under noisy i/o. arXiv preprint arXiv:1003.07469, 2017.
James Dundas. Announcement of results and awards. Championship Branch Prediction (CBP-5) in Conjunction with ISCA-43, 2016.
J. D. Gindele. Buffer block prefetching method. IBM Technical Disclosure Bulletin, 20(2):696 July 1977.
Greg Hamerly, Erez Perelman, Jeremy Lau, and Brad Calder. Simpoint 3.0: Faster and more flexible program phase analysis. Journal of Instruction Level Parallelism, 7(4), 2005.
Milad Hashemi, Kevin Swersky, Jamie A Smith, Grant Ayers, Heiner Litz, Jichuan Chang, Christos Kozyrakis. and Parthasarathy Ranganathan. Learning memory access patterns. arXiv preprint arXiv:1803.02329, 2018.
Jordan Henkel, Shuvendu Lahiri, Ben Liblit, and Thomas Reps. Code vectors: Understanding programs through embedded abstracted symbolic traces. arXiv preprini arXiv: 1803:06800666 2018.
John L Hennessy and David A Patterson. A new golden age for computer architecture. Communica- tions of the ACM, 62(2), 2019.
Geoffrey Hinton, Li Deng, Dong Yu, George Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Brian Kingsbury, and Tara Sainath. Deep neural networks for acoustic modeling in speech recognition. IEEE Signal processing magazine, 2012.
Suresh Jagannathan and Andrew Wright. Flow-directed inlining. In ACM SIGPLAN Notices, 1996.
Daniel A Jiménez. Multiperspective perceptron predictor. Championship Branch Prediction (CBP-5) in Conjunction with ISCA-43, 2016.
Daniel A Jiménez and Calvin Lin. Dynamic branch prediction with perceptrons. In The International Symposium on Perfomaaeo Computer Architecture, 2001.
Norman Jouppi.
Improving direct-mapped cache performance by the addition of small fully- associative cache and prefetch buffers. In The International Symposium on Computer Architecture, 1990.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. Advances in Neural Information Processing Systems, 2012.
Chih-Chieh Lee, I-Cheng K Chen, and Trevor N Mudge. The bi-mode branch predictor. In The International Symposium on Microarchitecture, 1997.
Yujia Li, Daniel Tarlow, Mare Brockschmidr, and Richard Zemel. Gated graph sequence neural networks. arXiv preprint arX1v:1511.05433, 2015.
Lu Zhang Tao Wang Zhi Jin Lili Mou, Ge Li. Convolutional neural networks over tree structures for programming language processing, AAAI Conference on Artificial Intelligence, 2016.
Chi-Keung Luk, Robert Cohn, Robert Muth, Harish Patil, Artur Klauser, Geoff Lowney, Steven Wallace, Vijay Janapa Reddi, and Kim Hazelwood. Pin: building customized program analysis tools with dynamic instrumentation. In ACM SIGPLAN notices, 2005.
Charith Mendis, Alex Renda, Saman Amarasinghe, and Michael Carbin. Ithemal: Accurate, portable and fast basic block throughput estimation using deep neural networks. arXiv preprint arXiv:1808.077 07 2018.
Aditya Menon, Omer Tamuz, Sumit Gulwani, Butler Lampson, and Adam Kalai. A machine learning framework for programming by example. In International Conference Machine Learning, 20133
10

---

Published as a conference paper at ICLR 2020
Aaron van den Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George van den Driessche, Edward Lockhart, Luis C Cobo, Florian Stimberg, Norman Casagrande, Dominik Grewe, Seb Noury. Sander Dieleman, Erich Elsen, Nal Kalchbrenner, Heiga Zen, Alex Graves, Helen King, Tom Walters, Dan Belov, and Demis Hassabis. Parallel wavenet: Fast high-fidelity speech synthesis. arXiv preprint arXiv: /7/ 10.433, 2017.
Emilio Parisotto, Abdel-rahman Mohamed, Rishabh Singh, Lihong Li, Dengyong Zhou, and Pushmeet Kohli. Neuro-symbolic program synthesis. arXiv preprint 1rXiv:1611.01855 2016.
Parthasarathy Ranganathan. More moore: Thinking outside the server box. ISCA-44 Keynote, 2017.
Amir Roth, Andreas Moshovos, and Gurindar S. Sohi. Dependence based prefetching for linked data structures. In The International Conference on Architectural Support for Programming Languages and Operating Systems, 1998.
André Seznec. new case for the tage branch predictor. In The International Symposium on Microarchiteectrre, 2011.
James E Smith. A study of branch prediction strategies. In The International Symposium Computer Architecture, 1981.
SPEC CPU2006 http://www..ppcbeecchorrrg
The Standard Performance
Evaluation
Corporation,
2006.
llya Sutskever, Oriol Vinyals, and Quoc Le. Sequence sequence learning with neural networks. In Advances in Neural Information Processing Systems, 2014.
Andrew Trask, Felix Hill, Scott E Reed, Jack Rae, Chris Dyer, and Phil Blunsom. Neural arithmetic logic units. Advances in Neural Information Processing Systems, 2018.
Aäron Van Den Oord, Sander Dicleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew W Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. 2016.
Ashwin J Vijayakumar, Abhishek Mohta, Oleksandr Polozov, Dhruv Batra, Prateck Jain, and Sumit Gulwani. Neural-guided deductive search for real-time program synthesis from examples. arXiv preprint arXiv:1804.00 1866, 2018.
Ke Wang, Rishabh Singh, and Zhendong Su. Dynamic neural program embedding for program repair. arXiv preprint arXiv: 171 1,07 2017.
Thomas Wenisch, Michael Ferdman, Anastasia Ailamaki, Babak Falsafi, and Andreas Moshovos. Practical off-chip meta-data for temporal memory streaming. In The International Symposium High Performance Computer Architecture, 2009.
Michael E Wolf, Dror E Maydan, and Ding-Kai Chen. Combining loop transformations considering caches and scheduling. In The International Symposium on Microarchitecture, 1996.
Amit Zohar and Lior Wolf. Automatic program synthesis of long programs with a learned garbage collector. In Advances in Neural Information Processing Systems. 2018.

---

Published as a conference paper at ICLR 2020
A HYPERPARAMETERS
The hyperparameters for all models are given in Table
Table 1: Hyperparameters for models.
input feature size
64
hidden size
64
Fused GGNN
propagation steps
optimizer
adam
learning rate
0.01
hidden size
2nnpptt size
MLPs for generalization test
optimizer
adam
1.2 regularization
0.0001
loss
square hinge
SVM for algorithm classification
L2 regularization
0.01
bits
2
baseline: Bimodal
Resources
Unlimited
history length
64
baseline: Perceptron
L2 regularization
0.0001
Unlimited resources to store all strides (delta between addresses) for each load, predicting the most frequent stride Unlimited resources
baseline: stride
every pairwise correlation, predicting the most frequent pair
baseline: Address Correlation
12

---

Published as a conference paper at ICLR 2020
B NODE SUB-TYPES
We describe the node sub-types in Table Pscudo-nodes implement operations that are commonly known as the addressing modes of the Instruction Set Architecture. Note that node sub-types used to derive initial node embeddings and for interpretability. They do not factor into the computation of the graph neural network.
Table 2: Descriptions about sub node types.
Major node type
Sub-type
Description
source operand that does not involve memory load operation, obtained directly from register(s) and/or constant(s)
non-mem-
source operand that involves memory load operation, obtained from loading data from a memory location target operand that does not involve memory write operation, writing directly to register
Pseudo nodes
mem-sre
non-mem-
fgt
source operand that involves a memory write operation, writing data memory location
mem-tgt
base
base that is obtained directly from variable node an indirect base that obtained from certain operations on the child variable nodes, like multiplying register by a constant an offset value that is to be added to a base
ind-base
offset
Variable nodes
reg
register, value is dynamically changed during execution constant, value is specified the assembly
const
B.1 ABLATION STUDY
The effectiveness the GNN depends on the input graph. As pseudo nodes are large component of the static graph, we run additional experiments to understand their importance. In particular, we try to only use the pseudo nodes src and tgt, which are directly connected to instruction nodes. Our data shows that removing pscudo nodes other than and igt and connecting variable nodes directly to src and rgt has little impact on branch prediction (an MPKI increase of 0.26), but has a large impact on the data-flow accuracy (accuracy goes down by 12.1% 1%).
F shows the sensitivity of task performance to the number propagation steps during training for the GNN omnelpp. We find that prefetching is sensitive propagation steps than branch prediction, and requires 5-8 steps for peak accuracy. Due to the control flow of programs, we find that 5-8 steps propagates information for 50-60 instruction nodes across the graph's backbone for omnetpp (up to 6000 nodes for perlbench).
350
0.6
MPKI for Branch Prediction
300
0.5
Accuracy for Prefetching
(lower bettter)
250
(higher is better)
0.4
200
branch prediction
0.3
150
prefetching
0.2
100
50
0.1
0
4
5
6
7
Number of Propagation Steps
Figure 7: GNN performance vs. number of propagation steps.
13


Reviews:
[{"review_id": "HJlUX4sTFH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes using Graph Neural Networks to learn representations of source code and its execution. They test their method on the SPEC CPU benchmark suite and show substantial improvement over methods that do not use execution. \n\nThe paper's main question is to answer how to learn code representations. The main novelty introduced in their approach is to build a graph representation not of high level code but of assembly code. They also develop a way to use what they call a \"snapshot mechanism\" that feeds limited memory states into the graph. The downstream consequences of their methods are improved methods for example for branch prediction. Interestingly NCF can be also used to represent programs for use in downstream tasks. This is demonstrated via transfer learning in an algorithm classification problem. The paper is well written and the background / related work makes it easy for the reader to understand the problem's relevance within the related literature. \n\nThe results look well justified and empirically verified. ", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "N/A", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "H1e_TZJS9H", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Using Deep Learning and especially, GNNs seems to be a popular area of research. I am no \nexpert at optimizing code performance, so please take my review with a grain of salt. The algorithmic contributions of the paper are as following: \n\n(a) GNN that combines static code and dynamic execution trace. \n(b) Binary encoding of features leads to better performance in comparison to categorical and scalar representations. \n\nThe results show that the proposed method outperforms existing methods on standard benchmarks in the program execution community. \n\nFrom a machine learning stand point, the contributions are straightforward and the results make sense. I have the following questions: \n\n(I) Authors argue that binary representations are better because of their hierarchical nature. They mention that they can generalize even if not all combinations of bits are seen, but a subset is seen in a manner that every bit has been flipped a couple of times. I don’t agree with this reasoning, as seeing the individual bits flip has no guarantee that a NN would generalize to a new combination of bits unless the distance in the binary code makes sense. Is there some special way in which the binary code is constructed?\n \n(ii) Transfer learning experiments: Its unclear to me if the comparison presented in the paper is a fair one. Comparison is made against Ben-Nun et al. pre-training on LLVM IR. I am not sure how different is LLVM IR dataset from the algorithm classification dataset. If the dataset is very different, then obviously a lot of pre-training will only result in modest performance gain. What happens with Ben-Nun method is pre-trained on the same dataset as the proposed method? Also, what is the difference in performance between the cases when the proposed method is applied to algorithm classification with and without pre-training? \n\nOverall, the paper is a application of GNN to optimizing code execution. The technical innovations are domain-specific and do not inform the general machine learning community. Given lack of expertise in the area of program execution, I cannot judge the significance of the performance improvements reported in the paper. \n\nGiven my current concerns, I cannot recommend acceptance. I might change my ratings based on the review discussions and the author’s responses to the above questions. \n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "BkeiWrw1qH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a novel improvement in methodology for learning code execution (at the level of branch-predictions and prefetching).  They combine static program description with dynamic program state into one graph neural network, for the first time, to achieve significant performance gains on standard benchmarks.\n\nI would vote to accept this paper.  They appear to have developed a new model structure and interface to the program information (i.e. inputs to the model), and the design decisions appear thoughtful, sensible, and well-justified (e.g. use of assembly code).  The presentation is mostly clear, with a good balance of background material, method description, and experiment results.  \n\nTaken at face value, the results are impressive, although I am not familiar enough with this field to assess the fairness of comparison against the baselines.  For example, it's a little unclear what the difference is vs previous baselines just from switching to source-code-as-input to assembly-code-as-input?\n\nThe study on memory representations (categorical vs scalar vs binary) is a helpful component which adds its own value, and the context for popularity of the alternatives is described.\n\nFew details as to implementation are discussed, although the code is included in the submission, and after a quick glance appears substantial.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a method to learn representations of programs via code and execution.

The paper presents an interesting method, and results on branch prediction and address pre-fetching are conclusive. The only main critiques associated with this paper seemed to be (1) potential lack of interest to the ICLR community, and (2) lack of comparison to other methods that similarly improve performance using other varieties of information. I am satisfied by the authors' responses to these concerns, and believe the paper warrants acceptance.
Decision: Accept (Poster)