ncc
(Neural Code Comprehension) is a general Machine Learning technique to learn semantics from raw code in virtually any programming language. It relies on inst2vec
, an embedding space and graph representation of LLVM IR statements and their context.
This repository contains the code used in [paper]:
Neural Code Comprehension: A Learnable Representation of Code Semantics, Tal Ben-Nun, Alice Shoshana Jakobovits, Torsten Hoefler
Please cite as:
@incollection{ncc,
title = {Neural Code Comprehension: A Learnable Representation of Code Semantics},
author = {Ben-Nun, Tal and Jakobovits, Alice Shoshana and Hoefler, Torsten},
booktitle = {Advances in Neural Information Processing Systems 31},
editor = {S. Bengio and H. Wallach and H. Larochelle and K. Grauman and N. Cesa-Bianchi and R. Garnett},
pages = {3588--3600},
year = {2018},
publisher = {Curran Associates, Inc.},
url = {https://papers.nips.cc/paper/7617-neural-code-comprehension-a-learnable-representation-of-code-semantics.pdf}
}
For training inst2vec
embeddings:
- GNU / Linux or Mac OS
- Python (3.6.5)
- tensorflow (1.7.0) or preferably: tensorflow-gpu (1.7.0)
- networkx (2.1)
- scipy (1.1.0)
- absl-py (0.2.2)
- jinja2 (2.10)
- bokeh (0.12.16)
- umap (0.1.1)
- sklearn (0.0)
- wget (3.2)
Additionally, for training ncc
models:
- GNU / Linux or Mac OS
- Python (3.6.5)
- labm8 (0.1.2)
- keras (2.2.0)
By default, inst2vec
will be trained on publicly available code. Some additional datasets are available on demand and you may add them manually to the training data. For more information on how to do this as well as on the datasets in general, see datasets.
$ python train_inst2vec.py --helpfull # to see the full list of options
$ python train_inst2vec.py \
> # --context_width ... (default: 2)
> # --data ... (default: data/, automatically generated one. You may provide your own)
Alternatively, you may skip this step and use pre-trained embeddings.
$ python train_inst2vec.py \
> --embeddings_file ... (path to the embeddings p-file to evaluate)
> --vocabulary_folder ... (path to the associated vocabulary folder)
We provide the code for training three downstream tasks using the same neural architecture (ncc
) and inst2vec
embeddings.
Algorithm classification
Task: Classify applications into 104 classes given their raw code.
Code and classes provided by https://sites.google.com/site/treebasedcnn/ (see Convolutional neural networks over tree structures for programming language processing)
Train:
$ python train_task_classifyapp.py --helpfull # to see the full list of options
$ python train_task_classifyapp.py
Alternatively, display results from a pre-trained model.
Optimal device mapping prediction
Task: Predict the best-performing compute device (e.g., CPU, GPU) Code and classes provided by https://github.com/ChrisCummins/paper-end2end-dl (see End-to-end Deep Learning of Optimization Heuristics)
Train:
$ python train_task_devmap.py --helpfull # to see the full list of options
$ python train_task_devmap.py
Alternatively, display results from a pre-trained model.
Optimal thread coarsening factor prediction
Code and classes provided by https://github.com/ChrisCummins/paper-end2end-dl (see End-to-end Deep Learning of Optimization Heuristics)
Train:
$ python train_task_threadcoarsening.py --helpfull # to see the full list of options
$ python train_task_threadcoarsening.py
Alternatively, display results from a pre-trained model.
We would be thrilled if you used and built upon this work. Contributions, comments, and issues are welcome!
NCC is published under the New BSD license, see LICENSE.