diff --git a/apps/README.md b/apps/README.md index 685750633493c..cad9b30220730 100644 --- a/apps/README.md +++ b/apps/README.md @@ -26,3 +26,4 @@ If you are interested in writing optimized kernels with TVM, checkout [TOPI: TVM - [android_rpc](android_rpc) Android RPC server. - [benchmark](benchmark) Example end to end compilation benchmarks - [howto_deploy](howto_deploy) Tutorial on how to deploy TVM with minimum code dependency. +- [wasm-dlbackend-tvm](wasm-dlbackend-tvm) WebAssembly backend for deep learning framework with TVM runtime. diff --git a/apps/wasm-dlbackend-tvm/.cargo/config b/apps/wasm-dlbackend-tvm/.cargo/config new file mode 100644 index 0000000000000..e17b808e4b074 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/.cargo/config @@ -0,0 +1,3 @@ +[build] +target = "wasm32-wasi" +rustflags = ["-C", "link-arg=--whole-archive", "-C", "link-arg=-lops_wasm32"] diff --git a/apps/wasm-dlbackend-tvm/.gitignore b/apps/wasm-dlbackend-tvm/.gitignore new file mode 100644 index 0000000000000..54fb6c73048d3 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/.gitignore @@ -0,0 +1,8 @@ +# Built packages +**/lib/ + + +#Added by cargo + +**/target/ +**/Cargo.lock diff --git a/apps/wasm-dlbackend-tvm/Cargo.toml b/apps/wasm-dlbackend-tvm/Cargo.toml new file mode 100644 index 0000000000000..596890354f8f5 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "wasm-dlbackend-tvm" +version = "0.1.0" +authors = ["leonwanghui "] +edition = "2018" +description = "WebAssembly backend to deep learning framework using TVM runtime" +readme = "README.md" +repository = "https://github.com/apache/incubator-tvm" +license = "Apache-2.0" +keywords = ["wasm", "machine learning", "tvm"] + +[package.metadata] +wasm-opt = true +wasm-name-section = false +wasm-producers-section = false + +[profile.release] +lto = true +opt-level = 's' + +[lib] +crate-type = ['cdylib'] + +[dependencies] +serde = "1.0.53" +serde_derive = "1.0.53" +serde_json = "1.0.53" +ndarray = "0.12" +tvm-common = { version = "0.1", path = "../../rust/common" } +tvm-runtime = { version = "0.1", path = "../../rust/runtime" } diff --git a/apps/wasm-dlbackend-tvm/README.md b/apps/wasm-dlbackend-tvm/README.md new file mode 100644 index 0000000000000..42e29081606a8 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/README.md @@ -0,0 +1,137 @@ +# WebAssembly Backend 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 (such like [MindSpore](https://github.com/mindspore-ai/mindspore)) on [WebAssembly runtime](https://github.com/bytecodealliance/wasmtime) with [TVM stack](https://tvm.apache.org/). + +- [WebAssembly Backend for Deep Learning Framework with TVM Runtime](#webassembly-backend-for-deep-learning-framework-with-tvm-runtime) + - [Motivation](#motivation) + - [Framework Landscape](#framework-landscape) + - [Project Status](#project-status) + - [PoC Guidelines](#poc-guidelines) + - [Pre-installation](#pre-installation) + - [Build wasm-dlbackend-tvm package](#build-wasm-dlbackend-tvm-package) + - [Test](#test) + - [Future Work](#future-work) + - [Operator enhancement](#operator-enhancement) + - [Performance benchmark](#performance-benchmark) + - [Native TVM Rust runtime support](#native-tvm-rust-runtime-support) + - [Appendix](#appendix) + - [System packages install](#system-packages-install) + - [Contribution](#contribution) + +## Motivation + +TVM hardware support + +As demonstrated in TVM runtime [tutorials](https://tvm.apache.org/docs/tutorials/relay_quick_start.html), 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 operator backend for all deep learning frameworks. + +## Framework Landscape + +The figure below demonstrates the whole landscape of running deep learning framework on WASM runtime with TVM compiler stack. +``` + _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ +| | _ _ _ _ _ _ _ _ _ _ _ +| Framework Frontend Expression | | | +|_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _| | TVM (TE) Python API | + || |_ _ _ _ _ _ _ _ _ _ _| + \/ || + _ _ _ _ _ _ _ _ _ _ _ _ _ _ \/ + | | _ _ _ _ _ _ _ _ _ _ _ + | Framework WASM Backend | | | + | (WASM runtime) | | TVM Compiler Stack | + |_ _ _ _ _ _ _ _ _ _ _ _ _ _| |_ _ _ _ _ _ _ _ _ _ _| + || || + \/ \/ + _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ + | | | | llvm-ar | | + | TVM Runtime | <--- | libops_wasm32.a | <------- | add.o sub.o | + |_ _ _ _ _ _ _ _| |_ _ _ _ _ _ _ _ _| |_ _ _ _ _ _ _| +``` + +## 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: + +| Operator Name | FP32 | INT32 | INT8 | +| ------------- | ---- | ----- | ---- | +| Add | ✔️ |
—
|
—
| +| Sub | ✔️ |
—
|
—
| + +**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](#system-packages-install) has been installed. + + After Rust installed, execute the code below to add `wasm32-wasi` target: + ```shell + rustup target add wasm32-wasi + cargo install cargo-wasi + ``` + +* TVM + + Please follow TVM [installations](https://tvm.apache.org/docs/install/index.html), `export TVM_HOME=/path/to/tvm` and add `libtvm_runtime` to your `LD_LIBRARY_PATH`. + + *Note:* To run the end-to-end examples and tests, `tvm` and `topi` need to be added to your `PYTHONPATH` or it's automatic via an Anaconda environment when it is installed individually. + +### Build wasm-dlbackend-tvm package + +```shell +cd wasm-dlbackend-tvm && cargo wasi build --release +``` + +### Test + +Run the command below to install the frontend package for testing (`rust` REQUIRED): +```shell +cd wasm-dlfrontend/ && cargo build --release +cp ./target/release/wasm-dlfrontend /usr/local/bin/ +``` + +Check the usage of `wasm-dlfrontend`: + +```shell +~# wasm-dlfrontend -h + +Usage: wasm-dlfrontend [options] + +Options: + -c, --ms-backend-config FILE_PATH + set wasm backend config file + -o, --op-type VALUE set the operator type, currently ONLY support Add and + Sub, default: Add. + -h, --help print this help menu +``` + +## Future Work + +### Operator enhancement +TODO + +### Performance benchmark +TODO + +### 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](https://win.rustup.rs/), 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. + + ```shell + curl https://sh.rustup.rs -sSf | sh + ``` + +## Contribution + +Lastly very thanks [@kazum](https://github.com/kazum) for having offered a lot of help when implementing this project. diff --git a/apps/wasm-dlbackend-tvm/build.rs b/apps/wasm-dlbackend-tvm/build.rs new file mode 100644 index 0000000000000..20507fd0dd545 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/build.rs @@ -0,0 +1,32 @@ +use std::{path::PathBuf, process::Command}; + +fn main() { + let mut out_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + out_dir.push("lib"); + + if !out_dir.is_dir() { + std::fs::create_dir(&out_dir).unwrap(); + } + + Command::new(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tools/build_ops_lib.py" + )) + .arg(&out_dir) + .output() + .expect("Failed to execute command!"); + + let ar = option_env!("LLVM_AR").unwrap_or("llvm-ar-10"); + let add_obj_file = out_dir.join("add.o"); + let sub_obj_file = out_dir.join("sub.o"); + let lib_file = out_dir.join("libops_wasm32.a"); + Command::new(ar) + .arg("rcs") + .arg(&lib_file) + .arg(&add_obj_file) + .arg(&sub_obj_file) + .output() + .expect("Failed to execute command!"); + + println!("cargo:rustc-link-search=native={}", out_dir.display()); +} diff --git a/apps/wasm-dlbackend-tvm/src/lib.rs b/apps/wasm-dlbackend-tvm/src/lib.rs new file mode 100644 index 0000000000000..acc1937da4153 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/src/lib.rs @@ -0,0 +1,31 @@ +#[macro_use] +extern crate serde_derive; +#[macro_use] +extern crate tvm_runtime; + +mod ops; +use ops::types::Status; +mod utils; + +#[no_mangle] +pub extern "C" fn run(op_type: i32, in_addr: i32, in_size: i32, out_addr: i32) -> i32 { + let inputs = utils::load_inputs(in_addr, in_size as usize); + if ops::validate_inputs(&inputs) != Status::Succeed { + return 0i32; + } + + let op_instance = ops::operator_instantiate(op_type); + let (a_shape, b_shape, c_shape) = ops::parse_inputs_shape(&inputs); + if op_instance.init(a_shape, b_shape, c_shape) != Status::Succeed { + return 0i32; + }; + + let (in_tensors, out_tensor) = ops::parse_inputs_tensor(&inputs); + let (stat, output) = op_instance.launch(in_tensors, out_tensor); + if stat != Status::Succeed { + return 0i32; + } + + let out_size = utils::store_output(out_addr, output); + out_size as i32 +} diff --git a/apps/wasm-dlbackend-tvm/src/ops/add.rs b/apps/wasm-dlbackend-tvm/src/ops/add.rs new file mode 100644 index 0000000000000..cf9320be691ec --- /dev/null +++ b/apps/wasm-dlbackend-tvm/src/ops/add.rs @@ -0,0 +1,63 @@ +use super::types::*; +use tvm_runtime::{Module as _, SystemLibModule}; + +extern "C" { + fn __wasm_call_ctors(); +} + +pub struct TVMAddOp {} + +impl TVMAddOp { + pub fn new() -> Self { + Self {} + } +} + +impl Operator for TVMAddOp { + fn init(&self, a_shape: Vec, b_shape: Vec, c_shape: Vec) -> Status { + if !((a_shape.len() == b_shape.len() + && a_shape + .iter() + .zip(&b_shape) + .filter(|&(a, b)| a == b) + .count() + == a_shape.len()) + && (b_shape.len() == c_shape.len() + && b_shape + .iter() + .zip(&c_shape) + .filter(|&(b, c)| b == c) + .count() + == c_shape.len())) + { + println!("Both dimension size and shape for Add operator should be equal!"); + return Status::InitFailed; + } + + println!("TVM Add operator init success!"); + Status::Succeed + } + + fn launch(&self, mut inputs: Vec, output: Tensor) -> (Status, Tensor) { + if inputs.len() != 2 { + println!("Inputs tensor length should be 2!"); + return (Status::LaunchFailed, Tensor::default()); + } + let mut l_tensor = inputs.get_mut(0).unwrap().as_dltensor(); + let mut r_tensor = inputs.get_mut(1).unwrap().as_dltensor(); + let mut out_tensor = output.as_dltensor(); + + unsafe { + // This is necessary to invoke TVMBackendRegisterSystemLibSymbol + // API calls. + __wasm_call_ctors(); + } + let syslib = SystemLibModule::default(); + let add = syslib.get_function("add").expect("add function not found!"); + call_packed!(add, &mut l_tensor, &mut r_tensor, &mut out_tensor).unwrap(); + + let output: Tensor = out_tensor.into(); + println!("TVM Add operator run success!"); + (Status::Succeed, output) + } +} diff --git a/apps/wasm-dlbackend-tvm/src/ops/mod.rs b/apps/wasm-dlbackend-tvm/src/ops/mod.rs new file mode 100644 index 0000000000000..c4188b46fcafc --- /dev/null +++ b/apps/wasm-dlbackend-tvm/src/ops/mod.rs @@ -0,0 +1,46 @@ +mod add; +use add::TVMAddOp; +mod sub; +use sub::TVMSubOp; +pub mod types; +use types::*; + +use std::boxed::Box; + +pub fn operator_instantiate(op_type: i32) -> Box { + match OpType::from(op_type) { + OpType::Add => Box::new(TVMAddOp::new()), + OpType::Sub => Box::new(TVMSubOp::new()), + } +} + +pub fn validate_inputs(inputs: &Vec) -> Status { + if (inputs.len() == 3 + && !(inputs[0].dtype() == inputs[1].dtype() && inputs[0].dtype() == inputs[2].dtype())) + || (inputs.len() == 2 && inputs[0].dtype() != inputs[1].dtype()) + { + println!("The dtype of inputs and outputs is not equal!"); + Status::ValidateFailed + } else { + Status::Succeed + } +} + +pub fn parse_inputs_shape(inputs: &Vec) -> (Vec, Vec, Vec) { + if inputs.len() == 3 { + (inputs[0].shape(), inputs[1].shape(), inputs[2].shape()) + } else { + (inputs[0].shape(), inputs[1].shape(), Vec::new()) + } +} + +pub fn parse_inputs_tensor(inputs: &Vec) -> (Vec, Tensor) { + if inputs.len() == 3 { + ( + vec![inputs[0].clone(), inputs[1].clone()], + inputs[2].clone(), + ) + } else { + (vec![inputs[0].clone()], inputs[1].clone()) + } +} diff --git a/apps/wasm-dlbackend-tvm/src/ops/sub.rs b/apps/wasm-dlbackend-tvm/src/ops/sub.rs new file mode 100644 index 0000000000000..6e1f3aac78c9e --- /dev/null +++ b/apps/wasm-dlbackend-tvm/src/ops/sub.rs @@ -0,0 +1,63 @@ +use super::types::*; +use tvm_runtime::{Module as _, SystemLibModule}; + +extern "C" { + fn __wasm_call_ctors(); +} + +pub struct TVMSubOp {} + +impl TVMSubOp { + pub fn new() -> Self { + Self {} + } +} + +impl Operator for TVMSubOp { + fn init(&self, a_shape: Vec, b_shape: Vec, c_shape: Vec) -> Status { + if !((a_shape.len() == b_shape.len() + && a_shape + .iter() + .zip(&b_shape) + .filter(|&(a, b)| a == b) + .count() + == a_shape.len()) + && (b_shape.len() == c_shape.len() + && b_shape + .iter() + .zip(&c_shape) + .filter(|&(b, c)| b == c) + .count() + == c_shape.len())) + { + println!("Both dimension size and shape for Sub operator should be equal!"); + return Status::InitFailed; + } + + println!("TVM Sub operator init success!"); + Status::Succeed + } + + fn launch(&self, mut inputs: Vec, output: Tensor) -> (Status, Tensor) { + if inputs.len() != 2 { + println!("Inputs tensor length should be 2!"); + return (Status::LaunchFailed, Tensor::default()); + } + let mut l_tensor = inputs.get_mut(0).unwrap().as_dltensor(); + let mut r_tensor = inputs.get_mut(1).unwrap().as_dltensor(); + let mut out_tensor = output.as_dltensor(); + + unsafe { + // This is necessary to invoke TVMBackendRegisterSystemLibSymbol + // API calls. + __wasm_call_ctors(); + } + let syslib = SystemLibModule::default(); + let sub = syslib.get_function("sub").expect("sub function not found!"); + call_packed!(sub, &mut l_tensor, &mut r_tensor, &mut out_tensor).unwrap(); + + let output: Tensor = out_tensor.into(); + println!("TVM Sub operator run success!"); + (Status::Succeed, output) + } +} diff --git a/apps/wasm-dlbackend-tvm/src/ops/types.rs b/apps/wasm-dlbackend-tvm/src/ops/types.rs new file mode 100644 index 0000000000000..b87fa0de0c7a5 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/src/ops/types.rs @@ -0,0 +1,194 @@ +use std::{ + any::TypeId, + convert::From, + os::raw::{c_int, c_void}, + slice, +}; +pub use tvm_common::ffi::DLTensor; +use tvm_common::ffi::{ + DLContext, DLDataType, DLDataTypeCode_kDLFloat, DLDataTypeCode_kDLInt, DLDeviceType_kDLCPU, +}; + +pub trait Operator { + fn init(&self, a_shape: Vec, b_shape: Vec, c_shape: Vec) -> Status; + + fn launch(&self, inputs: Vec, output: Tensor) -> (Status, Tensor); +} + +#[derive(Debug, PartialEq)] +pub enum Status { + Succeed, + ValidateFailed, + InitFailed, + LaunchFailed, +} + +#[derive(Debug, PartialEq, Clone)] +pub enum OpType { + Add, + Sub, +} + +impl From for OpType { + fn from(op_type: i32) -> Self { + match op_type { + 0i32 => OpType::Add, + 1i32 => OpType::Sub, + _ => OpType::Add, + } + } +} + +#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)] +pub enum DataType { + FP32, + INT32, + INT8, +} + +impl DataType { + pub fn as_dldtype(&self) -> DLDataType { + match self { + DataType::INT32 => DLDataType { + code: DLDataTypeCode_kDLInt as u8, + bits: 32u8, + lanes: 1u16, + }, + DataType::INT8 => DLDataType { + code: DLDataTypeCode_kDLInt as u8, + bits: 8u8, + lanes: 1u16, + }, + DataType::FP32 => DLDataType { + code: DLDataTypeCode_kDLFloat as u8, + bits: 32u8, + lanes: 1u16, + }, + } + } + + /// Returns whether this `DataType` represents primitive type `T`. + pub fn is_type(&self) -> bool { + let typ = TypeId::of::(); + typ == TypeId::of::() || typ == TypeId::of::() || typ == TypeId::of::() + } +} + +impl From for DataType { + fn from(dl_dtype: DLDataType) -> Self { + if dl_dtype.code == DLDataTypeCode_kDLInt as u8 && dl_dtype.bits == 32u8 { + DataType::INT32 + } else if dl_dtype.code == DLDataTypeCode_kDLInt as u8 && dl_dtype.bits == 8u8 { + DataType::INT8 + } else if dl_dtype.code == DLDataTypeCode_kDLFloat as u8 && dl_dtype.bits == 32u8 { + DataType::FP32 + } else { + DataType::FP32 + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Tensor { + pub(crate) dtype: DataType, + pub(crate) shape: Vec, + pub(crate) strides: Option>, + pub(crate) data: Vec, +} + +#[allow(dead_code)] +impl Tensor { + pub fn new(dtype: DataType, shape: Vec, strides: Vec, data: Vec) -> Self { + Tensor { + dtype: dtype, + shape: shape, + strides: Some(strides), + data: data, + } + } + + pub fn dtype(&self) -> DataType { + self.dtype.clone() + } + + pub fn ndim(&self) -> usize { + self.shape.len() + } + + pub fn shape(&self) -> Vec { + self.shape.clone() + } + + pub fn data(&self) -> Vec { + self.data.clone() + } + + pub fn as_dltensor(&self) -> DLTensor { + DLTensor { + data: self.data.as_ptr() as *mut c_void, + ctx: DLContext { + device_type: DLDeviceType_kDLCPU, + device_id: 0 as c_int, + }, + ndim: self.shape.len() as c_int, + dtype: self.dtype().as_dldtype(), + shape: self.shape.as_ptr() as *mut i64, + strides: self.strides.as_ref().unwrap().as_ptr() as *mut i64, + byte_offset: 0, + ..Default::default() + } + } + + /// Returns the data of this `Tensor` as a `Vec`. + /// + /// # Panics + /// + /// Panics if the `Tensor` does not contain elements of type `T`. + pub fn to_vec(&self) -> Vec { + assert!(self.dtype().is_type::()); + + unsafe { + slice::from_raw_parts( + self.data().as_ptr() as *const T, + self.shape().iter().map(|v| *v as usize).product::() as usize, + ) + .to_vec() + } + } +} + +impl Default for Tensor { + fn default() -> Self { + Self { + dtype: DataType::FP32, + shape: Vec::new(), + strides: None, + data: Vec::new(), + } + } +} + +impl From for Tensor { + fn from(dlt: DLTensor) -> Self { + unsafe { + let shape = slice::from_raw_parts_mut(dlt.shape, dlt.ndim as usize).to_vec(); + let size = shape.iter().map(|v| *v as usize).product::() as usize; + let itemsize: usize = (dlt.dtype.bits >> 3).into(); + let data = slice::from_raw_parts(dlt.data as *const u8, size * itemsize).to_vec(); + + Self { + dtype: DataType::from(dlt.dtype), + shape, + strides: if dlt.strides.is_null() { + None + } else { + Some( + slice::from_raw_parts_mut(dlt.strides as *mut usize, dlt.ndim as usize) + .to_vec(), + ) + }, + data, + } + } + } +} diff --git a/apps/wasm-dlbackend-tvm/src/utils.rs b/apps/wasm-dlbackend-tvm/src/utils.rs new file mode 100644 index 0000000000000..05050981f4760 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/src/utils.rs @@ -0,0 +1,29 @@ +use super::ops::types::*; +use serde_json; +use std::ptr; + +pub fn load_inputs(in_addr: i32, in_size: usize) -> Vec { + let in_addr = in_addr as *mut u8; + + let mut data_vec = Vec::new(); + for i in 0..in_size { + data_vec.push(unsafe { ptr::read(in_addr.offset(i as isize)) }); + } + let inputs: Vec = serde_json::from_slice(&data_vec).unwrap(); + + inputs +} + +pub fn store_output(out_addr: i32, output: Tensor) -> usize { + let out_addr = out_addr as *mut u8; + + let data_vec = serde_json::to_vec(&output).unwrap(); + let data_size = data_vec.len(); + for i in 0..data_size { + unsafe { + ptr::write(out_addr.offset(i as isize), *data_vec.get(i).unwrap()); + } + } + + data_size +} diff --git a/apps/wasm-dlbackend-tvm/tools/build_ops_lib.py b/apps/wasm-dlbackend-tvm/tools/build_ops_lib.py new file mode 100644 index 0000000000000..e7a85ea1adc73 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/tools/build_ops_lib.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +"""Prepares a simple TVM library for operators.""" + +from os import path as osp +import sys + +import tvm +from tvm import te + + +def add(): + n = te.var('n') + A = te.placeholder((n,), name='A') + B = te.placeholder((n,), name='B') + C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') + s = tvm.te.create_schedule(C.op) + s[C].parallel(s[C].op.axis[0]) + m = tvm.lower(s, [A, B, C], name="add", simple_mode=True) + tvm.build(m, target="llvm -target=wasm32-unknown-unknown --system-lib").save(osp.join(sys.argv[1], 'add.o')) + + +def sub(): + n = te.var('n') + A = te.placeholder((n,), name='A') + B = te.placeholder((n,), name='B') + C = te.compute(A.shape, lambda *i: A(*i) - B(*i), name='C') + s = tvm.te.create_schedule(C.op) + s[C].parallel(s[C].op.axis[0]) + m = tvm.lower(s, [A, B, C], name="sub", simple_mode=True) + tvm.build(m, target="llvm -target=wasm32-unknown-unknown --system-lib").save(osp.join(sys.argv[1], 'sub.o')) + + +if __name__ == '__main__': + add() + sub() diff --git a/apps/wasm-dlbackend-tvm/wasm-dlfrontend/Cargo.toml b/apps/wasm-dlbackend-tvm/wasm-dlfrontend/Cargo.toml new file mode 100644 index 0000000000000..e5174f91e065c --- /dev/null +++ b/apps/wasm-dlbackend-tvm/wasm-dlfrontend/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "wasm-dlfrontend" +version = "0.1.0" +authors = ["leonwanghui "] +edition = "2018" + +[dependencies] +wasmtime = "0.17.0" +wasmtime-wasi = "0.17.0" +anyhow = "1.0.30" +serde = "1.0.53" +serde_json = "1.0.53" +serde_derive = "1.0.53" +getopts = "0.2.21" +ndarray = "0.12" diff --git a/apps/wasm-dlbackend-tvm/wasm-dlfrontend/src/main.rs b/apps/wasm-dlbackend-tvm/wasm-dlfrontend/src/main.rs new file mode 100644 index 0000000000000..0c4cc5f4f9564 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/wasm-dlfrontend/src/main.rs @@ -0,0 +1,132 @@ +#[macro_use] +extern crate serde_derive; + +pub mod types; +use types::*; + +use anyhow::Result; +use getopts::Options; +use ndarray::Array; +use serde_json; +use std::env; +use wasmtime::*; +use wasmtime_wasi::{Wasi, WasiCtx}; + +fn print_usage(program: &str, opts: Options) { + let brief = format!("Usage: {} [options]", program); + print!("{}", opts.usage(&brief)); +} + +fn main() { + let args: Vec = env::args().collect(); + let program = args[0].clone(); + + let mut opts = Options::new(); + opts.optopt( + "c", + "ms-backend-config", + "set wasm backend config file", + "FILE_PATH", + ); + opts.optopt( + "o", + "op-type", + "set the operator type, currently ONLY support Add and Sub, default: Add.", + "VALUE", + ); + + opts.optflag("h", "help", "print this help menu"); + let matches = match opts.parse(&args[1..]) { + Ok(m) => m, + Err(f) => panic!(f.to_string()), + }; + if matches.opt_present("h") { + print_usage(&program, opts); + return; + } + let wasm_backend_file: String = match matches.opt_str("c") { + Some(s) => s, + None => String::from("/opt/ms-backend-wasm/wasm_backend_tvm.wasi.wasm"), + }; + let op_type_str: String = match matches.opt_str("o") { + Some(s) => s, + None => String::from("Add"), + }; + let op_type: i32 = match op_type_str.as_str() { + "Add" => 0, + "Sub" => 1, + _ => 0, + }; + let a = Array::from_vec(vec![1f32, 2., 3., 4.]); + let b = Array::from_vec(vec![1f32; 4]); + let c = Array::from_vec(vec![0f32; 4]); + let a_tensor: Tensor = a.into(); + let b_tensor: Tensor = b.into(); + let c_tensor: Tensor = c.into(); + + let result: Tensor = match execute( + wasm_backend_file, + op_type, + vec![a_tensor, b_tensor, c_tensor], + ) { + Ok(m) => m, + Err(f) => panic!(f.to_string()), + }; + println!( + "{}", + serde_json::to_string_pretty(&result.to_vec::()).unwrap() + ); +} + +fn execute(wasm_backend_file: String, op_type: i32, input_data: Vec) -> Result { + let store = Store::default(); + + // First set up our linker which is going to be linking modules together. We + // want our linker to have wasi available, so we set that up here as well. + let mut linker = Linker::new(&store); + // Create an instance of `Wasi` which contains a `WasiCtx`. Note that + // `WasiCtx` provides a number of ways to configure what the target program + // will have access to. + let wasi = Wasi::new(&store, WasiCtx::new(std::env::args())?); + wasi.add_to_linker(&mut linker)?; + + let module = Module::from_file(store.engine(), &wasm_backend_file)?; + let instance = linker.instantiate(&module)?; + let memory = instance + .get_memory("memory") + .ok_or(anyhow::format_err!("failed to find `memory` export"))?; + + // Specify the input address and output address to access the wasm memory. + let in_addr = 0x1000; + let out_addr = 0x2000; + + // Serialize the data into a JSON string. + let in_data = serde_json::to_vec(&input_data)?; + let in_size = in_data.len(); + // Insert the input data into wasm memory. + for i in 0..in_size { + unsafe { + memory.data_unchecked_mut()[in_addr + i] = *in_data.get(i).unwrap(); + } + } + + // Invoke `run` export + let run = instance + .get_func("run") + .ok_or(anyhow::format_err!("failed to find `run` function export!"))? + .get4::()?; + + let out_size = run( + op_type.clone() as i32, + in_addr as i32, + in_size as i32, + out_addr as i32, + )?; + if out_size == 0 { + panic!("Opeartor {:?} run failed!", op_type); + } + + let out_data = unsafe { &memory.data_unchecked()[out_addr..][..out_size as usize] }; + let out_vec: Tensor = serde_json::from_slice(out_data).unwrap(); + Ok(out_vec.clone()) +} diff --git a/apps/wasm-dlbackend-tvm/wasm-dlfrontend/src/types.rs b/apps/wasm-dlbackend-tvm/wasm-dlfrontend/src/types.rs new file mode 100644 index 0000000000000..17168e9da4e55 --- /dev/null +++ b/apps/wasm-dlbackend-tvm/wasm-dlfrontend/src/types.rs @@ -0,0 +1,107 @@ +use std::{any::TypeId, convert::From, mem, slice}; + +#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)] +pub enum DataType { + FP32, + INT32, + INT8, +} + +impl DataType { + /// Returns whether this `DataType` represents primitive type `T`. + pub fn is_type(&self) -> bool { + let typ = TypeId::of::(); + typ == TypeId::of::() || typ == TypeId::of::() || typ == TypeId::of::() + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Tensor { + pub(crate) dtype: DataType, + pub(crate) shape: Vec, + pub(crate) strides: Option>, + pub(crate) data: Vec, +} + +#[allow(dead_code)] +impl Tensor { + pub fn new(dtype: DataType, shape: Vec, strides: Vec, data: Vec) -> Self { + Tensor { + dtype: dtype, + shape: shape, + strides: Some(strides), + data: data, + } + } + + pub fn dtype(&self) -> DataType { + self.dtype.clone() + } + + pub fn ndim(&self) -> usize { + self.shape.len() + } + + pub fn shape(&self) -> Vec { + self.shape.clone() + } + + pub fn data(&self) -> Vec { + self.data.clone() + } + + /// Returns the data of this `Tensor` as a `Vec`. + /// + /// # Panics + /// + /// Panics if the `Tensor` does not contain elements of type `T`. + pub fn to_vec(&self) -> Vec { + assert!(self.dtype().is_type::()); + + unsafe { + slice::from_raw_parts( + self.data().as_ptr() as *const T, + self.shape().iter().map(|v| *v as usize).product::() as usize, + ) + .to_vec() + } + } +} + +impl Default for Tensor { + fn default() -> Self { + Self { + dtype: DataType::FP32, + shape: Vec::new(), + strides: None, + data: Vec::new(), + } + } +} + +/// `From` conversions to `Tensor` for `ndarray::Array`. +/// Takes a reference to the `ndarray` since `Tensor` is not owned. +macro_rules! impl_tensor_from_ndarray { + ($type:ty, $typecode:expr) => { + impl From> for Tensor { + fn from(arr: ndarray::Array<$type, D>) -> Self { + Tensor { + dtype: $typecode, + shape: arr.shape().iter().map(|v| *v as i64).collect(), + strides: Some(arr.strides().iter().map(|v| *v as usize).collect()), + data: unsafe { + slice::from_raw_parts( + arr.as_ptr() as *const u8, + arr.len() * mem::size_of::<$type>(), + ) + .to_vec() + }, + } + } + } + }; +} + +impl_tensor_from_ndarray!(f32, DataType::FP32); +impl_tensor_from_ndarray!(i32, DataType::INT32); +impl_tensor_from_ndarray!(i8, DataType::INT8);