Skip to content

uwsampl/relay-aot

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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.

About

An experimental ahead of time compiler for Relay.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 4

  •  
  •  
  •  
  •  

Languages