IITH-Compilers / IR2Vec

Implementation of IR2Vec, LLVM IR Based Scalable Program Embeddings
Other
84 stars 39 forks source link
embeddings llvm

IR2Vec

IR2Vec is a LLVM IR based framework to generate distributed representations for the source code in an unsupervised manner, which can be used to represent programs as input to solve machine learning tasks that take programs as inputs.

This repo contains the source code and relevant information described in the paper (arXiv). Please see here for more details.

IR2Vec: LLVM IR Based Scalable Program Embeddings, S. VenkataKeerthy, Rohit Aggarwal, Shalini Jain, Maunendra Sankar Desarkar, Ramakrishna Upadrasta, and Y. N. Srikant

LLVM PyPI Version Tests Publish pre-commit checks

Image

LLVM Version Archive

LLVM Version Branch
LLVM 17.0.6 main
LLVM 16.0.1 llvm16
LLVM 14.0.1 llvm14
LLVM 12.0.0 llvm12
LLVM 10.0.1 llvm10
LLVM 8.0.1 llvm8

Table Of Contents

Installation

IR2Vec can be installed in different ways to accommodate individual preferences and requirements effectively. You may select to install via a user-friendly Python wheel setup if you are a Python user, or opt for a C++ based installation if you are looking to integrate with a compiler pass or necessitate advanced control and enhanced integration capabilities. The detailed setup steps are mentioned in the following sections.

Python

If you prefer working with Python, you can easily install IR2Vec using pip.

pip install -U ir2vec

Now, you can import and use IR2Vec in your Python projects. Make sure you have a good understanding of Python and its package management system.

We are actively working on improving the Python interfaces and providing better support. If you find any good-to-have interfaces that you may need for your use case missing, please feel free to raise a request.

Cpp

If you're a C++ developer and require low-level control, optimization, or integration with C++ projects, you can build IR2Vec from source. First, ensure the below requirements are satisfied, then follow the steps mentioned in the Building from source section.

Requirements

(Experiments are done on an Ubuntu 20.04 machine)

Building from source

  1. mkdir build && cd build
  2. IR2Vec uses Eigen library. If your system already have Eigen (3.3.7) setup, you can skip this step.
    1. Download and extract the released version.
      • wget https://gitlab.com/libeigen/eigen/-/archive/3.3.7/eigen-3.3.7.tar.gz
      • tar -xvzf eigen-3.3.7.tar.gz
    2. mkdir eigen-build && cd eigen-build
    3. cmake ../eigen-3.3.7 && make
    4. cd ../
  3. cmake -DLT_LLVM_INSTALL_DIR=<path_to_LLVM_build_dir> -DEigen3_DIR=<path_to_eigen_build_dir> [-DCMAKE_INSTALL_PREFIX=<install_dir>] ..
  4. make [&& make install]

This process would generate ir2vec binary under build/bin directory, libIR2Vec.a and libIR2Vec.so under build/lib directory.

To ensure the correctness, run make check_ir2vec

Generating program representations

IR2Vec can be used either as a stand-alone tool using binary or can be integrated with any third-party tools using libraries. Please see below for the usage instructions.

Using Binary

ir2vec -\<mode> -dim \<dimensions> -o \<output-file> -level \<p|f> -class \<class-number> -funcName=\<function-name> \<input-ll-file>

Command-Line options

Please use --help for further details.

Format of the output embeddings in output_file

<class-number> <Embeddings>

class-number would be printed only if it is not -1

<function-name> = <Embeddings>

Flow-Aware Embeddings

For all functions

For a specific function

Symbolic Embeddings

For all functions

Using Libraries

The libraries can be installed by passing the installation location to the CMAKE_INSTALL_PREFIX flag during cmake followed by make install. The interfaces are available in IR2Vec.h. External projects that would like to use IR2Vec can access the functionality using these exposed interfaces on including IR2Vec.h from the installed location after linking statically or dynamically.

The following template can be used to link IR2vec libraries on a cmake based project.

set(IR2VEC_INSTALL_DIR "" CACHE PATH "IR2Vec installation directory")
include_directories("${IR2VEC_INSTALL_DIR}/include")
target_link_libraries(<your_executable_or_library> PUBLIC ${IR2VEC_INSTALL_DIR}/lib/<libIR2Vec.a or libIR2Vec.so>)

And then pass the location of IR2Vec's install prefix to DIR2VEC_INSTALL_DIR during cmake.

The following example snippet shows how to query the exposed vector representations.

#include "IR2Vec.h"

// Creating object to generate FlowAware representation
auto ir2vec =
      IR2Vec::Embeddings(<LLVM Module>, IR2Vec::IR2VecMode::FlowAware, <DIM>);

