1 unstable release

0.1.0-alpha Feb 23, 2021

#521 in Science

Apache-2.0

97KB
2K SLoC

TVM Runtime Support

This crate provides an idiomatic Rust API for TVM runtime, see here for more details.

What Does This Crate Offer?

TVM is an end-to-end deep learning compiler which takes high level machine learning models or tensor computations and lowers them into executable code for a variety of heterogenous devices (e.g., CPU, GPU).

This crate provides access to the APIs for manipulating runtime data structures, as well as TVM's cross-language Object system which functions similarly to systems such as COM, enabling cross-language interoperability.

Installations

Please follow TVM installation instructions, export TVM_HOME=/path/to/tvm and add libtvm_runtime to your LD_LIBRARY_PATH.

Example of registering a cross-language closure.

One can use register! macro to expose a Rust closure with arguments which implement TryFrom<ArgValue> and return types which implement Into<RetValue>. Once registered with TVM these functions can be accessed via Python or C++, or any other language which implements the TVM packed function convention see the offcial documentation for more information.

use tvm_rt::{ArgValue, RetValue};
use tvm_rt::function::{Function, Result, register};

fn sum(x: i64, y: i64, z: i64) -> i64 {
    x + y + z
}

fn main() {
    register(sum, "mysum".to_owned()).unwrap();
    let func = Function::get("mysum").unwrap();
    let boxed_fn: Box<dyn Fn(i64, i64, i64) -> Result<i64>> = func.into();
    let ret = boxed_fn(10, 20, 30).unwrap();
    assert_eq!(ret, 60);
}

Dependencies

~3.5–8.5MB
~161K SLoC