Repository: apache/tvm

Issue:
[RUST] RFC and 0.5 Release Plans
The following are the current plans for [TVM Rust](https://github.com/dmlc/tvm/issues/1601) including both "runtime" and "frontend" as part of roadmap v0.5.

- [x] Create `common` crate for runtime and frontend
- [x] [Merge frontend](https://github.com/dmlc/tvm/pull/2292)
- [x] Comply with Rust [2018-edition](https://rust-lang-nursery.github.io/edition-guide/rust-2018/index.html).
- [x] Add function builder to runtime and make runtime `PackedFunc` compatible with frontend
- [x] Make `call_packed!` macro more general and compatible across both runtime and frontend sides as [discussed](https://github.com/dmlc/tvm/issues/2305).


Discussion:
Here're my thoughts about a `common` crate:

The main difference (besides some conversions and using a fat point) between frontend and runtime `TVMArgValue` and `TVMRetValue` is how they're wrapping/using `TVMValue`. The runtime `TVMValue`  is raw, however, frontend needs `Debug + Display` and some idiomatic conversions. Therefore, `common`'s `TVMValue` has to be wrapped *at least*. 

Because bindgen doesn't work for runtime, so either we should 
* include the generated raw c api in `common` entirely (same as runtime) 
or 
* include what needed such as raw `TVMValue` and `TVMTypeCode` and then, wrap and impl required conversions for *both* runtime and frontend (because of Orphan rule so that we won't need more wrapper boilerplates in either crates).
@nhynes I'm working on the common crate and it's completely non-obvious how to make it. I'm making minimum assumptions like even removing the debugging for `TVMValue` and including it raw same as in Rust runtime. The only idea that might lead to a solution is to include the entire `c_runtime_api.rs` in common and get rid of my `tvm-sys` which I don't like and it would impose very uncommon layout for a Rust binding. If I include some partial C runtime api such as `TVMValue` and `TVMTypeCode`, then Rust's orphan rule prevents me from conversion from `common::TVMValue` to `tvm_sys::TVMValue` and need to change all occurences of `tvm_sys::TVMValue` in the raw api to `common::TVMValue` which is not possible!

I'm now thinking maybe common crate is not a good idea here and we can have exactly the same functionalities with two different implementations because runtime and frontend needs are different when we include debugging for example. Note that this includes very basic value conversions API only and it's not significant.

Overall, I'll work on making a common crate this week and next maybe and if I wouldn't have found a solution by then, I'd just change the frontend `TVMRetValue` to exactly match the runtime and impl the same conversions and send the PR. 

Any comment?
an interesting point. what happens if the `common` crate has the headers behind one of two mutually exclusive feature flags: `runtime` and `frontend`? The former would pull in the definitions from `c_runtime_api.rs` and the latter would expose `tvm-sys`.

> removing the debugging 

I'm not sure I follow. Is it not as easy as `#[derive(Debug)]` or even `impl Debug for Struct`

Like even if we have to do out-of-band codegen, I think that it's worth having one API. If we don't, SWIM will surely come by in a few months, read the codebase, and think "wtf. why are there two different implementations of the same thing?!?!" much as I do when reading topi code :P 

I'd prefer for you to not struggle with this, so if it really becomes too onerous, just give me your acceptance criteria for the frontend "working" and I'll do my best to merge them.
> what happens if the common crate has the headers behind one of two mutually exclusive feature flags: runtime and frontend?

Well, I haven't thought about that. It seems a viable approach. 

> I'm not sure I follow. Is it not as easy as #[derive(Debug)] or even impl Debug for Struct

I take back my debugging concern. Previously, I sort of made `TVMValue` into `PartialEq + Eq` by attaching a type but is not necessary and it should be fine to just `impl Debug`.

> I'd  prefer for you to not struggle with this, so if it really becomes too  onerous, just give me your acceptance criteria for the frontend  "working" and I'll do my best to merge them.

Thanks for the support :+1: 

Right now, I'm using runtime impl in `common` with some changes in runtime. I'll update you for the frontend compatibility as it implies some good number of changes as well. Hopefully it won't be too much headache!

I should say that inevitably the downside of `common` is exposing all fields of `TVMArgValue` (include `_lifetime`) and `TVMRetValue` which might be ok for our case but need more care perhaps!

Resolving patch:
diff --git a/rust/frontend/examples/resnet/src/build_resnet.py b/rust/frontend/examples/resnet/src/build_resnet.py
new file mode 100755
--- /dev/null
+++ b/rust/frontend/examples/resnet/src/build_resnet.py
@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+
+import argparse
+import csv
+import logging
+from os import path as osp
+import sys
+
+import numpy as np
+
+import mxnet as mx
+from mxnet.gluon.model_zoo.vision import get_model
+from mxnet.gluon.utils import download
+
+import tvm
+from tvm.contrib import graph_runtime, cc
+import nnvm
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+parser = argparse.ArgumentParser(description='Resnet build example')
+aa = parser.add_argument
+aa('--batch-size', type=int, default=1, help='input image batch size')
+aa('--opt-level', type=int, default=3,
+   help='level of optimization. 0 is unoptimized and 3 is the highest level')
+aa('--target', type=str, default='llvm', help='target context for compilation')
+aa('--image-shape', type=str, default='3,224,224', help='input image dimensions')
+aa('--image-name', type=str, default='cat.png', help='name of input image to download')
+args = parser.parse_args()
+
+target_dir = osp.dirname(osp.dirname(osp.realpath(__file__)))
+batch_size = args.batch_size
+opt_level = args.opt_level
+target = tvm.target.create(args.target)
+image_shape = tuple(map(int, args.image_shape.split(",")))
+data_shape = (batch_size,) + image_shape
+
+def build(target_dir):
+    """ Compiles resnet18 with TVM"""
+    deploy_lib = osp.join(target_dir, 'deploy_lib.o')
+    if osp.exists(deploy_lib):
+        return
+    # download the pretrained resnet18 trained on imagenet1k dataset for
+    # image classification task
+    block = get_model('resnet18_v1', pretrained=True)
+
+    sym, params = nnvm.frontend.from_mxnet(block)
+    # add the softmax layer for prediction
+    net = nnvm.sym.softmax(sym)
+    # compile the model
+    with nnvm.compiler.build_config(opt_level=opt_level):
+        graph, lib, params = nnvm.compiler.build(
+            net, target, shape={"data": data_shape}, params=params)
+    # save the model artifacts
+    lib.save(deploy_lib)
+    cc.create_shared(osp.join(target_dir, "deploy_lib.so"),
+                    [osp.join(target_dir, "deploy_lib.o")])
+
+    with open(osp.join(target_dir, "deploy_graph.json"), "w") as fo:
+        fo.write(graph.json())
+
+    with open(osp.join(target_dir,"deploy_param.params"), "wb") as fo:
+        fo.write(nnvm.compiler.save_param_dict(params))
+
+def download_img_labels():
+    """ Download an image and imagenet1k class labels for test"""
+    img_name = 'cat.png'
+    synset_url = ''.join(['https://gist.githubusercontent.com/zhreshold/',
+                      '4d0b62f3d01426887599d4f7ede23ee5/raw/',
+                      '596b27d23537e5a1b5751d2b0481ef172f58b539/',
+                      'imagenet1000_clsid_to_human.txt'])
+    synset_name = 'synset.txt'
+    download('https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true', img_name)
+    download(synset_url, synset_name)
+
+    with open(synset_name) as fin:
+        synset = eval(fin.read())
+
+    with open("synset.csv", "w") as fout:
+        w = csv.writer(fout)
+        w.writerows(synset.items())
+
+def test_build(target_dir):
+    """ Sanity check with random input"""
+    graph = open(osp.join(target_dir, "deploy_graph.json")).read()
+    lib = tvm.module.load(osp.join(target_dir, "deploy_lib.so"))
+    params = bytearray(open(osp.join(target_dir,"deploy_param.params"), "rb").read())
+    input_data = tvm.nd.array(np.random.uniform(size=data_shape).astype("float32"))
+    ctx = tvm.cpu()
+    module = graph_runtime.create(graph, lib, ctx)
+    module.load_params(params)
+    module.run(data=input_data)
+    out = module.get_output(0).asnumpy()
+
+
+if __name__ == '__main__':
+    logger.info("building the model")
+    build(target_dir)
+    logger.info("build was successful")
+    logger.info("test the build artifacts")
+    test_build(target_dir)
+    logger.info("test was successful")
+    download_img_labels()
+    logger.info("image and synset downloads are successful")


Test patch:
diff --git a/rust/frontend/tests/basics/.gitignore b/rust/frontend/tests/basics/.gitignore
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/basics/.gitignore
@@ -0,0 +1,7 @@
+/target
+**/*.rs.bk
+Cargo.lock
+*.o
+*.so
+*.ptx
+*.json
diff --git a/rust/frontend/tests/basics/Cargo.toml b/rust/frontend/tests/basics/Cargo.toml
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/basics/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+name = "basics"
+version = "0.0.0"
+authors = ["TVM Contributors"]
+license = "Apache-2.0"
+build = "build.rs"
+
+[dependencies]
+ndarray = "0.12.1"
+tvm-frontend = { path = "../../" }
+
+[features]
+default = ["cpu"]
+cpu = []
+gpu = []
diff --git a/rust/frontend/tests/basics/build.rs b/rust/frontend/tests/basics/build.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/basics/build.rs
@@ -0,0 +1,27 @@
+fn main() {
+    let out_dir = std::env::var("OUT_DIR").unwrap();
+
+    let output = std::process::Command::new(concat!(env!("CARGO_MANIFEST_DIR"), "/src/tvm_add.py"))
+        .args(&[
+            if cfg!(feature = "cpu") {
+                "llvm"
+            } else {
+                "cuda"
+            },
+            &std::env::var("OUT_DIR").unwrap(),
+        ])
+        .output()
+        .expect("Failed to execute command");
+    assert!(
+        std::path::Path::new(&format!("{}/test_add.so", out_dir)).exists(),
+        "Could not build tvm lib: {}",
+        String::from_utf8(output.stderr)
+            .unwrap()
+            .trim()
+            .split("\n")
+            .last()
+            .unwrap_or("")
+    );
+
+    println!("cargo:rustc-link-search=native={}", out_dir);
+}
diff --git a/rust/frontend/tests/basics/src/main.rs b/rust/frontend/tests/basics/src/main.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/basics/src/main.rs
@@ -0,0 +1,35 @@
+extern crate ndarray as rust_ndarray;
+extern crate tvm_frontend as tvm;
+
+use tvm::*;
+
+fn main() {
+    let shape = &mut [2];
+    let mut data = vec![3f32, 4.0];
+
+    let (ctx, ctx_name) = if cfg!(feature = "cpu") {
+        (TVMContext::cpu(0), "cpu")
+    } else {
+        (TVMContext::gpu(0), "gpu")
+    };
+    let dtype = TVMType::from("float32");
+    let mut arr = NDArray::empty(shape, ctx, dtype);
+    arr.copy_from_buffer(data.as_mut_slice());
+    let mut ret = NDArray::empty(shape, ctx, dtype);
+    let mut fadd = Module::load(&concat!(env!("OUT_DIR"), "/test_add.so")).unwrap();
+    if !fadd.enabled(ctx_name) {
+        return;
+    }
+    if cfg!(feature = "gpu") {
+        fadd.import_module(Module::load(&concat!(env!("OUT_DIR"), "/test_add.ptx")).unwrap());
+    }
+    function::Builder::from(&mut fadd)
+        .arg(&arr)
+        .arg(&arr)
+        .set_output(&mut ret)
+        .unwrap()
+        .invoke()
+        .unwrap();
+
+    assert_eq!(ret.to_vec::<f32>().unwrap(), vec![6f32, 8.0]);
+}
diff --git a/rust/frontend/tests/basics/src/tvm_add.py b/rust/frontend/tests/basics/src/tvm_add.py
new file mode 100755
--- /dev/null
+++ b/rust/frontend/tests/basics/src/tvm_add.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+
+import os.path as osp
+import sys
+
+import tvm
+from tvm.contrib import cc
+
+
+def main(target, out_dir):
+    n = tvm.var('n')
+    A = tvm.placeholder((n,), name='A')
+    B = tvm.placeholder((n,), name='B')
+    C = tvm.compute(A.shape, lambda i: A[i] + B[i], name='C')
+    s = tvm.create_schedule(C.op)
+
+    if target == 'cuda':
+        bx, tx = s[C].split(C.op.axis[0], factor=64)
+        s[C].bind(bx, tvm.thread_axis('blockIdx.x'))
+        s[C].bind(tx, tvm.thread_axis('threadIdx.x'))
+
+    fadd = tvm.build(s, [A, B, C], target, target_host='llvm', name='myadd')
+
+    fadd.save(osp.join(out_dir, 'test_add.o'))
+    if target == 'cuda':
+        fadd.imported_modules[0].save(os.path.join(out_dir, 'test_add.ptx'))
+    cc.create_shared(
+        osp.join(out_dir, 'test_add.so'), [osp.join(out_dir, 'test_add.o')])
+
+
+if __name__ == '__main__':
+    main(sys.argv[1], sys.argv[2])
+
diff --git a/rust/frontend/tests/callback/Cargo.toml b/rust/frontend/tests/callback/Cargo.toml
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/callback/Cargo.toml
@@ -0,0 +1,8 @@
+[package]
+name = "callback"
+version = "0.0.0"
+authors = ["TVM Contributors"]
+
+[dependencies]
+ndarray = "0.12.1"
+tvm-frontend = { path = "../../" }
diff --git a/rust/frontend/tests/callback/src/bin/array.rs b/rust/frontend/tests/callback/src/bin/array.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/callback/src/bin/array.rs
@@ -0,0 +1,44 @@
+#![feature(extern_crate_item_prelude, try_from)]
+#![allow(unused_imports)]
+
+extern crate ndarray as rust_ndarray;
+#[macro_use]
+extern crate tvm_frontend as tvm;
+
+use rust_ndarray::ArrayD;
+use std::convert::{TryFrom, TryInto};
+
+use tvm::*;
+
+fn main() {
+    register_global_func! {
+        fn sum(args: &[TVMArgValue]) -> Result<TVMRetValue> {
+            let mut ret = 0f32;
+            let shape = &mut [2];
+            for arg in args.iter() {
+                let e = NDArray::empty(shape, TVMContext::cpu(0), TVMType::from("float32"));
+                let arg: NDArray = arg.try_into()?;
+                let arr = arg.copy_to_ndarray(e)?;
+                let rnd: ArrayD<f32> = ArrayD::try_from(&arr)?;
+                ret += rnd.scalar_sum();
+            }
+            Ok(TVMRetValue::from(ret))
+        }
+    }
+
+    let shape = &mut [2];
+    let mut data = vec![3f32, 4.0];
+    let mut arr = NDArray::empty(shape, TVMContext::cpu(0), TVMType::from("float32"));
+    arr.copy_from_buffer(data.as_mut_slice());
+
+    let mut registered = function::Builder::default();
+    let ret: f32 = registered
+        .get_function("sum", true)
+        .arg(&arr)
+        .arg(&arr)
+        .invoke()
+        .unwrap()
+        .try_into()
+        .unwrap();
+    assert_eq!(ret, 14f32);
+}
diff --git a/rust/frontend/tests/callback/src/bin/error.rs b/rust/frontend/tests/callback/src/bin/error.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/callback/src/bin/error.rs
@@ -0,0 +1,43 @@
+#![feature(extern_crate_item_prelude, panic_info_message)]
+#![allow(unused_imports)]
+
+use std::panic;
+
+#[macro_use]
+extern crate tvm_frontend as tvm;
+
+use tvm::*;
+
+fn main() {
+    register_global_func! {
+        fn error(_args: &[TVMArgValue]) -> Result<TVMRetValue> {
+            Err(ErrorKind::TypeMismatch(
+                format!("{}", "i64".to_string()),
+                format!("{}", "f64".to_string()),
+            ).into())
+        }
+    }
+
+    let mut registered = function::Builder::default();
+    registered.get_function("error", true);
+    assert!(registered.func.is_some());
+    registered.args(&[10, 20]);
+
+    println!("expected error message is:");
+    panic::set_hook(Box::new(|panic_info| {
+        if let Some(msg) = panic_info.message() {
+            println!("{:?}", msg);
+        }
+        if let Some(location) = panic_info.location() {
+            println!(
+                "panic occurred in file '{}' at line {}",
+                location.file(),
+                location.line()
+            );
+        } else {
+            println!("panic occurred but can't get location information");
+        }
+    }));
+
+    let _result = registered.invoke();
+}
diff --git a/rust/frontend/tests/callback/src/bin/float.rs b/rust/frontend/tests/callback/src/bin/float.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/callback/src/bin/float.rs
@@ -0,0 +1,32 @@
+#![feature(extern_crate_item_prelude, try_from)]
+#![allow(unused_imports)]
+
+#[macro_use]
+extern crate tvm_frontend as tvm;
+
+use std::convert::TryInto;
+use tvm::*;
+
+fn main() {
+    register_global_func! {
+        fn sum(args: &[TVMArgValue]) -> Result<TVMRetValue> {
+            let mut ret = 0.0;
+            for arg in args.iter() {
+                let val: f64 = arg.try_into()?;
+                ret += val;
+            }
+            Ok(TVMRetValue::from(&ret))
+        }
+    }
+
+    let mut registered = function::Builder::default();
+    registered.get_function("sum", true);
+    assert!(registered.func.is_some());
+    let ret: f64 = registered
+        .args(&[10.0f64, 20.0, 30.0])
+        .invoke()
+        .unwrap()
+        .try_into()
+        .unwrap();
+    assert_eq!(ret, 60f64);
+}
diff --git a/rust/frontend/tests/callback/src/bin/int.rs b/rust/frontend/tests/callback/src/bin/int.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/callback/src/bin/int.rs
@@ -0,0 +1,31 @@
+#![feature(extern_crate_item_prelude, try_from)]
+#![allow(unused_imports)]
+
+extern crate tvm_frontend as tvm;
+
+use std::convert::TryInto;
+use tvm::*;
+
+fn main() {
+    fn sum(args: &[TVMArgValue]) -> Result<TVMRetValue> {
+        let mut ret = 0i64;
+        for arg in args.iter() {
+            let val: i64 = arg.try_into()?;
+            ret += val;
+        }
+        Ok(TVMRetValue::from(&ret))
+    }
+
+    tvm::function::register(sum, "mysum".to_owned(), false).unwrap();
+
+    let mut registered = function::Builder::default();
+    registered.get_function("mysum", true);
+    assert!(registered.func.is_some());
+    let ret: i64 = registered
+        .args(&[10, 20, 30])
+        .invoke()
+        .unwrap()
+        .try_into()
+        .unwrap();
+    assert_eq!(ret, 60);
+}
diff --git a/rust/frontend/tests/callback/src/bin/string.rs b/rust/frontend/tests/callback/src/bin/string.rs
new file mode 100644
--- /dev/null
+++ b/rust/frontend/tests/callback/src/bin/string.rs
@@ -0,0 +1,34 @@
+#![feature(extern_crate_item_prelude, try_from)]
+#![allow(unused_imports)]
+
+#[macro_use]
+extern crate tvm_frontend as tvm;
+use std::convert::TryInto;
+use tvm::*;
+
+// FIXME
+fn main() {
+    register_global_func! {
+        fn concate_str(args: &[TVMArgValue]) -> Result<TVMRetValue> {
+            let mut ret = "".to_string();
+            for arg in args.iter() {
+                let val: String = arg.try_into()?;
+                ret += val.as_str();
+            }
+            Ok(TVMRetValue::from(ret))
+        }
+    }
+    let mut registered = function::Builder::default();
+    registered.get_function("concate_str", true);
+    assert!(registered.func.is_some());
+    let a = "a".to_string();
+    let b = "b".to_string();
+    let c = "c".to_string();
+    let ret: String = registered
+        .args(&[a, b, c])
+        .invoke()
+        .unwrap()
+        .try_into()
+        .unwrap();
+    assert_eq!(ret, "abc".to_owned());
+}
diff --git a/rust/tests/.gitignore b/rust/runtime/tests/.gitignore
similarity index 100%
rename from rust/tests/.gitignore
rename to rust/runtime/tests/.gitignore
diff --git a/rust/tests/build_model.py b/rust/runtime/tests/build_model.py
old mode 100644
new mode 100755
similarity index 98%
rename from rust/tests/build_model.py
rename to rust/runtime/tests/build_model.py
--- a/rust/tests/build_model.py
+++ b/rust/runtime/tests/build_model.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """Builds a simple NNVM graph for testing."""
 
 from os import path as osp
diff --git a/rust/runtime/tests/test_graph_serde.rs b/rust/runtime/tests/test_graph_serde.rs
new file mode 100644
--- /dev/null
+++ b/rust/runtime/tests/test_graph_serde.rs
@@ -0,0 +1,39 @@
+#![feature(try_from)]
+
+extern crate serde;
+extern crate serde_json;
+
+extern crate tvm_runtime;
+
+use std::{convert::TryFrom, fs, io::Read};
+
+use tvm_runtime::Graph;
+
+#[test]
+fn test_load_graph() {
+    let mut params_bytes = Vec::new();
+    fs::File::open(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/graph.params"))
+        .expect("Could not find TVM graph. Did you run `tests/build_model.py`?")
+        .read_to_end(&mut params_bytes)
+        .unwrap();
+    let _params = tvm_runtime::load_param_dict(&params_bytes);
+
+    let graph = Graph::try_from(
+        &fs::read_to_string(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/graph.json")).unwrap(),
+    )
+    .unwrap();
+
+    assert_eq!(graph.nodes[3].op, "tvm_op");
+    assert_eq!(
+        graph.nodes[3]
+            .attrs
+            .as_ref()
+            .unwrap()
+            .get("func_name")
+            .unwrap(),
+        "fuse_dense"
+    );
+    assert_eq!(graph.nodes[5].inputs[0].index, 0);
+    assert_eq!(graph.nodes[6].inputs[0].index, 1);
+    assert_eq!(graph.heads.len(), 2);
+}
diff --git a/rust/tests/test_nnvm/Cargo.toml b/rust/runtime/tests/test_nnvm/Cargo.toml
similarity index 71%
rename from rust/tests/test_nnvm/Cargo.toml
rename to rust/runtime/tests/test_nnvm/Cargo.toml
--- a/rust/tests/test_nnvm/Cargo.toml
+++ b/rust/runtime/tests/test_nnvm/Cargo.toml
@@ -2,13 +2,13 @@
 name = "test-nnvm"
 version = "0.0.0"
 license = "Apache-2.0"
-authors = ["Nick Hynes <nhynes@berkeley.edu>"]
+authors = ["TVM Contributors"]
 
 [dependencies]
 ndarray = "0.11.2"
-tvm = { path = "../../" }
 serde = "1.0.59"
 serde_json = "1.0.17"
+tvm-runtime = { path = "../../" }
 
 [build-dependencies]
 ar = "0.6.0"
diff --git a/rust/runtime/tests/test_nnvm/build.rs b/rust/runtime/tests/test_nnvm/build.rs
new file mode 100644
--- /dev/null
+++ b/rust/runtime/tests/test_nnvm/build.rs
@@ -0,0 +1,33 @@
+extern crate ar;
+
+use std::{env, fs::File, path::Path, process::Command};
+
+use ar::Builder;
+
+fn main() {
+    let out_dir = env::var("OUT_DIR").unwrap();
+
+    let output = Command::new(concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/src/build_test_graph.py"
+    ))
+    .arg(&out_dir)
+    .output()
+    .expect("Failed to execute command");
+    assert!(
+        Path::new(&format!("{}/graph.o", out_dir)).exists(),
+        "Could not build graph lib: {}",
+        String::from_utf8(output.stderr)
+            .unwrap()
+            .trim()
+            .split("\n")
+            .last()
+            .unwrap_or("")
+    );
+
+    let mut builder = Builder::new(File::create(format!("{}/libgraph.a", out_dir)).unwrap());
+    builder.append_path(format!("{}/graph.o", out_dir)).unwrap();
+
+    println!("cargo:rustc-link-lib=static=graph");
+    println!("cargo:rustc-link-search=native={}", out_dir);
+}
diff --git a/rust/tests/test_nnvm/src/build_test_graph.py b/rust/runtime/tests/test_nnvm/src/build_test_graph.py
similarity index 99%
rename from rust/tests/test_nnvm/src/build_test_graph.py
rename to rust/runtime/tests/test_nnvm/src/build_test_graph.py
--- a/rust/tests/test_nnvm/src/build_test_graph.py
+++ b/rust/runtime/tests/test_nnvm/src/build_test_graph.py
@@ -23,6 +23,7 @@ def _get_model(dshape):
 def _init_params(graph, input_shapes, initializer=init.Xavier(), seed=10):
     if isinstance(graph, sym.Symbol):
         graph = nnvm.graph.create(graph)
+
     ishapes, _ = graph_util.infer_shape(graph, **input_shapes)
     param_shapes = dict(zip(graph.index.input_names, ishapes))
     np.random.seed(seed)
@@ -40,6 +41,7 @@ def _init_params(graph, input_shapes, initializer=init.Xavier(), seed=10):
         initializer(param, init_value)
         # init_value /= init_value.sum() + 1e-10
         params[param] = tvm.nd.array(init_value)
+
     return params
 
 def main():
@@ -56,6 +58,7 @@ def main():
     lib.save(osp.join(sys.argv[1], 'graph.o'))
     with open(osp.join(out_dir, 'graph.json'), 'w') as f_resnet:
         f_resnet.write(graph.json())
+        
     with open(osp.join(out_dir, 'graph.params'), 'wb') as f_params:
         f_params.write(nnvm.compiler.save_param_dict(params))
 
diff --git a/rust/runtime/tests/test_nnvm/src/main.rs b/rust/runtime/tests/test_nnvm/src/main.rs
new file mode 100644
--- /dev/null
+++ b/rust/runtime/tests/test_nnvm/src/main.rs
@@ -0,0 +1,82 @@
+#![feature(try_from)]
+
+#[macro_use]
+extern crate ndarray;
+extern crate serde;
+extern crate serde_json;
+
+extern crate tvm_runtime;
+use std::{collections::HashMap, convert::TryFrom, fs, io::Read};
+
+use ndarray::Array;
+use tvm_runtime::{Graph, GraphExecutor, SystemLibModule, Tensor};
+
+const BATCH_SIZE: usize = 4;
+const IN_DIM: usize = 8;
+
+macro_rules! check_sum {
+    ($e:expr, $a:ident, $b:ident) => {
+        let a = Array::try_from($e.get_input(stringify!($a)).unwrap()).unwrap();
+        check_sum!(a, $b);
+    };
+    ($e:expr, $a:expr, $b:ident) => {
+        let a = Array::try_from($e.get_output($a).unwrap()).unwrap();
+        check_sum!(a, $b);
+    };
+    ($a:ident, $b:ident) => {
+        let a_sum: f32 = $a.scalar_sum();
+        let b_sum: f32 = $b.scalar_sum();
+        assert!((a_sum - b_sum).abs() < 1e-2, "{} != {}", a_sum, b_sum);
+    };
+}
+
+fn main() {
+    let syslib = SystemLibModule::default();
+
+    let mut params_bytes = Vec::new();
+    fs::File::open(concat!(env!("OUT_DIR"), "/graph.params"))
+        .unwrap()
+        .read_to_end(&mut params_bytes)
+        .unwrap();
+    let params = tvm_runtime::load_param_dict(&params_bytes)
+        .unwrap()
+        .into_iter()
+        .map(|(k, v)| (k, v.to_owned()))
+        .collect::<HashMap<String, Tensor<'static>>>();
+
+    let graph =
+        Graph::try_from(&fs::read_to_string(concat!(env!("OUT_DIR"), "/graph.json")).unwrap())
+            .unwrap();
+    let mut exec = GraphExecutor::new(graph, &syslib).unwrap();
+
+    let x = Array::from_shape_vec(
+        (BATCH_SIZE, IN_DIM),
+        (0..BATCH_SIZE * IN_DIM)
+            .map(|x| x as f32)
+            .collect::<Vec<f32>>(),
+    )
+    .unwrap();
+    let w = Array::try_from(params.get("dense0_weight").unwrap())
+        .unwrap()
+        .into_shape((IN_DIM * 2, IN_DIM))
+        .unwrap();
+    let b = Array::try_from(params.get("dense0_bias").unwrap()).unwrap();
+    let dense = x.dot(&w.t()) + &b;
+    let left = dense.slice(s![.., 0..IN_DIM]);
+    let right = dense.slice(s![.., IN_DIM..]);
+    let expected_o0 = &left + 1f32;
+    let expected_o1 = &right - 1f32;
+
+    exec.load_params(params);
+    exec.set_input("data", (&x).into());
+
+    check_sum!(exec, data, x);
+    check_sum!(exec, dense0_weight, w);
+    check_sum!(exec, dense0_bias, b);
+
+    exec.run();
+
+    check_sum!(exec, 0, expected_o0);
+    check_sum!(exec, 1, expected_o1);
+    check_sum!(exec, 2, dense);
+}
diff --git a/rust/tests/test_tvm_basic/Cargo.toml b/rust/runtime/tests/test_tvm_basic/Cargo.toml
similarity index 66%
rename from rust/tests/test_tvm_basic/Cargo.toml
rename to rust/runtime/tests/test_tvm_basic/Cargo.toml
--- a/rust/tests/test_tvm_basic/Cargo.toml
+++ b/rust/runtime/tests/test_tvm_basic/Cargo.toml
@@ -2,11 +2,11 @@
 name = "test-tvm-basic"
 version = "0.0.0"
 license = "Apache-2.0"
-authors = ["Nick Hynes <nhynes@berkeley.edu>"]
+authors = ["TVM Contributors"]
 
 [dependencies]
 ndarray = "0.11.2"
-tvm = { path = "../../" }
+tvm-runtime = { path = "../../" }
 
 [build-dependencies]
 ar = "0.6.0"
diff --git a/rust/runtime/tests/test_tvm_basic/build.rs b/rust/runtime/tests/test_tvm_basic/build.rs
new file mode 100644
--- /dev/null
+++ b/rust/runtime/tests/test_tvm_basic/build.rs
@@ -0,0 +1,34 @@
+extern crate ar;
+
+use std::{env, path::Path, process::Command};
+
+use ar::Builder;
+use std::fs::File;
+
+fn main() {
+    let out_dir = env::var("OUT_DIR").unwrap();
+
+    let output = Command::new(concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/src/build_test_lib.py"
+    ))
+    .arg(&out_dir)
+    .output()
+    .expect("Failed to execute command");
+    assert!(
+        Path::new(&format!("{}/test.o", out_dir)).exists(),
+        "Could not build tvm lib: {}",
+        String::from_utf8(output.stderr)
+            .unwrap()
+            .trim()
+            .split("\n")
+            .last()
+            .unwrap_or("")
+    );
+
+    let mut builder = Builder::new(File::create(format!("{}/libtest.a", out_dir)).unwrap());
+    builder.append_path(format!("{}/test.o", out_dir)).unwrap();
+
+    println!("cargo:rustc-link-lib=static=test");
+    println!("cargo:rustc-link-search=native={}", out_dir);
+}
diff --git a/rust/tests/test_tvm_basic/src/build_test_lib.py b/rust/runtime/tests/test_tvm_basic/src/build_test_lib.py
similarity index 100%
rename from rust/tests/test_tvm_basic/src/build_test_lib.py
rename to rust/runtime/tests/test_tvm_basic/src/build_test_lib.py
diff --git a/rust/runtime/tests/test_tvm_basic/src/main.rs b/rust/runtime/tests/test_tvm_basic/src/main.rs
new file mode 100644
--- /dev/null
+++ b/rust/runtime/tests/test_tvm_basic/src/main.rs
@@ -0,0 +1,22 @@
+extern crate ndarray;
+#[macro_use]
+extern crate tvm_runtime;
+
+use ndarray::Array;
+use tvm_runtime::{DLTensor, Module, SystemLibModule};
+
+fn main() {
+    let syslib = SystemLibModule::default();
+    let add = syslib
+        .get_function("default_function")
+        .expect("main function not found");
+    let mut a = Array::from_vec(vec![1f32, 2., 3., 4.]);
+    let mut b = Array::from_vec(vec![1f32, 0., 1., 0.]);
+    let mut c = Array::from_vec(vec![0f32; 4]);
+    let e = Array::from_vec(vec![2f32, 2., 4., 4.]);
+    let mut a_dl: DLTensor = (&mut a).into();
+    let mut b_dl: DLTensor = (&mut b).into();
+    let mut c_dl: DLTensor = (&mut c).into();
+    call_packed!(add, &mut a_dl, &mut b_dl, &mut c_dl);
+    assert!(c.all_close(&e, 1e-8f32));
+}
diff --git a/rust/tests/test_graph_serde.rs b/rust/tests/test_graph_serde.rs
deleted file mode 100644
--- a/rust/tests/test_graph_serde.rs
+++ /dev/null
@@ -1,39 +0,0 @@
-#![feature(try_from)]
-
-extern crate serde;
-extern crate serde_json;
-
-extern crate tvm;
-
-use std::{convert::TryFrom, fs, io::Read};
-
-use tvm::runtime::Graph;
-
-#[test]
-fn test_load_graph() {
-  let mut params_bytes = Vec::new();
-  fs::File::open(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/graph.params"))
-    .expect("Could not find TVM graph. Did you run `tests/build_model.py`?")
-    .read_to_end(&mut params_bytes)
-    .unwrap();
-  let _params = tvm::runtime::load_param_dict(&params_bytes);
-
-  let graph = Graph::try_from(
-    &fs::read_to_string(concat!(env!("CARGO_MANIFEST_DIR"), "/tests/graph.json")).unwrap(),
-  )
-  .unwrap();
-
-  assert_eq!(graph.nodes[3].op, "tvm_op");
-  assert_eq!(
-    graph.nodes[3]
-      .attrs
-      .as_ref()
-      .unwrap()
-      .get("func_name")
-      .unwrap(),
-    "fuse_dense"
-  );
-  assert_eq!(graph.nodes[5].inputs[0].index, 0);
-  assert_eq!(graph.nodes[6].inputs[0].index, 1);
-  assert_eq!(graph.heads.len(), 2);
-}
diff --git a/rust/tests/test_nnvm/build.rs b/rust/tests/test_nnvm/build.rs
deleted file mode 100644
--- a/rust/tests/test_nnvm/build.rs
+++ /dev/null
@@ -1,40 +0,0 @@
-extern crate ar;
-
-use std::{
-  env,
-  fs::File,
-  path::{Path, PathBuf},
-  process::Command,
-};
-
-use ar::Builder;
-
-fn main() {
-  let out_dir = env::var("OUT_DIR").unwrap();
-
-  let output = Command::new(concat!(
-    env!("CARGO_MANIFEST_DIR"),
-    "/src/build_test_graph.py"
-  ))
-  .arg(&out_dir)
-  .output()
-  .expect("Failed to execute command");
-  assert!(
-    Path::new(&format!("{}/graph.o", out_dir)).exists(),
-    "Could not build graph lib: {}",
-    String::from_utf8(output.stderr)
-      .unwrap()
-      .trim()
-      .split("\n")
-      .last()
-      .unwrap_or("")
-  );
-
-  let in_path: PathBuf = [&out_dir, "graph.o"].iter().collect();
-  let out_path: PathBuf = [&out_dir, "libgraph.a"].iter().collect();
-  let mut builder = Builder::new(File::create(out_path.to_str().unwrap()).unwrap());
-  builder.append_path(in_path.to_str().unwrap()).unwrap();
-
-  println!("cargo:rustc-link-lib=static=graph");
-  println!("cargo:rustc-link-search=native={}", out_dir);
-}
diff --git a/rust/tests/test_nnvm/src/main.rs b/rust/tests/test_nnvm/src/main.rs
deleted file mode 100644
--- a/rust/tests/test_nnvm/src/main.rs
+++ /dev/null
@@ -1,80 +0,0 @@
-#![feature(try_from)]
-
-#[macro_use]
-extern crate ndarray;
-extern crate serde;
-extern crate serde_json;
-
-extern crate tvm;
-use std::{collections::HashMap, convert::TryFrom, fs, io::Read};
-
-use ndarray::Array;
-use tvm::runtime::{Graph, GraphExecutor, SystemLibModule, Tensor};
-
-const BATCH_SIZE: usize = 4;
-const IN_DIM: usize = 8;
-
-macro_rules! check_sum {
-  ($e:expr, $a:ident, $b:ident) => {
-    let a = Array::try_from($e.get_input(stringify!($a)).unwrap()).unwrap();
-    check_sum!(a, $b);
-  };
-  ($e:expr, $a:expr, $b:ident) => {
-    let a = Array::try_from($e.get_output($a).unwrap()).unwrap();
-    check_sum!(a, $b);
-  };
-  ($a:ident, $b:ident) => {
-    let a_sum: f32 = $a.scalar_sum();
-    let b_sum: f32 = $b.scalar_sum();
-    assert!((a_sum - b_sum).abs() < 1e-2, "{} != {}", a_sum, b_sum);
-  };
-}
-
-fn main() {
-  let syslib = SystemLibModule::default();
-
-  let mut params_bytes = Vec::new();
-  fs::File::open(concat!(env!("OUT_DIR"), "/graph.params"))
-    .unwrap()
-    .read_to_end(&mut params_bytes)
-    .unwrap();
-  let params = tvm::runtime::load_param_dict(&params_bytes)
-    .unwrap()
-    .into_iter()
-    .map(|(k, v)| (k, v.to_owned()))
-    .collect::<HashMap<String, Tensor<'static>>>();
-
-  let graph =
-    Graph::try_from(&fs::read_to_string(concat!(env!("OUT_DIR"), "/graph.json")).unwrap()).unwrap();
-  let mut exec = GraphExecutor::new(graph, &syslib).unwrap();
-
-  let x = Array::from_shape_vec(
-    (BATCH_SIZE, IN_DIM),
-    (0..BATCH_SIZE * IN_DIM)
-      .map(|x| x as f32)
-      .collect::<Vec<f32>>(),
-  ).unwrap();
-  let w = Array::try_from(params.get("dense0_weight").unwrap())
-    .unwrap()
-    .into_shape((IN_DIM * 2, IN_DIM))
-    .unwrap();
-  let b = Array::try_from(params.get("dense0_bias").unwrap()).unwrap();
-  let dense = x.dot(&w.t()) + &b;
-  let left = dense.slice(s![.., 0..IN_DIM]);
-  let right = dense.slice(s![.., IN_DIM..]);
-  let expected_o0 = &left + 1f32;
-  let expected_o1 = &right - 1f32;
-
-  exec.load_params(params);
-  exec.set_input("data", x.clone().into());
-
-  check_sum!(exec, data, x);
-  check_sum!(exec, dense0_weight, w);
-  check_sum!(exec, dense0_bias, b);
-
-  exec.run();
-
-  check_sum!(exec, 0, expected_o0);
-  check_sum!(exec, 1, expected_o1);
-  check_sum!(exec, 2, dense);
-}
diff --git a/rust/tests/test_tvm_basic/build.rs b/rust/tests/test_tvm_basic/build.rs
deleted file mode 100644
--- a/rust/tests/test_tvm_basic/build.rs
+++ /dev/null
@@ -1,28 +0,0 @@
-extern crate ar;
-
-use std::{env, path::PathBuf, process::Command};
-
-use ar::Builder;
-use std::fs::File;
-
-fn main() {
-  let out_dir = env::var("OUT_DIR").unwrap();
-
-  let output = Command::new(concat!(
-    env!("CARGO_MANIFEST_DIR"),
-    "/src/build_test_lib.py"
-  )).arg(&out_dir)
-    .output()
-    .expect("Failed to execute command");
-  if output.stderr.len() > 0 {
-    panic!(String::from_utf8(output.stderr).unwrap());
-  }
-
-  let in_path: PathBuf = [&out_dir, "test.o"].iter().collect();
-  let out_path: PathBuf = [&out_dir, "libtest.a"].iter().collect();
-  let mut builder = Builder::new(File::create(out_path.to_str().unwrap()).unwrap());
-  builder.append_path(in_path.to_str().unwrap()).unwrap();
-
-  println!("cargo:rustc-link-lib=static=test");
-  println!("cargo:rustc-link-search=native={}", out_dir);
-}
diff --git a/rust/tests/test_tvm_basic/src/main.rs b/rust/tests/test_tvm_basic/src/main.rs
deleted file mode 100644
--- a/rust/tests/test_tvm_basic/src/main.rs
+++ /dev/null
@@ -1,25 +0,0 @@
-extern crate ndarray;
-#[macro_use]
-extern crate tvm;
-
-use ndarray::Array;
-use tvm::{
-  ffi::runtime::DLTensor,
-  runtime::{Module, SystemLibModule},
-};
-
-fn main() {
-  let syslib = SystemLibModule::default();
-  let add = syslib
-    .get_function("default_function")
-    .expect("main function not found");
-  let mut a = Array::from_vec(vec![1f32, 2., 3., 4.]);
-  let mut b = Array::from_vec(vec![1f32, 0., 1., 0.]);
-  let mut c = Array::from_vec(vec![0f32; 4]);
-  let e = Array::from_vec(vec![2f32, 2., 4., 4.]);
-  let mut a_dl: DLTensor = (&mut a).into();
-  let mut b_dl: DLTensor = (&mut b).into();
-  let mut c_dl: DLTensor = (&mut c).into();
-  call_packed!(add, &mut a_dl, &mut b_dl, &mut c_dl);
-  assert!(c.all_close(&e, 1e-8f32));
-}
diff --git a/tests/scripts/task_rust.sh b/tests/scripts/task_rust.sh
--- a/tests/scripts/task_rust.sh
+++ b/tests/scripts/task_rust.sh
@@ -2,24 +2,60 @@
 
 set -e
 
-export LD_LIBRARY_PATH=lib:$LD_LIBRARY_PATH
+export TVM_HOME="$(git rev-parse --show-toplevel)"
 
-tvm_root="$(git rev-parse --show-toplevel)"
-export PYTHONPATH="$tvm_root/python":"$tvm_root/nnvm/python":"$tvm_root/topi/python"
+export LD_LIBRARY_PATH="$TVM_HOME/lib":"$TVM_HOME/build":"$TVM_HOME/nnvm":$LD_LIBRARY_PATH
+export PYTHONPATH="$TVM_HOME/python":"$TVM_HOME/nnvm/python":"$TVM_HOME/topi/python"
+export RUST_DIR="$TVM_HOME/rust"
 
-#cd rust
-#cargo fmt -- --check
+cd $RUST_DIR
+cargo fmt -- --check
+
+# test common
+cd $RUST_DIR/common
+cargo build --features runtime
+cargo test --features runtime --tests
+
+cargo build --features frontend
+cargo test --features frontend --tests
+
+# test runtime
+cd $RUST_DIR/runtime
 
 # run basic tests
-#python3 tests/build_model.py
-#cargo test --tests
+python3 tests/build_model.py
+cargo test --tests
 
 # run TVM module test
-#cd tests/test_tvm_basic
-#cargo run
-#cd -
+cd tests/test_tvm_basic
+cargo run
+cd -
 
 # run NNVM graph test
-#cd tests/test_nnvm
-#cargo run
-#cd -
+cd tests/test_nnvm
+cargo run
+cd -
+
+# test frontend
+cd $RUST_DIR/frontend
+
+cargo test --tests -- --test-threads=1
+
+# run basic tests on cpu
+cd tests/basics
+cargo build --features cpu
+cargo run --features cpu
+# uncomment when have more CI resources
+# cargo build --features gpu
+# cargo run --features gpu
+# fi
+cd -
+
+# run callback tests separately: https://discuss.tvm.ai/t/are-global-functions-need-to-be-accessed-in-separate-processes/1075
+cd tests/callback
+cargo build
+cargo run --bin int
+cargo run --bin float
+cargo run --bin array
+cargo run --bin string
+cd -