WebAssembly Standalone for Deep Learning Framework with TVM Runtime

Experimental notice: This project is still experimental and only serves as a proof of concept for running deep learning frameworks on WebAssembly runtime with TVM stack.

WebAssembly Standalone for Deep Learning Framework with TVM Runtime

Motivation

As demonstrated in TVM runtime tutorials, TVM already supports WASM as the optional hardware backend, so we can leverage the features of WebAssembly (portability, security) and TVM runtime (domain-specific, optimization) to build a flexible and auto-optimized graph compiler for all deep learning frameworks.

Framework Landscape

The figures below demonstrate the whole landscape of running deep learning frameworks on WASM runtime with TVM compiler stack.

WASM graph generation

   _ _ _ _ _ _ _ _ _ _        _ _ _ _ _ _ _        _ _ _ _ _ _ _ _ _ _ _ _
  |                   |      |             |      |                       |
  |  Framework Model  | ---> |  ONNX Model | ---> |  TVM Relay Python API |
  |_ _ _ _ _ _ _ _ _ _|      |_ _ _ _ _ _ _|      |_ _ _ _ _ _ _ _ _ _ _ _|
                                                             ||
                                                             \/
             _ _ _ _ _ _ _ _ _ _ _                  _ _ _ _ _ _ _ _ _ _ _
            |                     |                |                     |
            | WASM Graph Builder  |                |  TVM Compiler Stack |
            |    (TVM runtime)    |                |_ _ _ _ _ _ _ _ _ _ _|
            |_ _ _ _ _ _ _ _ _ _ _|                          ||
                      ||                                     \/
  _ _ _ _ _ _ _ _ _   ||   _ _ _ _ _ _ _ _ _ _            _ _ _ _ _
 |                 |  \/  |                   |  llvm-ar |         |
 | wasm_graph.wasm | <--- | libgraph_wasm32.a | <------- | graph.o |
 |_ _ _ _ _ _ _ _ _|      |_ _ _ _ _ _ _ _ _ _|          |_ _ _ _ _|

WASM graph loading

     _ _ _ _ _ _ _ _ _ _ _
    |                     |
    |  WASM Graph Loader  |
    |   (WASM runtime)    |
    |_ _ _ _ _ _ _ _ _ _ _|
              ||
              \/
      _ _ _ _ _ _ _ _ _ _
     |                   |
     |  wasm_graph.wasm  |
     |_ _ _ _ _ _ _ _ _ _|

Project Status

This project should be considered experimental at the very early stage, all rich features are under active development. Here is the current operator support matrix:

Model Name	Status
ResNet50	✔️
LeNet	—

NOTICE: Currently this project is ONLY tested on Ubuntu system, so Ubuntu 16.04+ should be prepared as the testing environment.

PoC Guidelines

Pre-installation

Rust
Before running this demo, please make sure Rust has been installed.
After Rust installed, execute the code below to add wasm32-wasi target:
```
rustup target add wasm32-wasi
```
TVM
Please follow TVM installations for the detailed instruction.
LLVM
LLVM 10.0 or later is REQUIRED.

Build ResNet50 model

Build DL library in the WebAssembly format.

Download model

cd wasm-graph/tools && wget https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet50v1/resnet50v1.onnx

Compile

LLVM_AR=llvm-ar-10 python ./build_graph_lib.py -O3 ./resnet50v1.onnx

Build wasm-graph package

cd wasm-graph && cargo build --release
cp ./target/wasm32-wasi/release/wasm_graph.wasm ./lib/wasm_graph_resnet50.wasm

Test

Before running this demo, please make sure Rust has been installed.

Next run the command below to install the runtime package for testing (rust REQUIRED):

cd wasm-runtime/tests/test_graph_resnet50 && cargo build

Check the usage of test_graph_resnet50:

~# ./target/debug/test_graph_resnet50 -h

Usage: ./target/debug/test_graph_resnet50 [options]

Options:
    -g, --wasm-graph-file FILE_PATH
                        set the path to wasm graph file
    -i, --input-data-file FILE_PATH
                        set the path to input image file
    -l, --label-class-file FILE_PATH
                        set the path to label class file
    -h, --help          print this help menu

Next perform model inference using these commands below:

$ cp ../../../wasm-graph/lib/wasm_graph_resnet50.wasm ./
$ wget -O cat.png https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true
$ wget -O synset.csv https://raw.githubusercontent.com/kazum/tvm-wasm/master/synset.csv
$ ./target/debug/test_graph_resnet50 -g ./wasm_graph_resnet50.wasm -i ./cat.png -l ./synset.csv
original image dimensions: (256, 256)
resized image dimensions: (224, 224)
input image belongs to the class `tabby, tabby cat`

Future Work

More networks support

TODO

Performance benchmark

We are working on several improvements on performances:

WebAssembly simd128 support (Done)
Auto-tvm enhancement for llvm target

Native TVM Rust runtime support

TODO

Appendix

System packages install

Rust (latest version)
If you are running Windows, to install Rust, download and run the RUST-INIT.EXE, and then follow the onscreen instructions.
If you are a Linux user, run the following in your terminal, then follow the on-screen instructions to install Rust.
```
curl https://sh.rustup.rs -sSf | sh
```