ehsanmok / tvm-rust

(MERGED) Rust bindings for TVM runtime

DEPRECATED

The RFC is closed and this has been merge into TVM.

TVM Runtime Frontend Support

This crate provides an idiomatic Rust API for TVM runtime frontend as part of the ongoing RFC. Currently this requires Nightly Rust.

Checkout the docs.

What Does This Crate Offer?

Here is a major workflow

1. Train your Deep Learning model using any major framework such as PyTorch, Apache MXNet or TensorFlow
2. Use TVM to build optimized model artifacts on a supported context such as CPU, GPU, OpenCL, Vulkan, VPI, ROCM, etc.
3. Deploy your models using Rust ❤️

Example: Deploy Image Classification from Pretrained Resnet18 on ImageNet1k

Please checkout examples/resnet for the complete end-to-end example.

Here's a Python snippet for downloading and building a pretrained Resnet18 via MXNet and TVM

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)
# same the model artifacts
lib.save(os.path.join(target_dir, "deploy_lib.o"))
cc.create_shared(os.path.join(target_dir, "deploy_lib.so"),
                [os.path.join(target_dir, "deploy_lib.o")])

with open(os.path.join(target_dir, "deploy_graph.json"), "w") as fo:
    fo.write(graph.json())
with open(os.path.join(target_dir,"deploy_param.params"), "wb") as fo:
    fo.write(nnvm.compiler.save_param_dict(params))

Now, we need to input the artifacts to create and run the Graph Runtime to detect our input cat image

as demostrated in the following Rust snippet

let graph = fs::read_to_string("deploy_graph.json")?;
// load module
let lib = Module::load(&Path::new("deploy_lib.so"))?;
// get the global TVM graph runtime function
let runtime_create_fn = Function::get_function("tvm.graph_runtime.create", true).unwrap();

let runtime_create_fn_ret = call_packed!(
    runtime_create_fn,
    &graph,
    &lib,
    &ctx.device_type,
    &ctx.device_id
)?;
// get graph runtime module
let graph_runtime_module = runtime_create_fn_ret.to_module();
// get the registered `load_params` from runtime module
let load_param_fn = graph_runtime_module
    .get_function("load_params", false)
    .unwrap();
// parse parameters and convert to TVMByteArray
let params: Vec<u8> = fs::read("deploy_param.params")?;
let barr = TVMByteArray::from(&params);
// load the parameters
call_packed!(load_param_fn, &barr)?;
// get the set_input function
let set_input_fn = graph_runtime_module
    .get_function("set_input", false)
    .unwrap();

call_packed!(set_input_fn, "data", &input)?;
// get `run` function from runtime module
let run_fn = graph_runtime_module.get_function("run", false).unwrap();
// execute the run function. Note that it has no argument.
call_packed!(run_fn,)?;
// prepare to get the output
let output_shape = &mut [1, 1000];
let output = empty(output_shape, TVMContext::cpu(0), TVMType::from("float"));
// get the `get_output` function from runtime module
let get_output_fn = graph_runtime_module
    .get_function("get_output", false)
    .unwrap();
// execute the get output function
call_packed!(get_output_fn, &0, &output)?;
// flatten the output as Vec<f32>
let output = output.to_vec::<f32>()?;

Installations

Please follow TVM installations, 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, nnvm and topi need to be added to your PYTHONPATH or it's automatic via an Anaconda environment when install individually.

Supported TVM Functionalities

Use TVM to Generate Shared Library

One can use the following Python snippet to generate add_gpu.so which add two vectors on GPU.

import os
import tvm
from tvm.contrib import cc

def test_add(target_dir):
    if not tvm.module.enabled("cuda"):
        print(f"skip {__file__} because cuda is not enabled...")
        return
    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)
    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_cuda = tvm.build(s, [A, B, C], "cuda", target_host="llvm", name="myadd")

    fadd_cuda.save(os.path.join(target_dir, "add_gpu.o"))
    fadd_cuda.imported_modules[0].save(os.path.join(target_dir, "add_gpu.ptx"))
    cc.create_shared(os.path.join(target_dir, "add_gpu.so"),
            [os.path.join(target_dir, "add_gpu.o")])


if __name__ == "__main__":
    import sys
    if len(sys.argv) != 2:
        sys.exit(-1)
    test_add(sys.argv[1])

Run the Generated Shared Library

The following code snippet demonstrates how to load and test the generated shared library (add_gpu.so) in Rust.

extern crate tvm_frontend as tvm;

use tvm::*;

fn main() {
    let shape = &mut [2];
    let mut data = vec![3f32, 4.0];
    let mut arr = empty(shape, TVMContext::gpu(0), TVMType::from("float"));
    arr.copy_from_buffer(data.as_mut_slice());
    let mut ret = empty(shape, TVMContext::gpu(0), TVMType::from("float"));
    let path = Path::new("add_gpu.so");
    let ptx = Path::new("add_gpu.ptx");
    let mut fadd = Module::load(path).unwrap();
    let fadd_dep = Module::load(ptx).unwrap();
    assert!(fadd.enabled("gpu"));
    fadd.import_module(fadd_dep);
    fadd.entry_func();
    function::Builder::from(&mut fadd)
        .arg(&arr)
        .arg(&arr)
        .set_output(&mut ret)
        .invoke()
        .unwrap();

    assert_eq!(ret.to_vec::<f32>().unwrap(), vec![6f32, 8.0]);
}

Note: it is required to instruct the rustc to link to the generated add_gpu.so in runtime, for example by cargo:rustc-link-search=native=add_gpu.

See the tests and examples custom build.rs for more details.

Convert and Register a Rust Function as a TVM Packed Function

One can use register_global_func! macro to convert and register a Rust function of type fn(&[TVMArgValue]) -> Result<TVMRetValue> to a global TVM packed function as follows

#[macro_use]
extern crate tvm_frontend as tvm;

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 = empty(shape, TVMContext::cpu(0), TVMType::from("float"));
                let arr = arg.to_ndarray().copy_to_ndarray(e).unwrap();
                let rnd: ArrayD<f32> = ArrayD::try_from(&arr).unwrap();
                ret += rnd.scalar_sum();
            }
            let ret_val = TVMRetValue::from(&ret);
            Ok(ret_val)
        }
    }

    let shape = &mut [2];
    let mut data = vec![3f32, 4.0];
    let mut arr = empty(shape, TVMContext::cpu(0), TVMType::from("float"));
    arr.copy_from_buffer(data.as_mut_slice());
    let mut registered = function::Builder::default();
    registered
        .get_function("sum", true)
        .arg(&arr)
        .arg(&arr);

    assert_eq!(registered.invoke().unwrap().to_float(), 14f64);
    }

tvm-rust's Issues

Get rid of ValueKind in favor of to_type methods

Make TVMRetValue owned

Similar to this

Update tvm-sys to master and the new added types

Done

Add docs

TVMRetValue to TVMByteArray and other TODOs

Ergonimic change of Function to FunctionBase and Builder to Function

Simplify various impl macros

Resolved

Add TVMByteArray and conversions

Done

Better error handling

Create end-to-end examples

Congratulations!

Congratulations on this, great work!!