Install and use NNVM Compiler

Reading time: 20 minutes

NNVM compiler is a graph compiler for the TVM Stack that takes in models in NNVM Intermediate Representation format and compiles them for various backends such as LLVM, METAL, CUDA and others.

Install/ Build NNVM

  • Step 1: Clone the source
git clone --recursive https://github.com/dmlc/nnvm.git
cd nnvm
  • Step 2: Build NNVM
make -j4
  • Step 3: Build the Python Library
cd python
python setup.py install --user
cd ..

Build Graph to Runtime Library

Using NNVM Compiler, you can build the model graph in NNVM Intermediate Representation to a runtime library.


# Compile Model on NNVM compiler
import nnvm.compiler
target = 'llvm'
input_name = sym.list_input_names()[0]
shape_dict = {input_name: x.shape}
with nnvm.compiler.build_config(opt_level=3, 
          add_pass=['AlterOpLayout', 'FoldScaleAxis', 'OpFusion', 
          'PrecomputePrune', 'SimplifyInference']):
    graph, lib, params = nnvm.compiler.build(graph = sym, 
             target_host = 'llvm', target = target, 
             shape = shape_dict, 
             dtype={"int64": "int64", "float32": "float32"}, 
             params = params, layout = "NCHW")

Details/ Configurations

NNVM Compiler supports the following targets:

  • LLVM
  • CUDA
  • METAL
  • OpenCL

NNVM Compiler has four optimization levels:

  • 0
  • 1
  • 2 (default optimization level)
  • 3

NNVM Compiler has several pass optimizations that you may apply such as:

  • AlterOpLayout
  • FoldScaleAxis
  • OpFusion
  • PrecomputePrune
  • SimplifyInference

Input

NNVM Compiler takes the model as two inputs:

  • Graph in NNVM Intermediate Representation
  • Params: parameters of the graph such as weights and bias

Output

NNVM Compiler gives three outputs which we can save for later use and run:

  • graph: Compiler NNVM Graph
  • lib: defines the target environment such as LLVM
  • params: parameters of the compiler NNVM Graph such as weights and bias