Neural Code Comprehension: A Learnable Representation of Code Semantics

Tal Ben-Nun; Alice Shoshana Jakobovits; Torsten Hoefler

Source-linked AI summary

Neural Code Comprehension: A Learnable Representation of Code Semantics

Tal Ben-Nun, Alice Shoshana Jakobovits, Torsten Hoefler

arXiv:1806.07336v3Jun 19, 2018cs.LGcs.NEcs.PLstat.ML

TL;DR

Existing code-embedding methods do not robustly represent semantics across language variation, branching, function calls, and nonsequential relationships. The paper addresses this with language-independent LLVM IR, contextual flow combining data and control dependence, and inst2vec embeddings, then reports strong performance across three high-level tasks. Its conclusion is that statement semantics can be recovered from context using appropriate statement granularity and mixed data/control flow.

Problem
Existing methods are limited by source-language dependence and sequential processing that does not adequately represent code with functions and loops.
Method
The pipeline converts code to LLVM IR, builds contextual flow graphs from data- and control-flow, and learns inst2vec statement embeddings for RNN-based tasks.
Results
The representation outperforms manually extracted features and matches or exceeds specialized deep neural network solutions across three high-level classification and prediction tasks.
Takeaways & Limitations
The results support recovering statement semantics from context when statement granularity and grouping combine filtered LLVM IR with data- and control-flow.
Takeaways & Limitations
The semantic-similarity formulation assumes statements are similar when they modify system state or consume resources similarly, including across variable-type changes.

Abstract

from arXiv · show

With the recent success of embeddings in natural language processing, research has been conducted into applying similar methods to code analysis. Most works attempt to process the code directly or use a syntactic tree representation, treating it like sentences written in a natural language. However, none of the existing methods are sufficient to comprehend program semantics robustly, due to structural features such as function calls, branching, and interchangeable order of statements. In this paper, we propose a novel processing technique to learn code semantics, and apply it to a variety of program analysis tasks. In particular, we stipulate that a robust distributional hypothesis of code applies to both human- and machine-generated programs. Following this hypothesis, we define an embedding space, inst2vec, based on an Intermediate Representation (IR) of the code that is independent of the source programming language. We provide a novel definition of contextual flow for this IR, leveraging both the underlying data- and control-flow of the program. We then analyze the embeddings qualitatively using analogies and clustering, and evaluate the learned representation on three different high-level tasks. We show that even without fine-tuning, a single RNN architecture and fixed inst2vec embeddings outperform specialized approaches for performance prediction (compute device mapping, optimal thread coarsening); and algorithm classification from raw code (104 classes), where we set a new state-of-the-art.

1 Introduction

The paper develops a language-independent pipeline for learning robust code semantics from LLVM IR and contextual flow, targeting code-comprehension tasks across varied programs. It evaluates this representation on clustering, analogies, and three high-level prediction or classification tasks.

Motivation: Code analysis is increasingly necessary because accessible code is growing rapidly and must be sifted, categorized, and understood across applications.The paper cites approximately 1 billion GitHub commits in 2017 across 337 programming languages.
Motivation: Existing approaches struggle with language specificity and sequential processing, which limits coverage of code containing functions and loops.The authors note that many real applications are not function- and loop-free.
Approach: The pipeline converts code from various source languages into LLVM IR, constructs XFGs from data- and control-flow, and trains inst2vec statement embeddings.XFGs support loops and function calls, and inst2vec embeddings are supplied to RNNs for high-level tasks.
Evaluation: The evaluation covers inst2vec clustering and analogies, algorithm classification, heterogeneous device mapping, and optimal thread coarsening prediction.The datasets include CPU and GPU code in C, C++, OpenCL, and FORTRAN.
Results: A single LSTM with fixed pre-trained embeddings matches or surpasses the best-performing specialized approaches across the evaluated tasks.The contribution claim includes specialized deep neural network architectures as comparators.

2 Related Work

Prior work embeds code tokens, statements, or lower-level representations and uses lexical or structural contexts for downstream tasks. This paper distinguishes its approach through compiler-IR embeddings, contextual flow, and direct evaluation of embedding quality.

Code Representation: Previous research embeds high-level languages such as Java, C, and OpenCL as tokens or statements, alongside lower-level object-code representations.The paper positions its compiler-IR embedding work against these representation levels.
Code Representation: Earlier contextual representations use lexical locality or structural graphs including data-flow graphs, control-flow graphs, ASTs, and augmented ASTs.The paper contrasts these with a graph combining data and control dependencies.
Code Representation: The paper introduces contextual flow for statement embeddings, combining data and control dependencies while allowing ambiguity that preserves contextual relationships.Unlike compiler-oriented graphs, these graphs are not designed for optimization or machine-code translation.
Automated Tasks on Code: Code embeddings have supported semantic tasks such as naming, clone detection, completion, summarization, and algorithm classification, as well as optimization-related prediction tasks.Prior work includes device mapping and thread-coarsening prediction for OpenCL code.
Embedding Evaluation: The paper is the first reported work to quantify code-embedding quality through clustering, syntactic analogies, semantic analogies, and categorical distance tests.Most prior studies evaluated embeddings only through downstream task performance.

