Open Deep Learning Compiler Stack

Documentation | Contributors | Community | Release Notes

Apache TVM is a compiler stack for deep learning systems. It is designed to close the gap between the productivity-focused deep learning frameworks, and the performance- and efficiency-focused hardware backends. TVM works with deep learning frameworks to provide end to end compilation to different backends.

(See the original README for more information)

Setup

The following instructions have been adapted from TVM's official "Installing TVM from source" guide

Getting started

Start by cloning our TVM fork.

git clone --recursive https://github.com/agamkohli9/relax

Update submodules

cd ./relax
git submodule init
git submodule update

Build TVM

# Setup build directory
mkdir build
cp cmake/config.cmake build

# Build TVM
make -C ./build

Python package setup

Update your python path to include your local build of TVM.

export TVM_HOME=/path/to/relax
export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}

Python scripting environment We recommend using a virtual environment. See the official documentation here.

Create a virtual environment and install dependencies.

cd ./demo

# Create and activate a virtual environment
python -m venv env
source env/bin/activate

# Install dependencies
pip install -r requirements.txt

Project Structure

├── definitions.py
├── gen_lib.py                          <- generate compiled libs for benchmarking
├── gen_plots.py                        <- generate plots from compiled libs
├── gen_relax_ir.py                     <- generate optimized Relax IR from inputs
├── logger.py
├── plot.png
├── relax
│   ├── __init__py
│   ├── modules.py
│   └── optimize.py
├── relax-ir
│   ├── ModuleBasic-opt.relax
│   ├── ModuleBasic-raw.relax
|   ...
├── relay
│   ├── __init__.py
│   ├── config.py
│   └── lib.py
└── requirements.txt

Optimizing models

For all of the following, ensure you are in the ./demo directory with a python virtual environment active and all dependencies installed..

Testing optimizations

Run an optimization pass defined in demo/relax/optimize.py for all example programs included in demo/relax/modules.py.

python3 ./gen_relax_ir.py

This produces two .relax files for each input (one containing the origininal, unmodified IR, and one containing the optimized IR). All generated .relax files are saved to demo/relax-ir.

Example:

Input	Output
# ModuleBasic-raw.relax @R.function def foo() -> R.Tensor(None, dtype="int32", ndim=0): # block 0 res: R.Tensor((), dtype="int32") = R.add(30, 40) return res	# ModuleBasic-opt.relax @R.function def foo() -> R.Tensor(None, dtype="int32", ndim=0): # block 0 res: R.Tensor((), dtype="int32") = 70 return 70

Input

Output

# ModuleBasic-raw.relax

@R.function
def foo() -> R.Tensor(None, dtype="int32", ndim=0):
    # block 0
    res: R.Tensor((), dtype="int32") = R.add(30, 40)
    return res

# ModuleBasic-opt.relax

@R.function
def foo() -> R.Tensor(None, dtype="int32", ndim=0):
    # block 0
    res: R.Tensor((), dtype="int32") = 70
    return 70