Deep Data Flow Analysis

1 INTRODUCTION . Compiler implementation is a complex and expensive activity ( Cooper & Torczon , 2012 ) . For this reason , there has been significant interest in using machine learning to automate various compiler tasks ( Allamanis et al. , 2018 ) . Most works have restricted their attention to selecting compiler heuristics or making optimization decisions ( Ashouri et al. , 2018 ; Wang & O ’ Boyle , 2018 ) . Whether learned or engineered by human experts , these decisions naturally require reasoning about the program and its behavior . Human experts most often rely upon data flow analyses ( Kildall , 1973 ; Kam & Ullman , 1976 ) . These are algorithms on abstract interpretations of the program , propagating information of interest through the program ’ s control-flow graph until a fixed point is reached ( Kam & Ullman , 1977 ) . Two examples out of many data flow analyses are : liveness – determining when resources become dead ( unused ) and may be reclaimed ; and available expressions – discovering which expressions have been computed on all paths to points in the program . Prior machine learning works , on the other hand , have typically represented the entirety of the program ’ s behavior as a fixed-length , statically computed feature vector ( Ashouri et al. , 2018 ) . Typical feature values might be the number of instructions in a loop or the dependency depth . The weakness of these techniques is shown by the fact that they are trivially confused by the addition of dead code , which changes their feature vectors without changing the program ’ s behavior or its response to optimizations . Such learning algorithms are unable to learn their own abstract interpretations of the program and so can not avoid these pitfalls or more subtle versions thereof ( Barchi et al. , 2019 ) . Recently , there have been attempts to develop representations that allow finer-grain program reasoning . Many , however , are limited both by how inputs are represented as well as how inputs are processed . Representations based on source code and its direct artifacts ( e.g. , AST ) ( Alon et al. , 2018a ; Yin et al. , 2018 ; Haj-Ali et al. , 2020 ) put unnecessary emphasis on naming and stylistic choices that may not correlate with the functionality of the code ( e.g. , Fig . 2a ) . Approaches based on intermediate representations ( IR ) ( Ben-Nun et al. , 2018 ; Mirhoseini et al. , 2017 ; Brauckmann et al. , 2020 ) remove such noise but fail to capture information about the program that is important for analysis ( e.g. , Fig . 2b variables , Fig . 2c commutativity ) . In both cases , models are expected to reason about the flow of information in programs using representations that do not directly encode this information . Clearly , a program representation is needed that enables machine learning algorithms to reason about the execution of a program by developing its own data flow analyses . Since current approaches are ill-suited to program-wide data flow analysis , we propose overcoming their limitations by making the program ’ s control , data , and call dependencies a central part of the program ’ s representation and a primary consideration when processing it . We achieve this by seeing the program as a graph in which individual statements are connected to other statements through relational dependencies . Each statement in the program is understood only in the context of the statements interacting with it . Through relational reasoning ( Battaglia et al. , 2018 ) , a latent representation of each statement is learned that is a function of not just the statement itself , but also of the ( latent ) representations of its graph neighborhood . Notably , this formulation has a striking similarity to the IRs used by compilers , and the iterative propagation of information resembles the transfer functions and meet operators in traditional data flow analyses ( Kildall , 1973 ) . Recently proposed techniques for learning over graphs have shown promise in a number of domains ( Schlichtkrull et al. , 2018 ; Ziwei et al. , 2020 ) . With a suitable representation and graph-based model , we extend these approaches to the domain of compiler analysis , enabling downstream tasks built on top of such graph models to natively incorporate reasoning about data flow into their decision making . This improves performance on downstream tasks without requiring additional features . We make the following contributions : • We propose a portable , language-independent graph representation of programs derived from compiler IRs . PROGRAML is the first representation to capture whole-program control- , data- , and call relations between instructions and operands as well as their order and data types . PROGRAML is a compiler-agnostic design for use at all points in the optimization pipeline ; we provide implementations for LLVM and XLA IRs . • We introduce a benchmark dataset that poses a suite of established compiler analysis tasks as supervised machine learning problems . DEEPDATAFLOW comprises five tasks that require , in combination , the ability to model : control- and data-flow , function boundaries , instruction types , and the type and order of operands over complex programs . DEEPDATAFLOW is constructed from 461k real-world program IRs covering a diverse range of domains and source languages , totaling 8.5 billion data flow analysis classification labels . • We adapt Gated-Graph Neural Networks ( GGNN ) to the PROGRAML representation . We show that , within a bounded problem size , our approach achieves ≥ 0.939 F1 score on all analysis tasks , a significant improvement over state-of-the-art representations . In evaluating the limits of this approach we propose directions to better learn over programs . 2 RELATED WORK . Data flow analysis is a long established area of work firmly embedded in modern compilers . Despite its central role , there has been limited work in learning such analysis . Bielik et al . ( 2017 ) use ASTs and code synthesis to learn rule-sets for static analyses , some of which are dataflow-related . Our approach does not require a program generator or a hand-crafted DSL for rules . Shi et al . ( 2020 ) and Wang & Su ( 2020 ) use dynamic information ( e.g. , register snapshots and traces ) from instrumented binaries to embed an assembler graph representation . We propose a static approach that does not need runtime features . Si et al . ( 2018 ) use a graph embedding of an SSA form to generate invariants . The lack of phi nodes and function call/return edges means that the representation is not suitable for interprocedural analysis as it stands . Kanade et al . ( 2020 ) explore a large-scale , context-dependent vector embedding . This is done at a token level , however , and is unsuited for dataflow analysis . Prior work on learning over programs employed methods from Natural Language Processing that represented programs as a sequence of lexical tokens ( Allamanis , 2016 ; Cummins et al. , 2017a ) . However , source-level representations are not suited for analyzing partially optimized compiler IRs as the input source can not be recovered . In program analysis it is critical to capture the structured nature of programs ( Raychev et al. , 2015 ; Allamanis et al. , 2017 ; Alon et al. , 2018b ) . Thus , syntactic ( tree-based ) as well as semantic ( graph-based ) representations have been proposed ( Allamanis et al. , 2018 ; Brauckmann et al. , 2020 ) . Dam et al . ( 2018 ) annotate nodes in Abstract Syntax Trees ( ASTs ) with type information and employ Tree-Based LSTMs ( Tai et al. , 2015 ) for program defect prediction . Both Raychev et al . ( 2015 ) and Alon et al . ( 2018a ; b ) use path-based abstractions of the AST as program representations , while Allamanis et al . ( 2017 ) augment ASTs with a hand-crafted set of additional typed edges and use GGNNs ( Li et al. , 2015 ) to learn downstream tasks related to variable naming . Another line of research considers modelling binary similarity via control-flow graphs ( CFGs ) with an adaptation of GNNs called Graph Matching Networks ( Li et al. , 2019 ) . The history of IR-based graph representations for optimization goes back to Ferrante et al . ( 1987 ) , who remove superfluous control-flow edges to ease optimization with a compact graph representation . A more contemporary precursor to our approach is the ConteXtual Flow Graph ( XFG ) ( BenNun et al. , 2018 ) , which combines control-flow with data-flow relations in order to learn unsupervised embeddings of LLVM-IR statements . XFGs omit information that is critical to analysis including the notion of argument order , vertices for both variables and constants , and all control-flow edges . PROGRAML , in combining call-graphs ( CG ) , control-flow graphs , and data-flow graphs ( DFG ) , offers an IR-level program representation that is designed to be useful for a variety of purposes from specific program analyses to downstream optimization tasks . Control and Data Flow Graphs ( CDFG ) ( Brauckmann et al. , 2020 ) use graph vertices for statements and have bi-directional edges for control and data dependencies . The CDFG uses only the instruction opcode to represent a statement , omitting operands , variables , data types , and constants . This prohibits the reasoning about variables and expressions that are required for many data flow analyses , including 3 out of the 5 benchmark tasks that we establish below . Mendis et al . ( 2019 ) represent LLVM-IR using a graph that is specialized to a vectorization task . They use unique edge types to differentiate the first five operand positions and augment the graph structure with vectorization opportunities that they compute a priori . Our approach is not specialized to a task , enabling such opportunities learned ( e.g. , subexpression detection ) , and uses a embedding weighting to differentiate edge positions without having to learn separate edge transfer weights for each . Finally , an alternate approach is taken by IR2Vec ( Keerthy S et al. , 2019 ) , an LLVM-IR-specific representation that elegantly models part-ofstatements as relations . However , in order to compute the values of the embeddings , IR2Vec requires access to the type of data flow analyses that our approach learns from data alone . 3 A GRAPHICAL REPRESENTATION FOR DEEP PROGRAM ANALYSIS . This section presents PROGRAML , a novel IR-based program representation that closely matches the data structures used traditionally in inter-procedural data flow analysis and can be processed natively by deep learning models . We represent programs as directed multigraphs where instructions , variables , and constants are vertices , and relations between vertices are edges . Edges are typed to differentiate control- , data- , and call-flow . Additionally , we augment edges with a local position attribute to encode the order of operands to instructions , and to differentiate between divergent branches in control-flow . We construct a PROGRAML graph G = ( V , E ) by traversing a compiler IR . An initially empty graph G = ∅ is populated in three stages : control-flow , data-flow , and call-flow , shown in Figure 3 . In practice the three stages of graph construction can be combined in a single O ( |V |+ |E| ) pass . ( I ) Control Flow We construct the full-flow graph of an IR by inserting a vertex for each instruction and connecting control-flow edges ( Fig . 3a , 3b ) . Control edges are augmented with a numeric position using an ascending sequence based on their order in the list of an instruction ’ s successors . ( II ) Data Flow We introduce constant values and variables as graph vertices ( Fig . 3c ) . Dataflow edges are inserted to capture the relation from constants and variables to the instructions that use them as operands , and from instructions to produced variables . As each unique variable and constant is a vertex , variables can be distinguished by their scope , and unlike the source-level representations of prior works , variables in different scopes map to distinct vertices and can thus be discerned . Data edges have a position attribute that encodes the order of operands for instructions . The latent representation of a statement ( e.g. , % 1 = add i32 % 0 , 1 ) is thus a function of the vertex representing the instruction and the vertices of any operand variables or constants , modulated by their order in the list of operands . ( III ) Call Flow Call edges capture the relation between an instruction that calls a function and the entry instruction of the called function ( Fig . 3d ) . Return call edges are added from each of the terminal instructions of a function to the calling statement . Control edges do not span functions , such that an IR with functions F produces |F | disconnected subgraphs ( the same is not true for data edges which may cross function boundaries , e.g. , in the case of a global constant which is used across many parts of a program ) . For IRs that support external linkage , an additional vertex is created representing an external call site and connected to all externally visible functions . If a call site references a function not defined in the current IR , a dummy function is created consisting of a single instruction vertex and connected through call edges to all call sites in the current IR . A unique dummy function is created for each externally defined function .

Summary: A methodology for learning representations of programs using graph neural networks applied to graphs extracted from a compiler intermediate representation is presented. The graph representation captures both control flow as well as data dependencies and also represents calls to/returns from functions. A large new dataset to evaluate if models can learn to perform dataflow analyses is introduced, and experimental results show that the described method can be trained to perform these with high accuracy (at least for small programs). Further experiments in the appendix illustrate that the proposed graph structure is also useful for more standard tasks that cannot be solved well using existing algorithms.