3 A Robust Distributional Hypothesis of Code

The paper proposes a distributional hypothesis for code in which contextual statements tend to have similar semantics, then defines context and similarity around program execution and state effects. Its representation uses language- and hardware-independent statements with combined data and execution dependence.

Distributional Hypothesis: The proposed code distributional hypothesis states that statements occurring in the same contexts tend to have similar semantics.It adapts the linguistic distributional hypothesis from words to processor operations.
Distributional Hypothesis: The hypothesis is intended as a basis for a learnable representation, but its key terms require explicit definitions.The paper identifies ambiguity in the meanings of statements, contexts, and similarity.
Statements: Statements should be independent of both source language and hardware architecture to support universality and uniformity.The paper contrasts source-language dependence with target-specific machine code.
Context: Context is defined as the union of data dependence and execution dependence, capturing relationships that consecutive code order may miss.Dataflow alone may omit branch and function-call relationships, while long-term dependencies can span a function.
Similarity: Statement similarity is defined through comparable effects on system state or comparable resource consumption under an operational-semantics view.Different variable types in two versions of the same algorithm can therefore correspond to synonymous statements.

4 Contextual Flow Processing

The paper converts source code into LLVM IR and builds contextual flow graphs that represent both data and execution dependencies, including branches, loops, functions, and recursion.

Compilation, Static Single Assignment, and LLVM IR: LLVM provides a language- and hardware-oriented Intermediate Representation that separates source-language frontends from optimization and target backends.LLVM IR supports multiple programming languages, architectures, and optimized constructs.
Compilation, Static Single Assignment, and LLVM IR: SSA assigns each variable once, making dataflow traceable while using φ-expressions to represent values selected by runtime control flow.A φ-expression can combine alternative values across branches.
Contextual Flow Graphs: XFGs represent each program as directed multigraphs whose edges encode either data dependence or execution dependence.Nodes represent variables or label identifiers such as basic blocks and function names.
Contextual Flow Graphs: XFG construction connects within-block data dependencies, inter-block dependencies, and otherwise unparented identifiers to roots.The construction uses two passes over LLVM IR statements.
Contextual Flow Graphs: XFGs create paths through dataflow, branches, loops, and functions, including recursion, with O(n) construction time for n SSA statements.This linear-time construction is intended to support learning over large code corpora.
Contextual Flow Graphs: External statically linked code is traversed as LLVM IR, whereas dynamically linked libraries are represented as call statements.The distinction depends on whether library code is available during compilation.

5 inst2vec: Embedding Statements in Continuous Space

inst2vec learns statement embeddings from preprocessed LLVM IR using XFG-derived context. The resulting space captures semantic relationships, with XFG context size 2 selected as the best-performing configuration in the reported tests.

Statement Preprocessing and Training: inst2vec trains LLVM IR statement embeddings with skip-gram over contextual-flow neighborhoods.A context includes statement pairs connected by paths up to a chosen size.
Statement Preprocessing and Training: Preprocessing removes comments and metadata, replaces identifiers and immediate values with abstract tokens, and inlines data structures.Immediate values are supplied separately to downstream RNNs.
Statement Preprocessing and Training: The training corpus spans multiple disciplines, source languages, and CPU/GPU targets, including C, C++, FORTRAN, and OpenCL.It includes human-written and synthetically generated programs.
Evaluation: The learned space groups semantically similar data-structure operations, including std::complex<float> with its two-float vector representation.This generalizes to user-defined structures containing two floats.
Evaluation: Analogy evaluation tests type, option, conversion, and data-structure relations using vector arithmetic and nearest-neighbor retrieval.Relative-distance tests compare statements using different versus shared resources.
Evaluation: Context size 2 statements produces the best overall analogy and test scores, while combining data and control flow yields better embeddings than either alone.Context sizes 1 and 3 produce larger multi-type clusters and outliers.

6 Code Comprehension Experiments

The experiments use one fixed-embedding LSTM pipeline across algorithm classification, device mapping, and thread-coarsening prediction, comparing it with specialized and manually engineered approaches.

Experimental setup: The same neural network architecture and fixed pre-trained inst2vec matrix are used across all three evaluation tasks.The embedding matrix remains fixed during task training.
Datasets: The algorithm classification benchmark contains 104 program classes written by 500 different people, while the other tasks use an OpenCL code dataset.The classification data come from POJ-104; device mapping and thread coarsening use Cummins et al.'s dataset.
Algorithm classification: 13.8 % lower error establishes a new state-of-the-art for inst2vec on POJ-104 despite embeddings trained without that dataset.The comparison includes TBCNN, a 600-cell RNN, and manually extracted surface features.
Heterogeneous device mapping: Inst2vec outperforms Grewe et al.'s features and matches DeepTune for heterogeneous device mapping, while immediate-value handling surpasses DeepTune's specialized training.The comparisons use runtime-prediction accuracy and resulting speedup.
Thread coarsening prediction: For thread coarsening, inst2vec consistently outperforms manually extracted features, while inst2vec-imm is generally on par with DeepTune but does not surpass DeepTune-TL.The task has only 17 programs per platform, which the paper identifies as a possible explanation for inconsistent gains.