// Getting Instruction vectors corresponding to the instructions in <LLVM Module>
auto instVecMap = ir2vec.getInstVecMap();
// Access the generated vectors
for (auto instVec : instVecMap) {
    outs() << "Instruction : ";
    instVec.first->print(outs());
    outs() << ": ";

    for (auto val : instVec.second)
      outs() << val << "\t";
}

// Getting vectors corresponding to the functions in <LLVM Module>
auto funcVecMap = ir2vec.getFunctionVecMap();
// Access the generated vectors
for (auto funcVec : funcVecMap) {
    outs() << "Function : " << funcVec.first->getName() << "\n";
    for (auto val : funcVec.second)
      outs() << val << "\t";
  }

// Getting the program vector
auto pgmVec = ir2vec.getProgramVector();
// Access the generated vector
for (auto val : pgmVec)
    outs() << val << "\t";

Using Python package (IR2Vec-Wheels)

Initialization -ir2vec.initEmbedding

Description: Initialize IR2Vec embedding for an LLVM IR file.

Parameters:

Returns:

Example:

import ir2vec

# Approach 1
initObj = ir2vec.initEmbedding("/path/to/file.ll", "fa", "p")

# Approach 2
initObj = ir2vec.initEmbedding("/path/to/file.ll", "fa", "p", 100)

# Approach 3
initObj = ir2vec.initEmbedding("/path/to/file.ll", "fa", "p", 100, "output.txt")

getProgramVector

Description: Gets the program-level vector representation.

Parameters: optional

Returns:

Example:

# Getting the program-level vector
progVector = initObj.getProgramVector()

getFunctionVectors

Description: Gets function-level vectors for all functions in the LLVM IR file.

Parameters: optional

Returns:

Example:

# Getting function-level vectors
functionVectorMap = initObj.getFunctionVectors()

getInstructionVectors

Description: Gets instruction-level vectors for all instructions in the LLVM IR file.

Parameters: optional

Returns:

Example:


# Getting instruction-level vectors
instructionVectorsList = initObj.getInstructionVectors()

Example

import ir2vec
import numpy as np

# IR2Vec Python APIs can be used in two ways. As shown below.
initObj = ir2vec.initEmbedding("/path/to/file.ll", "fa", "p")

#Approach 1
progVector1 = ir2vec.getProgramVector(initObj)
functionVectorMap1 = ir2vec.getFunctionVectors(initObj)
instructionVectorsList1 = ir2vec.getInstructionVectors(initObj)

#Approach 2
progVector2 = initObj.getProgramVector()
functionVectorMap2 = initObj.getFunctionVectors()
instructionVectorsList2 = initObj.getInstructionVectors()

# Both the approaches would result in same outcomes
assert(np.allclose(progVector1,progVector2))

for fun, funcObj in functionVectorMap1.items():
    assert fun == funcObj["demangledName"]
    functionOutput1 = ir2vec.getFunctionVectors(
        initObj,
        funcObj["actualName"],
    )
    functionOutput2 = initObj.getFunctionVectors(
        funcObj["actualName"]
    )
    assert(np.allclose(functionOutput1[fun]["vector"],functionOutput2[fun]["vector"]))

Binaries, Libraries and Wheels - Artifacts

Binaries, Libraries (.a and .so), and whl files are autogenerated for every relevant check-in using GitHub Actions. Such generated artifacts are tagged along with the successful runs of Publish and Build Wheels actions.

Experiments

Note

The results mentioned in the experiment's scripts/the published version are not updated for this branch. The experimental results for this branch would be different when compared to the published version. For comparison, use the release corresponding to v0.1.0.

Citation

@article{VenkataKeerthy-2020-IR2Vec,
author = {VenkataKeerthy, S. and Aggarwal, Rohit and Jain, Shalini and Desarkar, Maunendra Sankar and Upadrasta, Ramakrishna and Srikant, Y. N.},
title = {{IR2Vec: LLVM IR Based Scalable Program Embeddings}},
year = {2020},
issue_date = {December 2020},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
volume = {17},
number = {4},
issn = {1544-3566},
url = {https://doi.org/10.1145/3418463},
doi = {10.1145/3418463},
journal = {ACM Trans. Archit. Code Optim.},
month = dec,
articleno = {32},
numpages = {27},
keywords = {heterogeneous systems, representation learning, compiler optimizations, LLVM, intermediate representations}
}

Contributions

Please feel free to raise issues to file a bug, pose a question, or initiate any related discussions. Pull requests are welcome :)

License

IR2Vec is released under a Apache License v2.0 with LLVM Exceptions License. See the LICENSE file for more details.