Skip to content
An experimental ahead of time compiler for Relay.
Python
Branch: master
Clone or download
MarisaKirisame Merge pull request #25 from slyubomirsky/record-fix
Using wrong variable for time recording
Latest commit 1f2b0e0 Aug 22, 2019
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
aot Using wrong variable for time recording Aug 22, 2019
examples README tweaks, make sure example runs Apr 16, 2019
test Add test case for passing a closure of a locally recursive function Aug 8, 2019
.gitignore save Apr 17, 2019
LICENSE Add LICENSE Apr 16, 2019
Pipfile Hacking on bug fix Mar 25, 2019
Pipfile.lock Hacking on bug fix Mar 25, 2019
README.md grammar Aug 8, 2019

README.md

relay-aot

An experimental ahead of time compiler for Relay.

The ahead of time compiler enables the execution of Relay code without requiring a framework's interpreter written in C++ or Python.

The removal of framework and interpretation overhead combined with optimized operators produced by TVM's operators dramatically reduces execution time. Additionally, the compiler produces a single binary which only depends on TVM, simplifying the deployment story. This repository contains an external library which implements ahead of time compilation for Relay. The current approach is a proof of concept which lowers Relay to C++, and relies on a C++ compiler such as gcc or clang to produce an executable.

The ahead of time compiler comes as a standalone library which exposes a primitive compile function which compiles a relay.Function into a Python closure which wraps the compiled native code.

The compiler's design is straight forward. It lowers functions into a small C++-like IR, and generates a C++ program which can be compiled and dynamically linked. We extract the corresponding symbol from the dynamic library, and wrap it as a Python closure.

You can see an example below:

import numpy as np
import tvm
from tvm.relay import Module, GlobalVar, Function
from aot import compile

def double_example():
    # Declare a Relay module.
    mod = Module()

    # Implement the double function.
    x = var('x', shape=())
    double = GlobalVar('double')
    mod[double] = Function([x], x + x)

    # Generate a function which calls double twice.
    x = var('x', shape=())
    f = Function([x], double(double(x)))
    # Compile the function.
    cfunc = compile(f, mod)

    a = tvm.nd.array(np.array(1.5, dtype='float32'))
    output = cfunc(a).asnumpy() # array(6.)

Currently there is no Python package due to the lack of a package for tvm itself. You can use the ahead of time compiler by adding it to your PYTHONPATH.

export PYTHONPATH="THE_AOT_PATH:${PYTHONPATH}"

You must set the variable TVM_HOME in order to use the native compiler currently.

You can test to ensure you setup the ahead of time correctly by running:

TVM_HOME=~/Git/tvm python3 examples/readme_ex.py

Note: this only sets the TVM_HOME for this command you must reexport in your shell.

You can’t perform that action at this time.