7 Conclusion

The paper concludes that statement semantics can be recovered from contextual information using filtered LLVM IR and combined data- and control-flow, with strong results across three tasks.

Conclusion: Statement semantics can be recovered from context alone when statements use filtered LLVM IR at an appropriate granularity.The conclusion links recovery to both representation granularity and statement grouping.
Conclusion: Combining data- and control-flow groups statements for the proposed representation.This grouping is presented alongside filtered LLVM IR as a basis for semantic recovery.
Conclusion: Across three high-level classification and prediction tasks, the representation outperforms manually extracted features and matches or exceeds two specialized state-of-the-art DNN solutions.The paper frames these results as empirical support for mechanized code comprehension through machine learning.
Future directions: The proposed direction targets mechanized code comprehension for both human-authored and automatically generated code.The authors also identify neural code interpretation and part-based representation refinement as future directions.

A Statement Categories for inst2vec Clustering Results

The appendix defines the statement-category vocabulary and color legend used to interpret inst2vec clustering results, including types, aggregates, conversions, and pointer operations.

Legend: The statement-category legend maps colors to categories appearing in Fig. 6.Table 6 presents the mapping from colors to statement categories.
Type categories: A type operation denotes an operation, function call, or function definition that returns an instance of that type.The category uses the returned type as its organizing property.
Type categories: Pointer categories use type* notation, with chained asterisks representing pointers to pointers.The notation describes pointer depth.
Aggregate categories: Vector and array categories use <d x type> and [d x type> notation for d elements of a specified type.The source defines vectors with angle brackets and arrays with square brackets.
Aggregate categories: Struct/class categories represent aggregate structures containing multiple types, while int covers integers of any bit-width.Floating point includes single- and double-precision values.
Special categories: Void call and invoke categories identify function calls or invocations with no return value; conversion operations denote LLVM type conversions.Conversions do not necessarily translate directly into code.
Special categories: Load-function-pointer and store-function-pointer categories distinguish instructions that read or write function pointers in memory.The distinction is based on the direction of memory access.

B Neural Code Comprehension: Network Architecture

The network converts LLVM IR statements into fixed embeddings, summarizes arbitrary-length programs with stacked LSTMs, optionally adds runtime data, and predicts output categories.

Architecture overview: Fig. 7 depicts the neural network architecture used for the paper's high-level tasks.The appendix describes the underlying layers of the network.
Input and embedding lookup: LLVM IR statements are mapped to their inst2vec vectors, with out-of-vocabulary statements assigned a predefined unknown-token vector.The embedding layer remains fixed during task training and acts as a lookup matrix.
Program characterization: Two LSTM layers transform an arbitrary-length sequence of statement embeddings into a fixed-length vector representing processed-program properties.This fixed-length representation feeds the subsequent prediction layers.
Auxiliary input concatenation: Runtime-only information, such as hardware parameters or data size, can optionally be concatenated with the LSTM output.The auxiliary input enables predictive modeling to incorporate information unavailable from code alone.
Prediction layers: Batch normalization and a 32-unit ReLU dense layer precede an output layer whose unit count matches the possible categories.A sigmoid produces outputs between 0 and 1, with the largest activation determining the prediction.

C Training NCC with Immediate Values: Method Description

NCC handles immediate values through abstraction during inst2vec training and three alternative architectures for reintroducing them downstream. These architectures concatenate or separately process immediate-value sequences alongside statement representations.

Immediate-value abstraction: Raw LLVM IR statements replace immediate values with type tokens such as <INT>, <FLOAT>, and <STRING> before inst2vec training.This reduces vocabulary sparsity and maps nearly identical statements to the same embedding vector.
Architecture overview: Figure 8 compares three architectures for training inst2vec statement sequences with their immediate values in NCC.Dark-orange components represent immediate-value processing, and a yellow plus sign marks the concatenation stage.
Naïve concatenation: The naïve concatenation architecture feeds statement embeddings concatenated with immediate values into the first LSTM layer.Its input size is the embedding dimension plus the length of the immediate-value list, while later NCC layers remain unchanged.
Concatenate then embed: The concatenate-then-embed architecture concatenates statement embeddings with immediate values, then reduces the combined dimension through a fully connected layer.The reduced representation returns to the embedding dimension before program characterization and the remaining network.
Extract then concatenate: The extract-then-concatenate architecture processes the entire program’s immediate-value sequence with a separate LSTM before adding it to program-characterization outputs.This pipeline is designed to extract critical immediate-value information as auxiliary inputs.

D Training NCC with Immediate Values: Exhaustive Results

Tables 7 and 8 report exhaustive results for two performance-prediction tasks across the alternative immediate-value handling modes. The ignore condition provides the simplest NCC baseline for comparison.

Baseline: The ignore column reports results for the simplest NCC version, which omits immediate values.It serves as the baseline against which the other handling modes can be compared.