{"id":13995406,"url":"https://github.com/ehsanmok/tvm-rust","last_synced_at":"2025-07-22T21:32:48.184Z","repository":{"id":57670803,"uuid":"135318745","full_name":"ehsanmok/tvm-rust","owner":"ehsanmok","description":"(MERGED) Rust bindings for TVM runtime","archived":true,"fork":false,"pushed_at":"2019-02-03T07:20:41.000Z","size":266,"stargazers_count":27,"open_issues_count":0,"forks_count":2,"subscribers_count":3,"default_branch":"master","last_synced_at":"2024-11-11T03:44:15.192Z","etag":null,"topics":["compiler","deep-learning","nnvm","rust-library","tvm"],"latest_commit_sha":null,"homepage":"https://github.com/dmlc/tvm/tree/master/rust","language":"Rust","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ehsanmok.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2018-05-29T15:46:08.000Z","updated_at":"2023-01-28T22:15:13.000Z","dependencies_parsed_at":"2022-09-12T15:14:08.257Z","dependency_job_id":null,"html_url":"https://github.com/ehsanmok/tvm-rust","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ehsanmok%2Ftvm-rust","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ehsanmok%2Ftvm-rust/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ehsanmok%2Ftvm-rust/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ehsanmok%2Ftvm-rust/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ehsanmok","download_url":"https://codeload.github.com/ehsanmok/tvm-rust/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":227177948,"owners_count":17743208,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["compiler","deep-learning","nnvm","rust-library","tvm"],"created_at":"2024-08-09T14:03:23.369Z","updated_at":"2024-11-29T17:31:41.750Z","avatar_url":"https://github.com/ehsanmok.png","language":"Rust","funding_links":[],"categories":["Rust","Model Inference"],"sub_categories":[],"readme":"# DEPRECATED\n\n**The RFC is closed and this has been merge into [TVM](https://github.com/dmlc/tvm/tree/master/rust).**\n\n# TVM Runtime Frontend Support\n\nThis crate provides an idiomatic Rust API for [TVM](https://github.com/dmlc/tvm) runtime frontend as part of the ~~[ongoing RFC](https://github.com/dmlc/tvm/issues/1601)~~. Currently this requires **Nightly Rust**.\n\nCheckout the [docs](https://ehsanmok.github.io/tvm_frontend/tvm_frontend/index.html).\n\n## What Does This Crate Offer?\n\nHere is a major workflow\n\n1. Train your **Deep Learning** model using any major framework such as [PyTorch](https://pytorch.org/), [Apache MXNet](https://mxnet.incubator.apache.org/) or [TensorFlow](https://www.tensorflow.org/)\n2. Use **TVM** to build optimized model artifacts on a supported context such as CPU, GPU, OpenCL, Vulkan, VPI, ROCM, etc.\n3. Deploy your models using **Rust** :heart:\n\n### Example: Deploy Image Classification from Pretrained Resnet18 on ImageNet1k\n\nPlease checkout [examples/resnet](https://github.com/ehsanmok/tvm-rust/tree/master/examples/resnet) for the complete end-to-end example.\n\nHere's a Python snippet for downloading and building a pretrained Resnet18 via MXNet and TVM\n\n```python\nblock = get_model('resnet18_v1', pretrained=True)\n    \nsym, params = nnvm.frontend.from_mxnet(block)\n# add the softmax layer for prediction\nnet = nnvm.sym.softmax(sym)\n# compile the model\nwith nnvm.compiler.build_config(opt_level=opt_level):\n    graph, lib, params = nnvm.compiler.build(\n        net, target, shape={\"data\": data_shape}, params=params)\n# same the model artifacts\nlib.save(os.path.join(target_dir, \"deploy_lib.o\"))\ncc.create_shared(os.path.join(target_dir, \"deploy_lib.so\"),\n                [os.path.join(target_dir, \"deploy_lib.o\")])\n\nwith open(os.path.join(target_dir, \"deploy_graph.json\"), \"w\") as fo:\n    fo.write(graph.json())\nwith open(os.path.join(target_dir,\"deploy_param.params\"), \"wb\") as fo:\n    fo.write(nnvm.compiler.save_param_dict(params))\n```\n\nNow, we need to input the artifacts to create and run the *Graph Runtime* to detect our input cat image\n\n![cat](https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true)\n\nas demostrated in the following Rust snippet\n\n```rust\n\nlet graph = fs::read_to_string(\"deploy_graph.json\")?;\n// load module\nlet lib = Module::load(\u0026Path::new(\"deploy_lib.so\"))?;\n// get the global TVM graph runtime function\nlet runtime_create_fn = Function::get_function(\"tvm.graph_runtime.create\", true).unwrap();\n\nlet runtime_create_fn_ret = call_packed!(\n    runtime_create_fn,\n    \u0026graph,\n    \u0026lib,\n    \u0026ctx.device_type,\n    \u0026ctx.device_id\n)?;\n// get graph runtime module\nlet graph_runtime_module = runtime_create_fn_ret.to_module();\n// get the registered `load_params` from runtime module\nlet load_param_fn = graph_runtime_module\n    .get_function(\"load_params\", false)\n    .unwrap();\n// parse parameters and convert to TVMByteArray\nlet params: Vec\u003cu8\u003e = fs::read(\"deploy_param.params\")?;\nlet barr = TVMByteArray::from(\u0026params);\n// load the parameters\ncall_packed!(load_param_fn, \u0026barr)?;\n// get the set_input function\nlet set_input_fn = graph_runtime_module\n    .get_function(\"set_input\", false)\n    .unwrap();\n\ncall_packed!(set_input_fn, \"data\", \u0026input)?;\n// get `run` function from runtime module\nlet run_fn = graph_runtime_module.get_function(\"run\", false).unwrap();\n// execute the run function. Note that it has no argument.\ncall_packed!(run_fn,)?;\n// prepare to get the output\nlet output_shape = \u0026mut [1, 1000];\nlet output = empty(output_shape, TVMContext::cpu(0), TVMType::from(\"float\"));\n// get the `get_output` function from runtime module\nlet get_output_fn = graph_runtime_module\n    .get_function(\"get_output\", false)\n    .unwrap();\n// execute the get output function\ncall_packed!(get_output_fn, \u00260, \u0026output)?;\n// flatten the output as Vec\u003cf32\u003e\nlet output = output.to_vec::\u003cf32\u003e()?;\n```\n\n## Installations\n\nPlease follow TVM [installations](https://docs.tvm.ai/install/index.html), `export TVM_HOME=/path/to/tvm` and add `libtvm_runtime` to your `LD_LIBRARY_PATH`.\n\n*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.\n\n## Supported TVM Functionalities\n\n### Use TVM to Generate Shared Library\n\nOne can use the following Python snippet to generate `add_gpu.so` which add two vectors on GPU.\n\n```python\nimport os\nimport tvm\nfrom tvm.contrib import cc\n\ndef test_add(target_dir):\n    if not tvm.module.enabled(\"cuda\"):\n        print(f\"skip {__file__} because cuda is not enabled...\")\n        return\n    n = tvm.var(\"n\")\n    A = tvm.placeholder((n,), name='A')\n    B = tvm.placeholder((n,), name='B')\n    C = tvm.compute(A.shape, lambda i: A[i] + B[i], name=\"C\")\n    s = tvm.create_schedule(C.op)\n    bx, tx = s[C].split(C.op.axis[0], factor=64)\n    s[C].bind(bx, tvm.thread_axis(\"blockIdx.x\"))\n    s[C].bind(tx, tvm.thread_axis(\"threadIdx.x\"))\n    fadd_cuda = tvm.build(s, [A, B, C], \"cuda\", target_host=\"llvm\", name=\"myadd\")\n\n    fadd_cuda.save(os.path.join(target_dir, \"add_gpu.o\"))\n    fadd_cuda.imported_modules[0].save(os.path.join(target_dir, \"add_gpu.ptx\"))\n    cc.create_shared(os.path.join(target_dir, \"add_gpu.so\"),\n            [os.path.join(target_dir, \"add_gpu.o\")])\n\n\nif __name__ == \"__main__\":\n    import sys\n    if len(sys.argv) != 2:\n        sys.exit(-1)\n    test_add(sys.argv[1])\n```\n\n### Run the Generated Shared Library\n\nThe following code snippet demonstrates how to load and test the generated shared library (`add_gpu.so`) in Rust.\n\n```rust\nextern crate tvm_frontend as tvm;\n\nuse tvm::*;\n\nfn main() {\n    let shape = \u0026mut [2];\n    let mut data = vec![3f32, 4.0];\n    let mut arr = empty(shape, TVMContext::gpu(0), TVMType::from(\"float\"));\n    arr.copy_from_buffer(data.as_mut_slice());\n    let mut ret = empty(shape, TVMContext::gpu(0), TVMType::from(\"float\"));\n    let path = Path::new(\"add_gpu.so\");\n    let ptx = Path::new(\"add_gpu.ptx\");\n    let mut fadd = Module::load(path).unwrap();\n    let fadd_dep = Module::load(ptx).unwrap();\n    assert!(fadd.enabled(\"gpu\"));\n    fadd.import_module(fadd_dep);\n    fadd.entry_func();\n    function::Builder::from(\u0026mut fadd)\n        .arg(\u0026arr)\n        .arg(\u0026arr)\n        .set_output(\u0026mut ret)\n        .invoke()\n        .unwrap();\n\n    assert_eq!(ret.to_vec::\u003cf32\u003e().unwrap(), vec![6f32, 8.0]);\n}\n```\n\n**Note:** it is required to instruct the `rustc` to link to the generated `add_gpu.so` in runtime, for example by\n`cargo:rustc-link-search=native=add_gpu`. \n\nSee the tests and examples custom `build.rs` for more details.\n\n### Convert and Register a Rust Function as a TVM Packed Function\n\nOne can use `register_global_func!` macro to convert and register a Rust \nfunction of type `fn(\u0026[TVMArgValue]) -\u003e Result\u003cTVMRetValue\u003e` to a global TVM **packed function** as follows\n\n```rust\n#[macro_use]\nextern crate tvm_frontend as tvm;\n\nuse tvm::*;\n\nfn main() {\n    register_global_func! {\n        fn sum(args: \u0026[TVMArgValue]) -\u003e Result\u003cTVMRetValue\u003e {\n            let mut ret = 0f32;\n            let shape = \u0026mut [2];\n            for arg in args.iter() {\n                let e = empty(shape, TVMContext::cpu(0), TVMType::from(\"float\"));\n                let arr = arg.to_ndarray().copy_to_ndarray(e).unwrap();\n                let rnd: ArrayD\u003cf32\u003e = ArrayD::try_from(\u0026arr).unwrap();\n                ret += rnd.scalar_sum();\n            }\n            let ret_val = TVMRetValue::from(\u0026ret);\n            Ok(ret_val)\n        }\n    }\n\n    let shape = \u0026mut [2];\n    let mut data = vec![3f32, 4.0];\n    let mut arr = empty(shape, TVMContext::cpu(0), TVMType::from(\"float\"));\n    arr.copy_from_buffer(data.as_mut_slice());\n    let mut registered = function::Builder::default();\n    registered\n        .get_function(\"sum\", true)\n        .arg(\u0026arr)\n        .arg(\u0026arr);\n\n    assert_eq!(registered.invoke().unwrap().to_float(), 14f64);\n    }\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fehsanmok%2Ftvm-rust","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fehsanmok%2Ftvm-rust","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fehsanmok%2Ftvm-rust/lists"}