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: The paper proposes a new program representation to find semantic code embeddings (ProGraML). The authors use these embeddings to perform a number of traditional dataflow analysis tasks that you can find in modern day compilers. When trained under a supervised learning setting, Deep Data Flow (DDF) analysis learns to perform these tasks more accurately compared to other deep program representations.