1 个不稳定版本
0.1.0-alpha | 2021年2月23日 |
---|
#521 in 科学
97KB
2K SLoC
TVM 运行时支持
此 crate 提供了对 TVM 运行时的 Rust API 的支持,详见 此处 了解更多详细信息。
此 crate 提供了什么?
TVM 是一个端到端深度学习编译器,它可以将高级机器学习模型或张量计算转换为适用于各种异构设备(例如 CPU、GPU)的可执行代码。
此 crate 提供了对操作运行时数据结构的 API 的访问,以及类似 COM 系统的 TVM 跨语言对象系统,这有助于实现跨语言互操作性。
安装
请遵循 TVM 安装说明,将 export TVM_HOME=/path/to/tvm
添加到您的 LD_LIBRARY_PATH
。
跨语言闭包注册示例。
可以使用 register!
宏来公开具有实现 TryFrom<ArgValue>
和返回类型实现 Into<RetValue>
的 Rust 闭包。一旦注册到 TVM,这些函数就可以通过 Python 或 C++ 或其他实现了 TVM 打包函数约定的语言访问。请参阅官方文档获取更多信息。
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);
}
依赖项
~3.5–8.5MB
~161K SLoC