1 个不稳定版本

0.1.0-alpha2021年2月23日

#521 in 科学

Apache-2.0

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