You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2025/10/01 17:25:09 UTC
(tvm-ffi) branch main updated: [experimental] Initial bringup of rust support (#80)
This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 09477ce [experimental] Initial bringup of rust support (#80)
09477ce is described below
commit 09477ce10de566f8cf511cce7dfe56e77759f100
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Oct 1 13:25:03 2025 -0400
[experimental] Initial bringup of rust support (#80)
This PR brings up an initial foundation support for (experimental) rust
support.
Our primary focus remains to support high-polished python and c++
support.
Rust can be an useful system programming language for certain ml systems
use-cases. Having another typed system language would help us to future
proof some of the abi design. So having a third first-class language can
provide some values to the core.
The overall design focuses on the close-to-metal interopability for
efficiency while providing ergonomic rust APIs for the ffi libraries.
This PR support both exposing rust libraries as ffi abi and loading into
abi.
---
.github/workflows/ci_test.yml | 6 +
rust/Cargo.toml | 21 ++
rust/README.md | 72 ++++
rust/tvm-ffi-macros/Cargo.toml | 34 ++
rust/tvm-ffi-macros/src/lib.rs | 36 ++
rust/tvm-ffi-macros/src/object_macros.rs | 247 +++++++++++++
rust/tvm-ffi-macros/src/utils.rs | 78 ++++
rust/tvm-ffi-sys/Cargo.toml | 29 ++
rust/tvm-ffi-sys/build.rs | 66 ++++
rust/tvm-ffi-sys/src/c_api.rs | 432 ++++++++++++++++++++++
rust/tvm-ffi-sys/src/c_env_api.rs | 77 ++++
rust/tvm-ffi-sys/src/dlpack.rs | 124 +++++++
rust/tvm-ffi-sys/src/lib.rs | 25 ++
rust/tvm-ffi/Cargo.toml | 43 +++
rust/tvm-ffi/build.rs | 80 ++++
rust/tvm-ffi/examples/load_library.rs | 42 +++
rust/tvm-ffi/scripts/generate_example_lib.py | 56 +++
rust/tvm-ffi/src/any.rs | 362 ++++++++++++++++++
rust/tvm-ffi/src/collections/mod.rs | 21 ++
rust/tvm-ffi/src/collections/shape.rs | 139 +++++++
rust/tvm-ffi/src/collections/tensor.rs | 338 +++++++++++++++++
rust/tvm-ffi/src/derive.rs | 20 +
rust/tvm-ffi/src/device.rs | 96 +++++
rust/tvm-ffi/src/dtype.rs | 209 +++++++++++
rust/tvm-ffi/src/error.rs | 205 +++++++++++
rust/tvm-ffi/src/extra/mod.rs | 19 +
rust/tvm-ffi/src/extra/module.rs | 77 ++++
rust/tvm-ffi/src/function.rs | 272 ++++++++++++++
rust/tvm-ffi/src/function_internal.rs | 193 ++++++++++
rust/tvm-ffi/src/lib.rs | 52 +++
rust/tvm-ffi/src/macros.rs | 398 ++++++++++++++++++++
rust/tvm-ffi/src/object.rs | 473 ++++++++++++++++++++++++
rust/tvm-ffi/src/string.rs | 533 +++++++++++++++++++++++++++
rust/tvm-ffi/src/type_traits.rs | 333 +++++++++++++++++
rust/tvm-ffi/tests/test_any.rs | 349 ++++++++++++++++++
rust/tvm-ffi/tests/test_device.rs | 41 +++
rust/tvm-ffi/tests/test_dtype.rs | 106 ++++++
rust/tvm-ffi/tests/test_error.rs | 82 +++++
rust/tvm-ffi/tests/test_function.rs | 194 ++++++++++
rust/tvm-ffi/tests/test_object.rs | 147 ++++++++
rust/tvm-ffi/tests/test_shape.rs | 38 ++
rust/tvm-ffi/tests/test_string.rs | 373 +++++++++++++++++++
rust/tvm-ffi/tests/test_tensor.rs | 112 ++++++
43 files changed, 6650 insertions(+)
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
index 2edd8c2..aec3266 100644
--- a/.github/workflows/ci_test.yml
+++ b/.github/workflows/ci_test.yml
@@ -131,3 +131,9 @@ jobs:
- name: Run python tests
run: |
pytest -vvs tests/python
+
+ # Run Rust tests, must happen after installing the pip package.
+ - name: Run rust tests
+ working-directory: rust
+ run: |
+ cargo test
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
new file mode 100644
index 0000000..d924376
--- /dev/null
+++ b/rust/Cargo.toml
@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[workspace]
+members = ["tvm-ffi", "tvm-ffi-sys", "tvm-ffi-macros"]
+
+resolver = "2"
diff --git a/rust/README.md b/rust/README.md
new file mode 100644
index 0000000..31f3fb0
--- /dev/null
+++ b/rust/README.md
@@ -0,0 +1,72 @@
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements. See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership. The ASF licenses this file -->
+<!--- to you under the Apache License, Version 2.0 (the -->
+<!--- "License"); you may not use this file except in compliance -->
+<!--- with the License. You may obtain a copy of the License at -->
+
+<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
+
+<!--- Unless required by applicable law or agreed to in writing, -->
+<!--- software distributed under the License is distributed on an -->
+<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
+<!--- KIND, either express or implied. See the License for the -->
+<!--- specific language governing permissions and limitations -->
+<!--- under the License. -->
+
+# Rust Packages
+
+(Experimental) Rust support for the `tvm-ffi` ABI.
+Currently, the rust support is in an experimental stage.
+This workspace contains three crates:
+
+- `tvm-ffi`: Safe, ergonomic Rust bindings over the ABI.
+- `tvm-ffi-sys`: Low-level exposure of raw C ABIs.
+- `tvm-ffi-macros`: Procedural macros used by `tvm-ffi` (derive/object helpers and exported function helpers).
+
+The overall project focuses on low-level, direct access to the ABI when possible for maximum
+efficiency while maintaining interoperability.
+
+## Installation
+
+The Rust support depends on `libtvm_ffi`.
+Please install the `tvm-ffi` pip package by running:
+
+```bash
+pip install -v ..
+```
+
+Confirm that `tvm-ffi-config` is available with:
+
+```bash
+tvm-ffi-config --libdir
+```
+
+Then build the workspace with:
+
+```bash
+cargo build
+```
+
+The build will:
+
+- Query `tvm-ffi-config --libdir` to add the appropriate link search path.
+- Link against `tvm_ffi`.
+- Update the appropriate dynamic loader path environment variable for `cargo run` and `cargo test`.
+
+For running downstream applications, you need to set the `LD_LIBRARY_PATH` so `libtvm_ffi` is available in the path.
+
+```bash
+export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:`tvm-ffi-config --libdir`
+```
+
+## Running Examples
+
+You can run an optional library-loading example similar to the quick_start examples in [examples/quick_start](../examples/quick_start/).
+
+```bash
+cargo run --example load_library --features example
+```
+
+Check out the [load_library.rs](tvm-ffi/examples/load_library.rs) for details.
diff --git a/rust/tvm-ffi-macros/Cargo.toml b/rust/tvm-ffi-macros/Cargo.toml
new file mode 100644
index 0000000..f8d29d4
--- /dev/null
+++ b/rust/tvm-ffi-macros/Cargo.toml
@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "tvm-ffi-macros"
+description = "Procedural macro crate for tvm-ffi"
+
+version = "0.1.0-alpha.0"
+edition = "2021"
+license = "Apache-2.0"
+
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "^1.0"
+quote = "^1.0"
+syn = { version = "1.0.48", features = ["full", "parsing", "extra-traits"] }
+proc-macro-error = "^1.0"
diff --git a/rust/tvm-ffi-macros/src/lib.rs b/rust/tvm-ffi-macros/src/lib.rs
new file mode 100644
index 0000000..64fe3f1
--- /dev/null
+++ b/rust/tvm-ffi-macros/src/lib.rs
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+use proc_macro::TokenStream;
+use proc_macro_error::proc_macro_error;
+
+mod object_macros;
+mod utils;
+
+#[proc_macro_error]
+#[proc_macro_derive(Object, attributes(type_key, type_index))]
+pub fn derive_object(input: TokenStream) -> TokenStream {
+ TokenStream::from(object_macros::derive_object(input))
+}
+
+#[proc_macro_error]
+#[proc_macro_derive(ObjectRef, attributes(type_key, type_index))]
+pub fn derive_object_ref(input: TokenStream) -> TokenStream {
+ TokenStream::from(object_macros::derive_object_ref(input))
+}
diff --git a/rust/tvm-ffi-macros/src/object_macros.rs b/rust/tvm-ffi-macros/src/object_macros.rs
new file mode 100644
index 0000000..8154709
--- /dev/null
+++ b/rust/tvm-ffi-macros/src/object_macros.rs
@@ -0,0 +1,247 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::DeriveInput;
+
+use crate::utils::*;
+
+/// Derive Object trait for a struct to generate boilerplate code
+pub fn derive_object(input: proc_macro::TokenStream) -> TokenStream {
+ let tvm_ffi_crate = get_tvm_ffi_crate();
+ let derive_input = syn::parse_macro_input!(input as DeriveInput);
+ let struct_name = derive_input.ident.clone();
+
+ let type_key = get_attr(&derive_input, "type_key")
+ .map(attr_to_str)
+ .expect("Expect #[type_key = \"<my_type_key>\"] attribute");
+
+ // type index can be optional
+ // for now we make it required for static index
+ let type_index_tokens = match get_attr(&derive_input, "type_index").map(attr_to_expr) {
+ Some(type_index) => {
+ let type_index_expr =
+ type_index.expect("Expect #[type_index(TypeIndex::<my_type_index>)] attribute");
+ quote! {
+ #[inline]
+ fn type_index() -> i32 {
+ #type_index_expr as i32
+ }
+ }
+ }
+ None => {
+ quote! {
+ #[inline]
+ fn type_index() -> i32 {
+ static TYPE_INDEX: std::sync::LazyLock<i32> = std::sync::LazyLock::new(||
+ unsafe {
+ let type_key_arg =
+ #tvm_ffi_crate::tvm_ffi_sys::TVMFFIByteArray::from_str(#type_key);
+ let mut tindex = 0;
+ let ret = #tvm_ffi_crate::tvm_ffi_sys::TVMFFITypeKeyToIndex(
+ &type_key_arg, &mut tindex
+ );
+ if ret != 0 {
+ proc_macro_error::abort!("Failed to get type index for type key: {}", #type_key);
+ }
+ tindex
+ }
+ );
+ *TYPE_INDEX
+ }
+ }
+ }
+ };
+ // search for field name base and derive the base type
+ // we expect base always to be the first field
+ let base_def_tokens = match &derive_input.data {
+ syn::Data::Struct(s) => s.fields.iter().next().and_then(|f| {
+ let (base_id, base_ty) = (f.ident.clone()?, f.ty.clone());
+ // The transitive case of subtyping
+ Some(quote! {
+ #[inline]
+ unsafe fn object_header_mut(
+ this: &mut Self
+ ) -> &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject {
+ const _: () = {
+ fn assert_impl<T: #tvm_ffi_crate::object::ObjectCore>() {}
+ let _ = assert_impl::<#base_ty>;
+ };
+ #base_ty::object_header_mut(&mut this.#base_id)
+ }
+ })
+ }),
+ _ => panic!("First field must be `<base_name>: <ObjectCoreType>`"),
+ };
+
+ let expanded = quote! {
+ unsafe impl #tvm_ffi_crate::object::ObjectCore for #struct_name {
+ const TYPE_KEY: &'static str = #type_key;
+
+ #type_index_tokens
+
+ #base_def_tokens
+ }
+ };
+ TokenStream::from(expanded)
+}
+
+/// Derive ObjectRef trait for a struct to generate boilerplate code
+pub fn derive_object_ref(input: proc_macro::TokenStream) -> TokenStream {
+ let tvm_ffi_crate = get_tvm_ffi_crate();
+ let derive_input = syn::parse_macro_input!(input as DeriveInput);
+ let struct_name = derive_input.ident.clone();
+
+ // search for field name base and derive the base type
+ // we expect base always to be the first field
+ let data_ty = match &derive_input.data {
+ syn::Data::Struct(s) => s.fields.iter().next().and_then(|f| {
+ let (data_id, data_ty) = (f.ident.clone()?, f.ty.clone());
+ if data_id == "data" {
+ // The transitive case of subtyping
+ Some(data_ty)
+ } else {
+ None
+ }
+ }),
+ _ => panic!("derive only works for structs"),
+ }
+ .expect("First field must be `data: ObjectArc<T>`");
+
+ let mut expanded = quote! {
+ unsafe impl #tvm_ffi_crate::object::ObjectRefCore for #struct_name {
+ type ContainerType = <#data_ty as std::ops::Deref>::Target;
+ #[inline]
+ fn data(this: &Self) -> &ObjectArc<Self::ContainerType> {
+ &this.data
+ }
+ #[inline]
+ fn into_data(this: Self) -> ObjectArc<Self::ContainerType> {
+ this.data
+ }
+ #[inline]
+ fn from_data(data: ObjectArc<Self::ContainerType>) -> Self {
+ Self { data}
+ }
+ }
+
+ // implement AnyCompatible for #struct_name
+ unsafe impl #tvm_ffi_crate::type_traits::AnyCompatible for #struct_name {
+ fn type_str() -> String {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ <ContainerType as #tvm_ffi_crate::object::ObjectCore>::TYPE_KEY.into()
+ }
+
+ unsafe fn copy_to_any_view(
+ src: &Self,
+ data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
+ ) {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ let type_index =
+ <ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
+ data.type_index = type_index as i32;
+ data.small_str_len = 0;
+ let data_ptr = #tvm_ffi_crate::object::ObjectArc::<ContainerType>::as_raw(
+ &src.data
+ );
+ data.data_union.v_obj =
+ data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
+ }
+
+ unsafe fn check_any_strict(data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny) -> bool {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ let type_index =
+ <ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
+ data.type_index == type_index as i32
+ }
+
+ unsafe fn copy_from_any_view_after_check(
+ data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
+ ) -> Self {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ let data_ptr = data.data_union.v_obj;
+ // need to increase ref because original weak ptr
+ // do not own the code
+ #tvm_ffi_crate::object::unsafe_::inc_ref(
+ data_ptr as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject
+ );
+ Self {
+ data : #tvm_ffi_crate::object::ObjectArc::from_raw(
+ data_ptr as *mut ContainerType
+ )
+ }
+ }
+
+ unsafe fn move_to_any(
+ src: Self,
+ data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
+ ) {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ let type_index =
+ <ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
+ data.type_index = type_index as i32;
+ data.small_str_len = 0;
+ let data_ptr = #tvm_ffi_crate::object::ObjectArc::into_raw(
+ src.data
+ );
+ data.data_union.v_obj =
+ data_ptr as *mut ContainerType as *mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIObject;
+ }
+
+ unsafe fn move_from_any_after_check(
+ data: &mut #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
+ ) -> Self {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ let data_ptr = data.data_union.v_obj as *mut ContainerType;
+ Self {
+ data : #tvm_ffi_crate::object::ObjectArc::<ContainerType>::from_raw(data_ptr)
+ }
+ }
+
+ unsafe fn try_cast_from_any_view(
+ data: & #tvm_ffi_crate::tvm_ffi_sys::TVMFFIAny
+ ) -> Result<Self, ()> {
+ type ContainerType = <#struct_name as #tvm_ffi_crate::object::ObjectRefCore>
+ ::ContainerType;
+ let type_index =
+ <ContainerType as #tvm_ffi_crate::object::ObjectCore>::type_index();
+ if data.type_index == type_index as i32 {
+ Ok(Self::copy_from_any_view_after_check(data))
+ } else {
+ Err(())
+ }
+ }
+ }
+ };
+ // skip ObjectRef since it can create circular dependency with any.rs
+ if struct_name != "ObjectRef" {
+ expanded.extend(quote! {
+ #tvm_ffi_crate::impl_try_from_any!(#struct_name);
+ #tvm_ffi_crate::impl_arg_into_ref!(#struct_name);
+ #tvm_ffi_crate::impl_into_arg_holder_default!(#struct_name);
+ });
+ }
+ TokenStream::from(expanded)
+}
diff --git a/rust/tvm-ffi-macros/src/utils.rs b/rust/tvm-ffi-macros/src/utils.rs
new file mode 100644
index 0000000..5b6f4c1
--- /dev/null
+++ b/rust/tvm-ffi-macros/src/utils.rs
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use proc_macro2::TokenStream;
+use quote::quote;
+use std::env;
+
+/// Get the tvm-rt crate name
+/// \return The tvm-rt crate name
+pub(crate) fn get_tvm_ffi_crate() -> TokenStream {
+ if env::var("CARGO_PKG_NAME").unwrap() == "tvm-ffi" {
+ quote!(crate)
+ } else {
+ quote!(tvm_ffi)
+ }
+}
+
+/// Get an attribute by name from a derive input
+///
+/// # Arguments
+/// * `derive_input` - The derive input to get the attribute from
+/// * `name` - The name of the attribute to get
+///
+/// # Returns
+/// * `Option<&syn::Attribute>` - The attribute if it exists
+pub(crate) fn get_attr<'a>(
+ derive_input: &'a syn::DeriveInput,
+ name: &str,
+) -> Option<&'a syn::Attribute> {
+ derive_input.attrs.iter().find(|a| a.path.is_ident(name))
+}
+
+/// Convert an attribute to a string
+///
+/// # Arguments
+/// * `attr` - The attribute to convert
+///
+/// # Returns
+/// * `syn::LitStr` - The string value of the attribute
+pub(crate) fn attr_to_str(attr: &syn::Attribute) -> syn::LitStr {
+ match attr.parse_meta() {
+ Ok(syn::Meta::NameValue(syn::MetaNameValue {
+ lit: syn::Lit::Str(s),
+ ..
+ })) => s,
+ Ok(m) => panic!("Expected a string literal, got"),
+ Err(e) => panic!("{}", e),
+ }
+}
+
+/// Convert an attribute to an integer
+///
+/// # Arguments
+/// * `attr` - The attribute to convert
+///
+/// # Returns
+/// * `syn::Result<syn::Expr>` - The integer value of the attribute
+pub(crate) fn attr_to_expr(attr: &syn::Attribute) -> syn::Result<syn::Expr> {
+ let parser = |input: syn::parse::ParseStream| {
+ input.parse::<syn::Expr>() // parse expression after '='
+ };
+ syn::parse::Parser::parse2(parser, attr.tokens.clone())
+}
diff --git a/rust/tvm-ffi-sys/Cargo.toml b/rust/tvm-ffi-sys/Cargo.toml
new file mode 100644
index 0000000..ef87038
--- /dev/null
+++ b/rust/tvm-ffi-sys/Cargo.toml
@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "tvm-ffi-sys"
+description = "Low-level sys crate for tvm-ffi"
+
+version = "0.1.0-alpha.0"
+edition = "2021"
+license = "Apache-2.0"
+
+
+[lib]
+name = "tvm_ffi_sys"
+crate-type = ["lib"]
diff --git a/rust/tvm-ffi-sys/build.rs b/rust/tvm-ffi-sys/build.rs
new file mode 100644
index 0000000..24e1f6c
--- /dev/null
+++ b/rust/tvm-ffi-sys/build.rs
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use std::env;
+use std::process::Command;
+
+/// Update the LD_LIBRARY_PATH environment variable
+/// so cargo run/test can directly pick it up
+/// note that it won't always work later consumption of the
+/// library, and we still need to figure out linking by setting LD_LIBRARY_PATH
+fn update_ld_library_path(lib_dir: &str) {
+ let os_env_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() {
+ Ok("windows") => "PATH",
+ Ok("macos") => "DYLD_LIBRARY_PATH",
+ Ok("linux") => "LD_LIBRARY_PATH",
+ _ => "",
+ };
+ if os_env_var.is_empty() {
+ return;
+ }
+ // Get the current value of the environment variable at build time (if any)
+ let current_val = env::var(os_env_var).unwrap_or_else(|_| String::new());
+ // Use platform-specific separator
+ let separator = if os_env_var == "PATH" { ";" } else { ":" };
+ let new_ld_path = if current_val.is_empty() {
+ lib_dir.to_string()
+ } else {
+ format!("{}{}{}", current_val, separator, lib_dir)
+ };
+ // this env is only used for cargo run/test
+ println!("cargo:rustc-env={}={}", os_env_var, new_ld_path);
+}
+
+fn main() {
+ // Run `mylib-config --libdir` to get the library path
+ let config_output = Command::new("tvm-ffi-config")
+ .arg("--libdir")
+ .output()
+ .expect("Failed to run tvm-ffi-config");
+ let lib_dir = String::from_utf8(config_output.stdout)
+ .expect("Invalid UTF-8 output from tvm-ffi-config")
+ .trim()
+ .to_string();
+
+ // add the library directory to the linker search path
+ println!("cargo:rustc-link-search=native={}", lib_dir);
+ // link the library
+ println!("cargo:rustc-link-lib=dylib=tvm_ffi");
+ // update the LD_LIBRARY_PATH environment variable
+ update_ld_library_path(&lib_dir);
+}
diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs
new file mode 100644
index 0000000..a55470d
--- /dev/null
+++ b/rust/tvm-ffi-sys/src/c_api.rs
@@ -0,0 +1,432 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+// NOTE: we manually write the C ABI as they are reasonably minimal
+// and we need to ensure clear control of the atomic access etc.
+#![allow(non_camel_case_types)]
+
+use std::ffi::c_void;
+use std::sync::atomic::AtomicU64;
+
+use crate::dlpack::DLDataType;
+use crate::dlpack::DLDevice;
+
+/// The index type of the FFI objects
+#[repr(i32)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum TVMFFITypeIndex {
+ /// None/nullptr value
+ kTVMFFINone = 0,
+ /// POD int value
+ kTVMFFIInt = 1,
+ /// POD bool value
+ kTVMFFIBool = 2,
+ /// POD float value
+ kTVMFFIFloat = 3,
+ /// Opaque pointer object
+ kTVMFFIOpaquePtr = 4,
+ /// DLDataType
+ kTVMFFIDataType = 5,
+ /// DLDevice
+ kTVMFFIDevice = 6,
+ /// DLTensor*
+ kTVMFFIDLTensorPtr = 7,
+ /// const char*
+ kTVMFFIRawStr = 8,
+ /// TVMFFIByteArray*
+ kTVMFFIByteArrayPtr = 9,
+ /// R-value reference to ObjectRef
+ kTVMFFIObjectRValueRef = 10,
+ /// Small string on stack
+ kTVMFFISmallStr = 11,
+ /// Small bytes on stack
+ kTVMFFISmallBytes = 12,
+ /// Start of statically defined objects.
+ kTVMFFIStaticObjectBegin = 64,
+ /// String object, layout = { TVMFFIObject, TVMFFIByteArray, ... }
+ kTVMFFIStr = 65,
+ /// Bytes object, layout = { TVMFFIObject, TVMFFIByteArray, ... }
+ kTVMFFIBytes = 66,
+ /// Error object.
+ kTVMFFIError = 67,
+ /// Function object.
+ kTVMFFIFunction = 68,
+ /// Shape object, layout = { TVMFFIObject, { const int64_t*, size_t }, ... }
+ kTVMFFIShape = 69,
+ /// Tensor object, layout = { TVMFFIObject, DLTensor, ... }
+ kTVMFFITensor = 70,
+ /// Array object.
+ kTVMFFIArray = 71,
+ //----------------------------------------------------------------
+ // more complex objects
+ //----------------------------------------------------------------
+ /// Map object.
+ kTVMFFIMap = 72,
+ /// Runtime dynamic loaded module object.
+ kTVMFFIModule = 73,
+ /// Opaque python object.
+ kTVMFFIOpaquePyObject = 74,
+}
+
+#[repr(i32)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum TVMFFIObjectDeleterFlagBitMask {
+ kTVMFFIObjectDeleterFlagBitMaskStrong = 1 << 0,
+ kTVMFFIObjectDeleterFlagBitMaskWeak = 1 << 1,
+ kTVMFFIObjectDeleterFlagBitMaskBoth = (1 << 0) | (1 << 1),
+}
+
+/// Handle to Object from C API's pov
+pub type TVMFFIObjectHandle = *mut c_void;
+pub type TVMFFIObjectDeleter = unsafe extern "C" fn(self_ptr: *mut c_void, flags: i32);
+
+// constants for working with combined reference count
+pub const COMBINED_REF_COUNT_MASK_U32: u64 = (1u64 << 32) - 1;
+pub const COMBINED_REF_COUNT_STRONG_ONE: u64 = 1;
+pub const COMBINED_REF_COUNT_WEAK_ONE: u64 = 1u64 << 32;
+pub const COMBINED_REF_COUNT_BOTH_ONE: u64 =
+ COMBINED_REF_COUNT_STRONG_ONE | COMBINED_REF_COUNT_WEAK_ONE;
+
+#[repr(C)]
+pub struct TVMFFIObject {
+ pub combined_ref_count: AtomicU64,
+ pub type_index: i32,
+ pub __padding: u32,
+ pub deleter: Option<TVMFFIObjectDeleter>,
+ // private padding to ensure 8 bytes alignment
+ #[cfg(target_pointer_width = "32")]
+ __padding: u32,
+}
+
+impl TVMFFIObject {
+ pub fn new() -> Self {
+ Self {
+ combined_ref_count: AtomicU64::new(0),
+ type_index: 0,
+ __padding: 0,
+ deleter: None,
+ }
+ }
+}
+
+/// Second union in TVMFFIAny - 8 bytes
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub union TVMFFIAnyDataUnion {
+ /// Integers
+ pub v_int64: i64,
+ /// Floating-point numbers
+ pub v_float64: f64,
+ /// Typeless pointers
+ pub v_ptr: *mut c_void,
+ /// Raw C-string
+ pub v_c_str: *const i8,
+ /// Ref counted objects
+ pub v_obj: *mut TVMFFIObject,
+ /// Data type
+ pub v_dtype: DLDataType,
+ /// Device
+ pub v_device: DLDevice,
+ /// Small string
+ pub v_bytes: [u8; 8],
+ /// uint64 repr mainly used for hashing
+ pub v_uint64: u64,
+}
+
+/// TVM FFI Any value - a union type that can hold various data types
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct TVMFFIAny {
+ /// Type index of the object.
+ /// The type index of Object and Any are shared in FFI.
+ pub type_index: i32,
+ /// small string length or zero padding
+ pub small_str_len: u32,
+ /// data union - 8 bytes
+ pub data_union: TVMFFIAnyDataUnion,
+}
+
+impl TVMFFIAny {
+ /// create a new instance of TVMFFIAny that represents None
+ pub fn new() -> Self {
+ Self {
+ type_index: TVMFFITypeIndex::kTVMFFINone as i32,
+ small_str_len: 0,
+ data_union: TVMFFIAnyDataUnion { v_int64: 0 },
+ }
+ }
+}
+
+/// Byte array data structure used by String and Bytes.
+#[repr(C)]
+pub struct TVMFFIByteArray {
+ pub data: *const u8,
+ pub size: usize,
+}
+
+impl TVMFFIByteArray {
+ pub fn new(data: *const u8, size: usize) -> Self {
+ Self { data, size }
+ }
+ /// Convert the TVMFFIByteArray to a str view
+ ///
+ /// # Arguments
+ /// * `self` - The TVMFFIByteArray to convert.
+ ///
+ /// # Returns
+ /// * `&str` - The converted str view.
+ pub fn as_str(&self) -> &str {
+ unsafe { std::str::from_utf8_unchecked(std::slice::from_raw_parts(self.data, self.size)) }
+ }
+ /// Unsafe function to create a TVMFFIByteArray from a string
+ /// This function is unsafe as it does not check lifetime of the string
+ /// the caller must ensure that the string is valid for the lifetime of the TVMFFIByteArray
+ ///
+ /// # Arguments
+ /// * `data` - The string to create the TVMFFIByteArray from.
+ ///
+ /// # Returns
+ /// * `TVMFFIByteArray` - The created TVMFFIByteArray.
+ pub unsafe fn from_str(data: &str) -> Self {
+ Self {
+ data: data.as_ptr(),
+ size: data.len(),
+ }
+ }
+}
+
+/// Safe call type for function ABI
+pub type TVMFFISafeCallType = unsafe extern "C" fn(
+ handle: *mut c_void,
+ args: *const TVMFFIAny,
+ num_args: i32,
+ result: *mut TVMFFIAny,
+) -> i32;
+
+/// Function cell
+#[repr(C)]
+pub struct TVMFFIFunctionCell {
+ /// A C API compatible call with exception catching.
+ pub safe_call: TVMFFISafeCallType,
+ pub cxx_call: *mut c_void,
+}
+
+unsafe impl Send for TVMFFIFunctionCell {}
+unsafe impl Sync for TVMFFIFunctionCell {}
+
+#[repr(i32)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum TVMFFIBacktraceUpdateMode {
+ kTVMFFIBacktraceUpdateModeReplace = 0,
+ kTVMFFIBacktraceUpdateModeAppend = 1,
+}
+
+/// Error cell used in error object following header.
+#[repr(C)]
+pub struct TVMFFIErrorCell {
+ pub kind: TVMFFIByteArray,
+ pub message: TVMFFIByteArray,
+ pub backtrace: TVMFFIByteArray,
+ pub update_backtrace: unsafe extern "C" fn(
+ self_ptr: *mut c_void,
+ backtrace: *const TVMFFIByteArray,
+ update_mode: i32,
+ ),
+}
+
+/// Shape cell used in shape object following header.
+#[repr(C)]
+pub struct TVMFFIShapeCell {
+ pub data: *const i64,
+ pub size: usize,
+}
+
+/// Field getter function pointer type
+pub type TVMFFIFieldGetter =
+ unsafe extern "C" fn(field: *mut c_void, result: *mut TVMFFIAny) -> i32;
+
+/// Field setter function pointer type
+pub type TVMFFIFieldSetter =
+ unsafe extern "C" fn(field: *mut c_void, value: *const TVMFFIAny) -> i32;
+
+/// Information support for optional object reflection
+#[repr(C)]
+pub struct TVMFFIFieldInfo {
+ /// The name of the field
+ pub name: TVMFFIByteArray,
+ /// The docstring about the field
+ pub doc: TVMFFIByteArray,
+ /// The metadata of the field in JSON string
+ pub metadata: TVMFFIByteArray,
+ /// bitmask flags of the field
+ pub flags: i64,
+ /// The size of the field
+ pub size: i64,
+ /// The alignment of the field
+ pub alignment: i64,
+ /// The offset of the field
+ pub offset: i64,
+ /// The getter to access the field
+ pub getter: Option<TVMFFIFieldGetter>,
+ /// The setter to access the field
+ /// The setter is set even if the field is readonly for serialization
+ pub setter: Option<TVMFFIFieldSetter>,
+ /// The default value of the field, this field hold AnyView,
+ /// valid when flags set kTVMFFIFieldFlagBitMaskHasDefault
+ pub default_value: TVMFFIAny,
+ /// Records the static type kind of the field
+ pub field_static_type_index: i32,
+}
+
+/// Object creator function pointer type
+pub type TVMFFIObjectCreator = unsafe extern "C" fn(result: *mut TVMFFIObjectHandle) -> i32;
+
+/// Method information that can appear in reflection table
+#[repr(C)]
+pub struct TVMFFIMethodInfo {
+ /// The name of the field
+ pub name: TVMFFIByteArray,
+ /// The docstring about the method
+ pub doc: TVMFFIByteArray,
+ /// Optional metadata of the method in JSON string
+ pub metadata: TVMFFIByteArray,
+ /// bitmask flags of the method
+ pub flags: i64,
+ /// The method wrapped as ffi::Function, stored as AnyView
+ /// The first argument to the method is always the self for instance methods
+ pub method: TVMFFIAny,
+}
+
+/// Extra information of object type that can be used for reflection
+///
+/// This information is optional and can be used to enable reflection based
+/// creation of the object.
+#[repr(C)]
+pub struct TVMFFITypeMetadata {
+ /// The docstring about the object
+ pub doc: TVMFFIByteArray,
+ /// An optional function that can create a new empty instance of the type
+ pub creator: Option<TVMFFIObjectCreator>,
+ /// Total size of the object struct, if it is fixed and known
+ ///
+ /// This field is set optional and set to 0 if not registered.
+ pub total_size: i32,
+ /// Optional meta-data for structural eq/hash
+ pub structural_eq_hash_kind: i32,
+}
+
+/// Column array that stores extra attributes about types
+///
+/// The attributes stored in a column array that can be looked up by type index.
+/// Note that the TypeAttr behaves like type_traits so column T so not contain
+/// attributes from base classes.
+#[repr(C)]
+pub struct TVMFFITypeAttrColumn {
+ /// The data of the column
+ pub data: *const TVMFFIAny,
+ /// The size of the column
+ pub size: usize,
+}
+
+/// Runtime type information for object type checking
+#[repr(C)]
+pub struct TVMFFITypeInfo {
+ /// The runtime type index
+ /// It can be allocated during runtime if the type is dynamic
+ pub type_index: i32,
+ /// number of parent types in the type hierachy
+ pub type_depth: i32,
+ /// the unique type key to identify the type
+ pub type_key: TVMFFIByteArray,
+ /// `type_acenstors[depth]` stores the type_index of the acenstors at depth level
+ /// To keep things simple, we do not allow multiple inheritance so the
+ /// hieracy stays as a tree
+ pub type_acenstors: *const *const TVMFFITypeInfo,
+ /// Cached hash value of the type key, used for consistent structural hashing
+ pub type_key_hash: u64,
+ /// number of reflection accessible fields
+ pub num_fields: i32,
+ /// number of reflection acccesible methods
+ pub num_methods: i32,
+ /// The reflection field information
+ pub fields: *const TVMFFIFieldInfo,
+ /// The reflection method
+ pub methods: *const TVMFFIMethodInfo,
+ /// The extra information of the type
+ pub metadata: *const TVMFFITypeMetadata,
+}
+
+unsafe extern "C" {
+ pub fn TVMFFITypeKeyToIndex(type_key: *const TVMFFIByteArray, out_tindex: *mut i32) -> i32;
+ pub fn TVMFFIFunctionGetGlobal(
+ name: *const TVMFFIByteArray,
+ out: *mut TVMFFIObjectHandle,
+ ) -> i32;
+ pub fn TVMFFIFunctionSetGlobal(
+ name: *const TVMFFIByteArray,
+ f: TVMFFIObjectHandle,
+ can_override: i32,
+ ) -> i32;
+ pub fn TVMFFIFunctionCreate(
+ self_ptr: *mut c_void,
+ safe_call: TVMFFISafeCallType,
+ deleter: Option<unsafe extern "C" fn(*mut c_void)>,
+ out: *mut TVMFFIObjectHandle,
+ ) -> i32;
+ pub fn TVMFFIAnyViewToOwnedAny(any_view: *const TVMFFIAny, out: *mut TVMFFIAny) -> i32;
+ pub fn TVMFFIFunctionCall(
+ func: TVMFFIObjectHandle,
+ args: *const TVMFFIAny,
+ num_args: i32,
+ result: *mut TVMFFIAny,
+ ) -> i32;
+ pub fn TVMFFIErrorMoveFromRaised(result: *mut TVMFFIObjectHandle);
+ pub fn TVMFFIErrorSetRaised(error: TVMFFIObjectHandle);
+ pub fn TVMFFIErrorSetRaisedFromCStr(kind: *const i8, message: *const i8);
+ pub fn TVMFFIErrorCreate(
+ kind: *const TVMFFIByteArray,
+ message: *const TVMFFIByteArray,
+ backtrace: *const TVMFFIByteArray,
+ out: *mut TVMFFIObjectHandle,
+ ) -> i32;
+ pub fn TVMFFITensorFromDLPack(
+ from: *mut c_void,
+ require_alignment: i32,
+ require_contiguous: i32,
+ out: *mut TVMFFIObjectHandle,
+ ) -> i32;
+ pub fn TVMFFITensorToDLPack(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
+ pub fn TVMFFITensorFromDLPackVersioned(
+ from: *mut c_void,
+ require_alignment: i32,
+ require_contiguous: i32,
+ out: *mut TVMFFIObjectHandle,
+ ) -> i32;
+ pub fn TVMFFITensorToDLPackVersioned(from: TVMFFIObjectHandle, out: *mut *mut c_void) -> i32;
+ pub fn TVMFFIStringFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
+ pub fn TVMFFIBytesFromByteArray(input: *const TVMFFIByteArray, out: *mut TVMFFIAny) -> i32;
+ pub fn TVMFFIDataTypeFromString(str: *const TVMFFIByteArray, out: *mut DLDataType) -> i32;
+ pub fn TVMFFIDataTypeToString(dtype: *const DLDataType, out: *mut TVMFFIAny) -> i32;
+ pub fn TVMFFITraceback(
+ filename: *const i8,
+ lineno: i32,
+ func: *const i8,
+ cross_ffi_boundary: i32,
+ ) -> *const TVMFFIByteArray;
+ pub fn TVMFFIGetTypeInfo(type_index: i32) -> *const TVMFFITypeInfo;
+}
diff --git a/rust/tvm-ffi-sys/src/c_env_api.rs b/rust/tvm-ffi-sys/src/c_env_api.rs
new file mode 100644
index 0000000..98b02ce
--- /dev/null
+++ b/rust/tvm-ffi-sys/src/c_env_api.rs
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+// NOTE: we manually write the C ABI as they are reasonably minimal
+// and we need to ensure clear control of the atomic access etc.
+#![allow(non_camel_case_types)]
+
+use std::ffi::c_void;
+use std::os::raw::c_char;
+
+use crate::c_api::TVMFFIObjectHandle;
+use crate::dlpack::DLTensor;
+
+// ----------------------------------------------------------------------------
+// Stream context
+// Focusing on minimalistic thread-local context recording stream being used.
+// We explicitly not handle allocation/de-allocation of stream here.
+// ----------------------------------------------------------------------------
+
+/// The type of the stream handle.
+pub type TVMFFIStreamHandle = *mut c_void;
+
+/// DLPack tensor allocator function type
+pub type DLPackTensorAllocator = unsafe extern "C" fn(
+ prototype: *mut DLTensor,
+ out: *mut *mut c_void, // DLManagedTensorVersioned**
+ error_ctx: *mut c_void,
+ set_error: unsafe extern "C" fn(*mut c_void, *const c_char, *const c_char),
+) -> i32;
+
+unsafe extern "C" {
+ pub fn TVMFFIEnvSetStream(
+ device_type: i32,
+ device_id: i32,
+ stream: TVMFFIStreamHandle,
+ opt_out_original_stream: *mut TVMFFIStreamHandle,
+ ) -> i32;
+
+ pub fn TVMFFIEnvGetStream(device_type: i32, device_id: i32) -> TVMFFIStreamHandle;
+
+ pub fn TVMFFIEnvSetTensorAllocator(
+ allocator: DLPackTensorAllocator,
+ write_to_global_context: i32,
+ opt_out_original_allocator: *mut DLPackTensorAllocator,
+ ) -> i32;
+
+ pub fn TVMFFIEnvGetTensorAllocator() -> DLPackTensorAllocator;
+
+ pub fn TVMFFIEnvCheckSignals() -> i32;
+
+ pub fn TVMFFIEnvRegisterCAPI(name: *const c_char, symbol: *mut c_void) -> i32;
+
+ pub fn TVMFFIEnvModLookupFromImports(
+ library_ctx: TVMFFIObjectHandle,
+ func_name: *const c_char,
+ out: *mut TVMFFIObjectHandle,
+ ) -> i32;
+
+ pub fn TVMFFIEnvModRegisterContextSymbol(name: *const c_char, symbol: *mut c_void) -> i32;
+
+ pub fn TVMFFIEnvModRegisterSystemLibSymbol(name: *const c_char, symbol: *mut c_void) -> i32;
+}
diff --git a/rust/tvm-ffi-sys/src/dlpack.rs b/rust/tvm-ffi-sys/src/dlpack.rs
new file mode 100644
index 0000000..e069ea8
--- /dev/null
+++ b/rust/tvm-ffi-sys/src/dlpack.rs
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+// DLPack C ABI declarations
+// NOTE: we manually write the C ABI as they are reasonably minimal
+// and we need to ensure clear control of the atomic access etc.
+#![allow(non_camel_case_types)]
+
+// DLPack related declarations
+#[repr(u32)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum DLDeviceType {
+ kDLCPU = 1,
+ kDLCUDA = 2,
+ kDLCUDAHost = 3,
+ kDLOpenCL = 4,
+ kDLVulkan = 7,
+ kDLMetal = 8,
+ kDLVPI = 9,
+ kDLROCM = 10,
+ kDLROCMHost = 11,
+ kDLExtDev = 12,
+ kDLCUDAManaged = 13,
+ kDLOneAPI = 14,
+ kDLWebGPU = 15,
+ kDLHexagon = 16,
+ kDLMAIA = 17,
+ kDLTrn = 18,
+}
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct DLDevice {
+ pub device_type: DLDeviceType,
+ pub device_id: i32,
+}
+
+/// DLPack data type code enum
+#[repr(u8)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum DLDataTypeCode {
+ kDLInt = 0,
+ kDLUInt = 1,
+ kDLFloat = 2,
+ kDLBfloat = 4,
+ kDLComplex = 5,
+ kDLOpaqueHandle = 3,
+ kDLBool = 6,
+ kDLFloat8_e3m4 = 7,
+ kDLFloat8_e4m3 = 8,
+ kDLFloat8_e4m3b11fnuz = 9,
+ kDLFloat8_e4m3fn = 10,
+ kDLFloat8_e4m3fnuz = 11,
+ kDLFloat8_e5m2 = 12,
+ kDLFloat8_e5m2fnuz = 13,
+ kDLFloat8_e8m0fnu = 14,
+ kDLFloat6_e2m3fn = 15,
+ kDLFloat6_e3m2fn = 16,
+ kDLFloat4_e2m1fn = 17,
+}
+
+/// DLPack data type struct
+#[repr(C)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct DLDataType {
+ pub code: u8,
+ pub bits: u8,
+ pub lanes: u16,
+}
+
+/// DLPack tensor struct - plain C tensor object, does not manage memory
+#[repr(C)]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DLTensor {
+ /// The data pointer points to the allocated data
+ pub data: *mut core::ffi::c_void,
+ /// The device of the tensor
+ pub device: DLDevice,
+ /// Number of dimensions
+ pub ndim: i32,
+ /// The data type of the pointer
+ pub dtype: DLDataType,
+ /// The shape of the tensor
+ pub shape: *mut i64,
+ /// Strides of the tensor (in number of elements, not bytes)
+ /// Can be NULL, indicating tensor is compact and row-majored
+ pub strides: *mut i64,
+ /// The offset in bytes to the beginning pointer to data
+ pub byte_offset: u64,
+}
+
+impl DLDevice {
+ pub fn new(device_type: DLDeviceType, device_id: i32) -> Self {
+ Self {
+ device_type: device_type,
+ device_id: device_id,
+ }
+ }
+}
+
+impl DLDataType {
+ pub fn new(code: DLDataTypeCode, bits: u8, lanes: u16) -> Self {
+ Self {
+ code: code as u8,
+ bits: bits,
+ lanes: lanes,
+ }
+ }
+}
diff --git a/rust/tvm-ffi-sys/src/lib.rs b/rust/tvm-ffi-sys/src/lib.rs
new file mode 100644
index 0000000..1530cc7
--- /dev/null
+++ b/rust/tvm-ffi-sys/src/lib.rs
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+pub mod c_api;
+pub mod c_env_api;
+pub mod dlpack;
+
+pub use crate::c_api::*;
+pub use crate::c_env_api::*;
+pub use crate::dlpack::*;
diff --git a/rust/tvm-ffi/Cargo.toml b/rust/tvm-ffi/Cargo.toml
new file mode 100644
index 0000000..b27c8c8
--- /dev/null
+++ b/rust/tvm-ffi/Cargo.toml
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+[package]
+name = "tvm-ffi"
+description = "tvm-ffi rust support"
+
+version = "0.1.0-alpha.0"
+edition = "2021"
+license = "Apache-2.0"
+
+
+[lib]
+name = "tvm_ffi"
+crate-type = ["lib"]
+
+[dependencies]
+paste = "1.0"
+tvm-ffi-sys = { version = "0.1.0-alpha.0", path = "../tvm-ffi-sys" }
+tvm-ffi-macros = { version = "0.1.0-alpha.0", path = "../tvm-ffi-macros" }
+
+# only needed for example building
+[features]
+example = []
+
+[[example]]
+name = "load_library"
+path = "examples/load_library.rs"
+required-features = ["example"]
diff --git a/rust/tvm-ffi/build.rs b/rust/tvm-ffi/build.rs
new file mode 100644
index 0000000..7209745
--- /dev/null
+++ b/rust/tvm-ffi/build.rs
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use std::env;
+use std::process::Command;
+
+fn generate_example_lib() {
+ let is_example_build = env::var("CARGO_FEATURE_EXAMPLE").is_ok();
+ if !is_example_build {
+ return;
+ }
+ println!("Running optional build step to generate example library");
+ let output_dir = env::var("OUT_DIR").unwrap();
+ let _ = Command::new("python")
+ .arg("scripts/generate_example_lib.py")
+ .arg(output_dir)
+ .output()
+ .expect("Failed to generate example library");
+ println!("cargo:rerun-if-changed=scripts/generate_example_lib.py");
+}
+
+/// Update the LD_LIBRARY_PATH environment variable
+/// so cargo run/test can directly pick it up
+/// note that it won't always work later consumption of the
+/// library, and we still need to figure out linking by setting LD_LIBRARY_PATH
+fn update_ld_library_path(lib_dir: &str) {
+ let os_env_var = match env::var("CARGO_CFG_TARGET_OS").as_deref() {
+ Ok("windows") => "PATH",
+ Ok("macos") => "DYLD_LIBRARY_PATH",
+ Ok("linux") => "LD_LIBRARY_PATH",
+ _ => "",
+ };
+ if os_env_var.is_empty() {
+ return;
+ }
+ // Get the current value of the environment variable at build time (if any)
+ let current_val = env::var(os_env_var).unwrap_or_else(|_| String::new());
+ // Use platform-specific separator
+ let separator = if os_env_var == "PATH" { ";" } else { ":" };
+ let new_ld_path = if current_val.is_empty() {
+ lib_dir.to_string()
+ } else {
+ format!("{}{}{}", current_val, separator, lib_dir)
+ };
+ // this env is only used for cargo run/test
+ println!("cargo:rustc-env={}={}", os_env_var, new_ld_path);
+}
+
+fn main() {
+ // Run `mylib-config --libdir` to get the library path
+ let config_output = Command::new("tvm-ffi-config")
+ .arg("--libdir")
+ .output()
+ .expect("Failed to run tvm-ffi-config");
+ let lib_dir = String::from_utf8(config_output.stdout)
+ .expect("Invalid UTF-8 output from tvm-ffi-config")
+ .trim()
+ .to_string();
+ // update the LD_LIBRARY_PATH environment variable
+ // note that we will also need to update ld_library_path for
+ // the cases here besides the tvm-ffi-sys crate so cargo test works out of the box
+ update_ld_library_path(&lib_dir);
+ // generate the example library
+ generate_example_lib();
+}
diff --git a/rust/tvm-ffi/examples/load_library.rs b/rust/tvm-ffi/examples/load_library.rs
new file mode 100644
index 0000000..47749e4
--- /dev/null
+++ b/rust/tvm-ffi/examples/load_library.rs
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::{Module, Result, Tensor};
+
+fn main() {
+ // lib_path is build using scripts/generate_example_lib.py [out_path]
+ // invoked automatically through build.rs
+ // you can replace it to any other ffi compatible library
+ let lib_path = concat!(env!("OUT_DIR"), "/add_one_cpu.so");
+ let lib: tvm_ffi::Module = Module::load_from_file(lib_path).unwrap();
+ // get the function from the library
+ let add_one_cpu: tvm_ffi::Function = lib.get_function("add_one_cpu").unwrap();
+ // load the function from the library
+ // and cast into a typed version of the function
+ let typed_add_one = tvm_ffi::into_typed_fn!(
+ add_one_cpu,
+ Fn(&Tensor, &Tensor) -> Result<()>
+ );
+ // run the function
+ let x_data: &[f32] = &[0.0, 1.0, 2.0, 3.0];
+ let x = Tensor::from_slice(x_data, &[4]).unwrap();
+ let y = Tensor::from_slice(&[0.0f32; 4], &[4]).unwrap();
+ println!("x: {:?}", x.data_as_slice::<f32>().unwrap());
+ typed_add_one(&x, &y).unwrap();
+ println!("y: {:?}", y.data_as_slice::<f32>().unwrap());
+}
diff --git a/rust/tvm-ffi/scripts/generate_example_lib.py b/rust/tvm-ffi/scripts/generate_example_lib.py
new file mode 100644
index 0000000..be38f17
--- /dev/null
+++ b/rust/tvm-ffi/scripts/generate_example_lib.py
@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Script to generate add_one_cpu.so used in examples."""
+
+import pathlib
+import shutil
+import sys
+
+import tvm_ffi.cpp
+
+
+def main() -> None:
+ """Generate the FFI library for examples."""
+ if len(sys.argv) != 2:
+ print("Usage: python generate_example_lib.py <output_dir>")
+ sys.exit(1)
+ output_dir = sys.argv[1]
+ output_lib_path = tvm_ffi.cpp.build_inline(
+ name="hello",
+ cpp_sources=r"""
+ void add_one_cpu(tvm::ffi::Tensor x, tvm::ffi::Tensor y) {
+ // implementation of a library function
+ TVM_FFI_ICHECK(x->ndim == 1) << "x must be a 1D tensor";
+ DLDataType f32_dtype{kDLFloat, 32, 1};
+ TVM_FFI_ICHECK(x->dtype == f32_dtype) << "x must be a float tensor";
+ TVM_FFI_ICHECK(y->ndim == 1) << "y must be a 1D tensor";
+ TVM_FFI_ICHECK(y->dtype == f32_dtype) << "y must be a float tensor";
+ TVM_FFI_ICHECK(x->shape[0] == y->shape[0]) << "x and y must have the same shape";
+ for (int i = 0; i < x->shape[0]; ++i) {
+ static_cast<float*>(y->data)[i] = static_cast<float*>(x->data)[i] + 1;
+ }
+ }
+ """,
+ functions=["add_one_cpu"],
+ )
+ target_lib_path = pathlib.Path(output_dir) / "add_one_cpu.so"
+ print(f"Generated FFI library at {target_lib_path}")
+ shutil.copy(output_lib_path, target_lib_path)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/rust/tvm-ffi/src/any.rs b/rust/tvm-ffi/src/any.rs
new file mode 100644
index 0000000..d5a4c85
--- /dev/null
+++ b/rust/tvm-ffi/src/any.rs
@@ -0,0 +1,362 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::error::Error;
+use crate::object;
+use crate::type_traits::AnyCompatible;
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIAnyViewToOwnedAny};
+
+/// Unmanaged Any that can hold reference to values
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub struct AnyView<'a> {
+ data: TVMFFIAny,
+ /// needs to explicit mark lifetime to avoid lifetime mismatch
+ _phantom: std::marker::PhantomData<&'a ()>,
+}
+
+/// Managed Any that can hold reference to values
+#[repr(C)]
+pub struct Any {
+ data: TVMFFIAny,
+}
+
+//---------------------
+// AnyView
+//---------------------
+impl<'a> AnyView<'a> {
+ pub fn new() -> Self {
+ Self {
+ data: TVMFFIAny::new(),
+ _phantom: std::marker::PhantomData,
+ }
+ }
+
+ #[inline]
+ pub fn type_index(&self) -> i32 {
+ self.data.type_index
+ }
+
+ /// More strict version than try_from/try_into
+ ///
+ /// This function will not try to cast the type
+ /// and ensures invariance that the return value is only Some(T)
+ /// Any::from(T) contains exactly the same value value
+ ///
+ /// will return Some(T) if the type is exactly compatible with T
+ /// will return None if the type is not compatible with T
+ #[inline]
+ pub fn try_as<T>(&self) -> Option<T>
+ where
+ T: AnyCompatible,
+ {
+ unsafe {
+ if T::check_any_strict(&self.data) {
+ Some(T::copy_from_any_view_after_check(&self.data))
+ } else {
+ None
+ }
+ }
+ }
+
+ /// Get the strong count of the underlying object for testing/debugging purposes
+ ///
+ /// If the underlying object is not ref counted, return None
+ pub fn debug_strong_count(&self) -> Option<usize> {
+ unsafe {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ Some(object::unsafe_::strong_count(self.data.data_union.v_obj))
+ } else {
+ None
+ }
+ }
+ }
+}
+
+impl<'a, T: AnyCompatible> From<&'a T> for AnyView<'a> {
+ #[inline]
+ fn from(value: &'a T) -> Self {
+ unsafe {
+ let mut data = TVMFFIAny::new();
+ T::copy_to_any_view(&value, &mut data);
+ Self {
+ data: data,
+ _phantom: std::marker::PhantomData,
+ }
+ }
+ }
+}
+
+impl Default for AnyView<'_> {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+/// Holder for Any value
+///
+/// This is used to define try_from rule while conforming to orphan rule
+/// Users should not use this directly
+pub struct TryFromTemp<T> {
+ value: T,
+}
+
+impl<T> TryFromTemp<T> {
+ /// Create a new holder for the value
+ #[inline(always)]
+ pub fn new(value: T) -> Self {
+ Self { value }
+ }
+
+ /// Move the value out of the holder
+ #[inline(always)]
+ pub fn into_value(this: Self) -> T {
+ this.value
+ }
+}
+
+//---------------------
+// Any
+//---------------------
+impl Any {
+ pub fn new() -> Self {
+ Self {
+ data: TVMFFIAny::new(),
+ }
+ }
+ #[inline]
+ pub fn type_index(&self) -> i32 {
+ self.data.type_index
+ }
+ /// Try to query if stored typed in Any exactly matches the type T
+ ///
+ /// This function is fast in the case of failure and can be used to check
+ /// if the type is compatible with T
+ ///
+ /// This function will not try to cast the type
+ /// and ensures invariance that the return value is only Some(T)
+ /// `Any::from(T)` contains exactly the same value value
+ /// `Any::try_as<T>()` contains exactly the same value value
+ ///
+ /// will return Some(T) if the type is exactly compatible with T
+ /// will return None if the type is not compatible with T
+ #[inline]
+ pub fn try_as<T>(&self) -> Option<T>
+ where
+ T: AnyCompatible,
+ {
+ unsafe {
+ if T::check_any_strict(&self.data) {
+ Some(T::copy_from_any_view_after_check(&self.data))
+ } else {
+ None
+ }
+ }
+ }
+
+ #[inline]
+ pub unsafe fn as_data_ptr(&mut self) -> *mut TVMFFIAny {
+ &mut self.data
+ }
+
+ #[inline]
+ pub unsafe fn into_raw_ffi_any(this: Self) -> TVMFFIAny {
+ this.data
+ }
+
+ #[inline]
+ pub unsafe fn from_raw_ffi_any(data: TVMFFIAny) -> Self {
+ Self { data }
+ }
+
+ /// Get the strong count of the underlying object for testing/debugging purposes
+ ///
+ /// If the underlying object is not ref counted, return None
+ pub fn debug_strong_count(&self) -> Option<usize> {
+ unsafe {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ Some(object::unsafe_::strong_count(self.data.data_union.v_obj))
+ } else {
+ None
+ }
+ }
+ }
+}
+
+impl Default for Any {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl Clone for Any {
+ #[inline]
+ fn clone(&self) -> Self {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { object::unsafe_::inc_ref(self.data.data_union.v_obj) }
+ }
+ Self { data: self.data }
+ }
+}
+
+impl Drop for Any {
+ #[inline]
+ fn drop(&mut self) {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { object::unsafe_::dec_ref(self.data.data_union.v_obj) }
+ }
+ }
+}
+
+// convert Any ref to AnyView
+impl<'a> From<&'a Any> for AnyView<'a> {
+ #[inline]
+ fn from(value: &'a Any) -> Self {
+ Self {
+ data: value.data,
+ _phantom: std::marker::PhantomData,
+ }
+ }
+}
+
+// convert AnyView to Any
+impl From<AnyView<'_>> for Any {
+ #[inline]
+ fn from(value: AnyView<'_>) -> Self {
+ unsafe {
+ let mut data = TVMFFIAny::new();
+ crate::check_safe_call!(TVMFFIAnyViewToOwnedAny(&value.data, &mut data)).unwrap();
+ Self { data }
+ }
+ }
+}
+
+impl<T: AnyCompatible> From<T> for Any {
+ #[inline]
+ fn from(value: T) -> Self {
+ unsafe {
+ let mut data = TVMFFIAny::new();
+ T::move_to_any(value, &mut data);
+ Self { data }
+ }
+ }
+}
+
+impl<'a, T: AnyCompatible> TryFrom<AnyView<'a>> for TryFromTemp<T> {
+ type Error = crate::error::Error;
+ #[inline]
+ fn try_from(value: AnyView<'a>) -> Result<Self, Self::Error> {
+ unsafe {
+ if T::check_any_strict(&value.data) {
+ Ok(TryFromTemp::new(T::copy_from_any_view_after_check(
+ &value.data,
+ )))
+ } else {
+ T::try_cast_from_any_view(&value.data)
+ .map_err(|_| {
+ let msg = format!(
+ "Cannot convert from type `{}` to `{}`",
+ T::get_mismatch_type_info(&value.data),
+ T::type_str()
+ );
+ crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
+ })
+ .map(TryFromTemp::new)
+ }
+ }
+ }
+}
+
+impl<T: AnyCompatible> TryFrom<Any> for TryFromTemp<T> {
+ type Error = crate::error::Error;
+ #[inline]
+ fn try_from(value: Any) -> Result<Self, Self::Error> {
+ unsafe {
+ if T::check_any_strict(&value.data) {
+ let mut value = std::mem::ManuallyDrop::new(value);
+ Ok(TryFromTemp::new(T::move_from_any_after_check(
+ &mut value.data,
+ )))
+ } else {
+ T::try_cast_from_any_view(&value.data)
+ .map_err(|_| {
+ let msg = format!(
+ "Cannot convert from type `{}` to `{}`",
+ T::get_mismatch_type_info(&value.data),
+ T::type_str()
+ );
+ crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
+ })
+ .map(TryFromTemp::new)
+ }
+ }
+ }
+}
+
+crate::impl_try_from_any!(
+ bool,
+ i8,
+ i16,
+ i32,
+ i64,
+ isize,
+ u8,
+ u16,
+ u32,
+ u64,
+ usize,
+ f32,
+ f64,
+ (),
+ *mut core::ffi::c_void,
+ crate::string::String,
+ crate::string::Bytes,
+ crate::object::ObjectRef,
+ tvm_ffi_sys::dlpack::DLDataType,
+ tvm_ffi_sys::dlpack::DLDevice,
+);
+
+crate::impl_try_from_any_for_parametric!(Option<T>);
+
+//------------------------------------------------------------
+/// ArgTryFromAnyView: Helper for function argument passing
+///-----------------------------------------------------------
+pub(crate) trait ArgTryFromAnyView: Sized {
+ fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self, Error>;
+}
+
+impl<T: AnyCompatible> ArgTryFromAnyView for T {
+ fn try_from_any_view(value: &AnyView, arg_index: usize) -> Result<Self, Error> {
+ unsafe {
+ if T::check_any_strict(&value.data) {
+ Ok(T::copy_from_any_view_after_check(&value.data))
+ } else {
+ T::try_cast_from_any_view(&value.data).map_err(|_| {
+ let msg = format!(
+ "Argument #{}: Cannot convert from type `{}` to `{}`",
+ arg_index,
+ T::get_mismatch_type_info(&value.data),
+ T::type_str()
+ );
+ crate::error::Error::new(crate::error::TYPE_ERROR, &msg, "")
+ })
+ }
+ }
+ }
+}
diff --git a/rust/tvm-ffi/src/collections/mod.rs b/rust/tvm-ffi/src/collections/mod.rs
new file mode 100644
index 0000000..85635a7
--- /dev/null
+++ b/rust/tvm-ffi/src/collections/mod.rs
@@ -0,0 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+pub mod shape;
+/// Collection types
+pub mod tensor;
diff --git a/rust/tvm-ffi/src/collections/shape.rs b/rust/tvm-ffi/src/collections/shape.rs
new file mode 100644
index 0000000..39d9a1d
--- /dev/null
+++ b/rust/tvm-ffi/src/collections/shape.rs
@@ -0,0 +1,139 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::derive::{Object, ObjectRef};
+use crate::object::{Object, ObjectArc, ObjectCoreWithExtraItems};
+use std::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
+use std::ops::Deref;
+use tvm_ffi_sys::TVMFFIShapeCell;
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+
+//-----------------------------------------------------
+// Shape
+//-----------------------------------------------------
+// ShapeObj for heap-allocated shape
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Shape"]
+#[type_index(TypeIndex::kTVMFFIShape)]
+pub struct ShapeObj {
+ object: Object,
+ data: TVMFFIShapeCell,
+}
+
+/// ABI stable owned Shape for ffi
+#[repr(C)]
+#[derive(ObjectRef, Clone)]
+pub struct Shape {
+ data: ObjectArc<ShapeObj>,
+}
+
+impl Shape {
+ /// Create a new empty Shape
+ pub fn new() -> Self {
+ let shape_obj = ShapeObj {
+ object: Object::new(),
+ data: TVMFFIShapeCell {
+ data: std::ptr::null(),
+ size: 0,
+ },
+ };
+ Self {
+ data: ObjectArc::new(shape_obj),
+ }
+ }
+
+ /// Get the shape as a slice
+ pub fn as_slice(&self) -> &[i64] {
+ unsafe { std::slice::from_raw_parts(self.data.data.data, self.data.data.size) }
+ }
+
+ /// Fill the strides from the shape
+ pub fn fill_strides_from_shape<T>(shape: T, strides: &mut [i64])
+ where
+ T: AsRef<[i64]>,
+ {
+ let shape = shape.as_ref();
+ let mut stride = 1;
+ for i in (0..shape.len()).rev() {
+ strides[i] = stride;
+ stride *= shape[i];
+ }
+ }
+}
+
+unsafe impl ObjectCoreWithExtraItems for ShapeObj {
+ type ExtraItem = i64;
+ fn extra_items_count(this: &Self) -> usize {
+ this.data.size
+ }
+}
+
+impl<T> From<T> for Shape
+where
+ T: AsRef<[i64]>,
+{
+ fn from(value: T) -> Self {
+ unsafe {
+ let value_slice: &[i64] = value.as_ref();
+ let mut obj_arc = ObjectArc::new_with_extra_items(ShapeObj {
+ object: Object::new(),
+ data: TVMFFIShapeCell {
+ data: std::ptr::null(),
+ size: value_slice.len(),
+ },
+ });
+ // reset the data ptr correctly after Arc is created
+ obj_arc.data.data = ShapeObj::extra_items(&obj_arc).as_ptr();
+ let extra_items = ShapeObj::extra_items_mut(&mut obj_arc);
+ extra_items.copy_from_slice(value_slice);
+ Self { data: obj_arc }
+ }
+ }
+}
+
+impl Deref for Shape {
+ type Target = [i64];
+ #[inline]
+ fn deref(&self) -> &[i64] {
+ self.as_slice()
+ }
+}
+
+impl PartialEq for Shape {
+ #[inline]
+ fn eq(&self, other: &Self) -> bool {
+ self.as_slice() == other.as_slice()
+ }
+}
+
+impl Eq for Shape {}
+
+impl PartialOrd for Shape {
+ #[inline]
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ self.as_slice().partial_cmp(other.as_slice())
+ }
+}
+
+impl Ord for Shape {
+ #[inline]
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.as_slice().cmp(other.as_slice())
+ }
+}
diff --git a/rust/tvm-ffi/src/collections/tensor.rs b/rust/tvm-ffi/src/collections/tensor.rs
new file mode 100644
index 0000000..6b34613
--- /dev/null
+++ b/rust/tvm-ffi/src/collections/tensor.rs
@@ -0,0 +1,338 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::collections::shape::Shape;
+use crate::derive::{Object, ObjectRef};
+use crate::dtype::AsDLDataType;
+use crate::dtype::DLDataTypeExt;
+use crate::error::Result;
+use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems};
+use tvm_ffi_sys::dlpack::{DLDataType, DLDevice, DLDeviceType, DLTensor};
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+
+//-----------------------------------------------------
+// NDAllocator Trait
+//-----------------------------------------------------
+/// Trait for n-dimensional array allocators
+pub unsafe trait NDAllocator: 'static {
+ /// The minimum alignment of the data allocated by the allocator
+ const MIN_ALIGN: usize;
+ /// Allocate data for the given DLTensor
+ ///
+ /// # Arguments
+ /// * `tensor` - The DLTensor to allocate data for
+ ///
+ /// This method should fill in the data pointer of the DLTensor.
+ unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut core::ffi::c_void;
+
+ /// Free data for the given DLTensor
+ ///
+ /// # Arguments
+ /// * `tensor` - The DLTensor to free data for
+ ///
+ /// This method should free the data pointer of the DLTensor.
+ unsafe fn free_data(&mut self, tensor: &DLTensor);
+}
+
+/// DLTensorExt trait
+/// This trait provides methods to get the number of elements and the item size of a DLTensor
+pub trait DLTensorExt {
+ fn numel(&self) -> usize;
+ fn item_size(&self) -> usize;
+}
+
+impl DLTensorExt for DLTensor {
+ fn numel(&self) -> usize {
+ unsafe {
+ std::slice::from_raw_parts(self.shape, self.ndim as usize)
+ .iter()
+ .product::<i64>() as usize
+ }
+ }
+
+ fn item_size(&self) -> usize {
+ (self.dtype.bits as usize * self.dtype.lanes as usize + 7) / 8
+ }
+}
+
+//-----------------------------------------------------
+// Shape
+//-----------------------------------------------------
+// ShapeObj for heap-allocated shape
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Tensor"]
+#[type_index(TypeIndex::kTVMFFITensor)]
+pub struct TensorObj {
+ object: Object,
+ dltensor: DLTensor,
+}
+
+/// ABI stable owned Shape for ffi
+#[repr(C)]
+#[derive(ObjectRef, Clone)]
+pub struct Tensor {
+ data: ObjectArc<TensorObj>,
+}
+
+impl Tensor {
+ /// Get the data pointer of the Tensor
+ ///
+ /// # Returns
+ /// * `*mut core::ffi::c_void` - The data pointer of the Tensor
+ pub fn data_ptr(&self) -> *const core::ffi::c_void {
+ self.data.dltensor.data
+ }
+ /// Get the data pointer of the Tensor
+ ///
+ /// # Returns
+ /// * `*mut core::ffi::c_void` - The data pointer of the Tensor
+ pub fn data_ptr_mut(&mut self) -> *mut core::ffi::c_void {
+ self.data.dltensor.data
+ }
+ /// Check if the Tensor is contiguous
+ ///
+ /// # Returns
+ /// * `bool` - True if the Tensor is contiguous, false otherwise
+ pub fn is_contiguous(&self) -> bool {
+ let strides = self.strides();
+ let shape = self.shape();
+ let mut expected_stride = 1;
+ for i in (0..self.ndim()).rev() {
+ if strides[i] != expected_stride {
+ return false;
+ }
+ expected_stride *= shape[i];
+ }
+ true
+ }
+
+ pub fn data_as_slice<T: AsDLDataType>(&self) -> Result<&[T]> {
+ let dtype = T::DL_DATA_TYPE;
+ if self.dtype() != dtype {
+ crate::bail!(
+ crate::error::TYPE_ERROR,
+ "Data type mismatch {} vs {}",
+ self.dtype().to_string(),
+ dtype.to_string()
+ );
+ }
+ if self.device().device_type != DLDeviceType::kDLCPU {
+ crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU");
+ }
+ crate::ensure!(
+ self.is_contiguous(),
+ crate::error::RUNTIME_ERROR,
+ "Tensor is not contiguous"
+ );
+
+ unsafe {
+ Ok(std::slice::from_raw_parts(
+ self.data.dltensor.data as *const T,
+ self.numel(),
+ ))
+ }
+ }
+ /// Get the data as a mutable slice
+ ///
+ /// Note that we do allow mutable data access to copies of the Tensor,
+ /// as in the case of low-level deep learning frameworks.
+ ///
+ /// # Arguments
+ /// * `T` - The type of the data
+ ///
+ /// # Returns
+ /// * `Result<&mut [T]>` - The data as a mutable slice
+ pub fn data_as_slice_mut<T: AsDLDataType>(&self) -> Result<&mut [T]> {
+ let dtype = T::DL_DATA_TYPE;
+ if self.dtype() != dtype {
+ crate::bail!(
+ crate::error::TYPE_ERROR,
+ "Data type mismatch: expected {}, got {}",
+ dtype.to_string(),
+ self.dtype().to_string()
+ );
+ }
+ if self.device().device_type != DLDeviceType::kDLCPU {
+ crate::bail!(crate::error::RUNTIME_ERROR, "Tensor is not on CPU");
+ }
+ crate::ensure!(
+ self.is_contiguous(),
+ crate::error::RUNTIME_ERROR,
+ "Tensor is not contiguous"
+ );
+ unsafe {
+ Ok(std::slice::from_raw_parts_mut(
+ self.data.dltensor.data as *mut T,
+ self.numel(),
+ ))
+ }
+ }
+
+ pub fn shape(&self) -> &[i64] {
+ unsafe { std::slice::from_raw_parts(self.data.dltensor.shape, self.ndim()) }
+ }
+
+ pub fn ndim(&self) -> usize {
+ self.data.dltensor.ndim as usize
+ }
+
+ pub fn numel(&self) -> usize {
+ self.data.dltensor.numel()
+ }
+
+ pub fn strides(&self) -> &[i64] {
+ unsafe { std::slice::from_raw_parts(self.data.dltensor.strides, self.ndim()) }
+ }
+
+ pub fn dtype(&self) -> DLDataType {
+ self.data.dltensor.dtype
+ }
+
+ pub fn device(&self) -> DLDevice {
+ self.data.dltensor.device
+ }
+}
+
+struct TensorObjFromNDAlloc<TNDAlloc>
+where
+ TNDAlloc: NDAllocator,
+{
+ base: TensorObj,
+ alloc: TNDAlloc,
+}
+
+unsafe impl<TNDAlloc: NDAllocator> ObjectCore for TensorObjFromNDAlloc<TNDAlloc> {
+ const TYPE_KEY: &'static str = TensorObj::TYPE_KEY;
+ fn type_index() -> i32 {
+ TensorObj::type_index()
+ }
+ unsafe fn object_header_mut(this: &mut Self) -> &mut tvm_ffi_sys::TVMFFIObject {
+ TensorObj::object_header_mut(&mut this.base)
+ }
+}
+
+unsafe impl<TNDAlloc: NDAllocator> ObjectCoreWithExtraItems for TensorObjFromNDAlloc<TNDAlloc> {
+ type ExtraItem = i64;
+ fn extra_items_count(this: &Self) -> usize {
+ (this.base.dltensor.ndim * 2) as usize
+ }
+}
+
+impl<TNDAlloc: NDAllocator> Drop for TensorObjFromNDAlloc<TNDAlloc> {
+ fn drop(&mut self) {
+ unsafe {
+ self.alloc.free_data(&self.base.dltensor);
+ }
+ }
+}
+
+impl Tensor {
+ // Create a Tensor from a NDAllocator
+ ///
+ /// # Arguments
+ /// * `alloc` - The NDAllocator
+ /// * `shape` - The shape of the Tensor
+ /// * `dtype` - The data type of the Tensor
+ /// * `device` - The device of the Tensor
+ ///
+ /// # Returns
+ /// * `Tensor` - The created Tensor
+ pub fn from_nd_alloc<TNDAlloc>(
+ alloc: TNDAlloc,
+ shape: &[i64],
+ dtype: DLDataType,
+ device: DLDevice,
+ ) -> Self
+ where
+ TNDAlloc: NDAllocator,
+ {
+ let tensor_obj = TensorObjFromNDAlloc {
+ base: TensorObj {
+ object: Object::new(),
+ dltensor: DLTensor {
+ data: std::ptr::null_mut(),
+ device: device,
+ ndim: shape.len() as i32,
+ dtype: dtype,
+ shape: std::ptr::null_mut(),
+ strides: std::ptr::null_mut(),
+ byte_offset: 0,
+ },
+ },
+ alloc: alloc,
+ };
+ unsafe {
+ let mut obj_arc = ObjectArc::new_with_extra_items(tensor_obj);
+ obj_arc.base.dltensor.shape =
+ TensorObjFromNDAlloc::extra_items(&obj_arc).as_ptr() as *mut i64;
+ obj_arc.base.dltensor.strides = obj_arc.base.dltensor.shape.add(shape.len());
+ let extra_items = TensorObjFromNDAlloc::extra_items_mut(&mut obj_arc);
+ extra_items[..shape.len()].copy_from_slice(shape);
+ Shape::fill_strides_from_shape(shape, &mut extra_items[shape.len()..]);
+ let dltensor_ptr = &obj_arc.base.dltensor as *const DLTensor;
+ obj_arc.base.dltensor.data = obj_arc.alloc.alloc_data(&*dltensor_ptr);
+ Self {
+ data: ObjectArc::from_raw(ObjectArc::into_raw(obj_arc) as *mut TensorObj),
+ }
+ }
+ }
+ /// Create a Tensor from a slice
+ ///
+ /// # Arguments
+ /// * `slice` - The slice to create the Tensor from
+ /// * `shape` - The shape of the Tensor
+ ///
+ /// # Returns
+ /// * `Tensor` - The created Tensor
+ pub fn from_slice<T: AsDLDataType>(slice: &[T], shape: &[i64]) -> Result<Self> {
+ let dtype = T::DL_DATA_TYPE;
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, shape, dtype, device);
+ if tensor.numel() != slice.len() {
+ crate::bail!(crate::error::VALUE_ERROR, "Slice length mismatch");
+ }
+ tensor.data_as_slice_mut::<T>()?.copy_from_slice(slice);
+ Ok(tensor)
+ }
+}
+
+/// Example CPU NDAllocator
+/// This allocator allocates data on the CPU
+pub struct CPUNDAlloc {}
+
+unsafe impl NDAllocator for CPUNDAlloc {
+ const MIN_ALIGN: usize = 64;
+
+ unsafe fn alloc_data(&mut self, prototype: &DLTensor) -> *mut core::ffi::c_void {
+ let numel = prototype.numel() as usize;
+ let item_size = prototype.item_size();
+ let size = numel * item_size as usize;
+ let layout = std::alloc::Layout::from_size_align(size, Self::MIN_ALIGN).unwrap();
+ let ptr = std::alloc::alloc(layout);
+ ptr as *mut core::ffi::c_void
+ }
+
+ unsafe fn free_data(&mut self, tensor: &DLTensor) {
+ let numel = tensor.numel() as usize;
+ let item_size = tensor.item_size();
+ let size = numel * item_size;
+ let layout = std::alloc::Layout::from_size_align(size, Self::MIN_ALIGN).unwrap();
+ std::alloc::dealloc(tensor.data as *mut u8, layout);
+ }
+}
diff --git a/rust/tvm-ffi/src/derive.rs b/rust/tvm-ffi/src/derive.rs
new file mode 100644
index 0000000..629b949
--- /dev/null
+++ b/rust/tvm-ffi/src/derive.rs
@@ -0,0 +1,20 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/// namespace to re-export derive macros
+pub use tvm_ffi_macros::{Object, ObjectRef};
diff --git a/rust/tvm-ffi/src/device.rs b/rust/tvm-ffi/src/device.rs
new file mode 100644
index 0000000..d6c0418
--- /dev/null
+++ b/rust/tvm-ffi/src/device.rs
@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::error::Result;
+use crate::type_traits::AnyCompatible;
+use tvm_ffi_sys::dlpack::DLDevice;
+use tvm_ffi_sys::{TVMFFIAny, TVMFFITypeIndex as TypeIndex};
+use tvm_ffi_sys::{TVMFFIEnvGetStream, TVMFFIEnvSetStream, TVMFFIStreamHandle};
+
+/// Get the current stream for a device
+pub fn current_stream(device: &DLDevice) -> TVMFFIStreamHandle {
+ unsafe { TVMFFIEnvGetStream(device.device_type as i32, device.device_id) }
+}
+/// call f with the stream set to the given stream
+pub fn with_stream<T>(
+ device: &DLDevice,
+ stream: TVMFFIStreamHandle,
+ f: impl FnOnce() -> Result<T>,
+) -> Result<T> {
+ let mut prev_stream: TVMFFIStreamHandle = std::ptr::null_mut();
+ unsafe {
+ crate::check_safe_call!(TVMFFIEnvSetStream(
+ device.device_type as i32,
+ device.device_id,
+ stream,
+ &mut prev_stream as *mut TVMFFIStreamHandle
+ ))?;
+ }
+ let result = f()?;
+ unsafe {
+ crate::check_safe_call!(TVMFFIEnvSetStream(
+ device.device_type as i32,
+ device.device_id,
+ prev_stream,
+ std::ptr::null_mut()
+ ))?;
+ }
+ Ok(result)
+}
+
+/// AnyCompatible for DLDevice
+unsafe impl AnyCompatible for DLDevice {
+ fn type_str() -> String {
+ // make it consistent with c++ representation
+ "Device".to_string()
+ }
+
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIDevice as i32;
+ data.small_str_len = 0;
+ data.data_union.v_uint64 = 0;
+ data.data_union.v_device = *src;
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIDevice as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = 0;
+ data.data_union.v_device = src;
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ return data.type_index == TypeIndex::kTVMFFIDevice as i32;
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ data.data_union.v_device
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ data.data_union.v_device
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIDevice as i32 {
+ Ok(data.data_union.v_device)
+ } else {
+ Err(())
+ }
+ }
+}
diff --git a/rust/tvm-ffi/src/dtype.rs b/rust/tvm-ffi/src/dtype.rs
new file mode 100644
index 0000000..eb471b2
--- /dev/null
+++ b/rust/tvm-ffi/src/dtype.rs
@@ -0,0 +1,209 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::error::Result;
+use crate::type_traits::AnyCompatible;
+/// Data type handling
+use tvm_ffi_sys::dlpack::{DLDataType, DLDataTypeCode};
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIByteArray, TVMFFIDataTypeFromString, TVMFFIDataTypeToString};
+
+/// Extra methods for DLDataType
+pub trait DLDataTypeExt: Sized {
+ /// Convert the DLDataType to a string representation
+ ///
+ /// # Returns
+ /// A string representation of the data type (e.g., "int32", "float64", "bool")
+ fn to_string(&self) -> crate::string::String;
+
+ /// Parse a string representation into a DLDataType
+ ///
+ /// # Arguments
+ /// * `dtype_str` - The string representation of the data type to parse
+ ///
+ /// # Returns
+ /// * `Ok(DLDataType)` - Successfully parsed data type
+ /// * `Err(Error)` - Failed to parse the string
+ ///
+ /// # Examples
+ /// ```
+ /// use tvm_ffi::{DLDataType, DLDataTypeExt};
+ ///
+ /// let dtype = DLDataType::try_from_str("int32").unwrap();
+ /// ```
+ fn try_from_str(dtype_str: &str) -> Result<Self>;
+}
+
+impl DLDataTypeExt for DLDataType {
+ fn to_string(&self) -> crate::string::String {
+ unsafe {
+ let mut ffi_any = TVMFFIAny::new();
+ crate::check_safe_call!(TVMFFIDataTypeToString(&*self, &mut ffi_any)).unwrap();
+ crate::any::Any::from_raw_ffi_any(ffi_any)
+ .try_into()
+ .unwrap()
+ }
+ }
+
+ fn try_from_str(dtype_str: &str) -> Result<Self> {
+ let mut dtype = DLDataType {
+ code: DLDataTypeCode::kDLOpaqueHandle as u8,
+ bits: 0,
+ lanes: 0,
+ };
+ unsafe {
+ let dtype_byte_array = TVMFFIByteArray::from_str(dtype_str);
+ crate::check_safe_call!(TVMFFIDataTypeFromString(&dtype_byte_array, &mut dtype))?;
+ }
+ Ok(dtype)
+ }
+}
+
+/// AnyCompatible implementation for DLDataType
+///
+/// This implementation allows DLDataType to be used with the TVM FFI Any system,
+/// enabling type-safe conversion between DLDataType and the generic Any type.
+unsafe impl AnyCompatible for DLDataType {
+ /// Get the type string identifier for DLDataType
+ ///
+ /// # Returns
+ /// The string "DataType" to match the C++ representation
+ fn type_str() -> String {
+ // make it consistent with c++ representation
+ "DataType".to_string()
+ }
+
+ /// Copy a DLDataType to an Any view
+ ///
+ /// # Arguments
+ /// * `src` - The DLDataType to copy from
+ /// * `data` - The Any view to copy to
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIDataType as i32;
+ data.small_str_len = 0;
+ data.data_union.v_uint64 = 0;
+ data.data_union.v_dtype = *src;
+ }
+
+ /// Move a DLDataType into an Any
+ ///
+ /// # Arguments
+ /// * `src` - The DLDataType to move from
+ /// * `data` - The Any to move into
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIDataType as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = 0;
+ data.data_union.v_dtype = src;
+ }
+
+ /// Check if an Any contains a DLDataType
+ ///
+ /// # Arguments
+ /// * `data` - The Any to check
+ ///
+ /// # Returns
+ /// `true` if the Any contains a DLDataType, `false` otherwise
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ return data.type_index == TypeIndex::kTVMFFIDataType as i32;
+ }
+
+ /// Copy a DLDataType from an Any view (after type check)
+ ///
+ /// # Arguments
+ /// * `data` - The Any view to copy from
+ ///
+ /// # Returns
+ /// The copied DLDataType
+ ///
+ /// # Safety
+ /// The caller must ensure that `data` contains a DLDataType
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ data.data_union.v_dtype
+ }
+
+ /// Move a DLDataType from an Any (after type check)
+ ///
+ /// # Arguments
+ /// * `data` - The Any to move from
+ ///
+ /// # Returns
+ /// The moved DLDataType
+ ///
+ /// # Safety
+ /// The caller must ensure that `data` contains a DLDataType
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ data.data_union.v_dtype
+ }
+
+ /// Try to cast an Any view to a DLDataType
+ ///
+ /// This method supports both direct DLDataType conversion and string parsing.
+ ///
+ /// # Arguments
+ /// * `data` - The Any view to cast from
+ ///
+ /// # Returns
+ /// * `Ok(DLDataType)` - Successfully cast to DLDataType
+ /// * `Err(())` - Failed to cast (wrong type or invalid string)
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIDataType as i32 {
+ Ok(data.data_union.v_dtype)
+ } else if let Ok(string) = crate::string::String::try_cast_from_any_view(data) {
+ DLDataType::try_from_str(string.as_str()).map_err(|_| ())
+ } else {
+ Err(())
+ }
+ }
+}
+
+/// Trait to convert standard data types to DLDataType
+///
+/// This trait provides a way to get the corresponding DLDataType for standard Rust types.
+/// It's implemented for common integer, unsigned integer, and floating-point types.
+pub trait AsDLDataType: Copy {
+ /// The corresponding DLDataType for this type
+ const DL_DATA_TYPE: DLDataType;
+}
+
+/// Macro to implement AsDLDataType for standard types
+///
+/// This macro generates implementations of the AsDLDataType trait for standard Rust types.
+/// It takes the type, the DLPack data type code, and the number of bits.
+macro_rules! impl_as_dl_data_type {
+ ($type: ty, $code: expr, $bits: expr) => {
+ impl AsDLDataType for $type {
+ const DL_DATA_TYPE: DLDataType = DLDataType {
+ code: $code as u8,
+ bits: $bits as u8,
+ lanes: 1,
+ };
+ }
+ };
+}
+
+impl_as_dl_data_type!(i8, DLDataTypeCode::kDLInt, 8);
+impl_as_dl_data_type!(i16, DLDataTypeCode::kDLInt, 16);
+impl_as_dl_data_type!(i32, DLDataTypeCode::kDLInt, 32);
+impl_as_dl_data_type!(i64, DLDataTypeCode::kDLInt, 64);
+impl_as_dl_data_type!(u8, DLDataTypeCode::kDLUInt, 8);
+impl_as_dl_data_type!(u16, DLDataTypeCode::kDLUInt, 16);
+impl_as_dl_data_type!(u32, DLDataTypeCode::kDLUInt, 32);
+impl_as_dl_data_type!(u64, DLDataTypeCode::kDLUInt, 64);
+impl_as_dl_data_type!(f32, DLDataTypeCode::kDLFloat, 32);
+impl_as_dl_data_type!(f64, DLDataTypeCode::kDLFloat, 64);
diff --git a/rust/tvm-ffi/src/error.rs b/rust/tvm-ffi/src/error.rs
new file mode 100644
index 0000000..cb6dabb
--- /dev/null
+++ b/rust/tvm-ffi/src/error.rs
@@ -0,0 +1,205 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::derive::{Object, ObjectRef};
+use crate::object::{Object, ObjectArc};
+use std::ffi::c_void;
+use tvm_ffi_sys::TVMFFIBacktraceUpdateMode::kTVMFFIBacktraceUpdateModeAppend;
+use tvm_ffi_sys::{
+ TVMFFIByteArray, TVMFFIErrorCell, TVMFFIErrorCreate, TVMFFIErrorMoveFromRaised,
+ TVMFFIErrorSetRaised, TVMFFIObjectHandle, TVMFFITypeIndex,
+};
+
+/// Error kind, wraps in a struct to be explicit
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ErrorKind<'a>(&'a str);
+
+impl<'a> ErrorKind<'a> {
+ pub fn as_str(&self) -> &str {
+ self.0
+ }
+}
+
+impl<'a> std::fmt::Display for ErrorKind<'a> {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "{}", self.0)
+ }
+}
+
+pub const VALUE_ERROR: ErrorKind = ErrorKind("ValueError");
+pub const TYPE_ERROR: ErrorKind = ErrorKind("TypeError");
+pub const RUNTIME_ERROR: ErrorKind = ErrorKind("RuntimeError");
+pub const ATTRIBUTE_ERROR: ErrorKind = ErrorKind("AttributeError");
+pub const KEY_ERROR: ErrorKind = ErrorKind("KeyError");
+pub const INDEX_ERROR: ErrorKind = ErrorKind("IndexError");
+
+/// error object
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Error"]
+#[type_index(TVMFFITypeIndex::kTVMFFIError)]
+pub struct ErrorObj {
+ object: Object,
+ cell: TVMFFIErrorCell,
+}
+
+/// Error reference class
+#[derive(Clone, ObjectRef)]
+pub struct Error {
+ data: ObjectArc<ErrorObj>,
+}
+
+/// Default result that uses Error as the error type
+pub type Result<T, E = Error> = std::result::Result<T, E>;
+
+impl Error {
+ pub fn new(kind: ErrorKind<'_>, message: &str, traceback: &str) -> Self {
+ unsafe {
+ let kind_data = TVMFFIByteArray::from_str(kind.as_str());
+ let message_data = TVMFFIByteArray::from_str(message);
+ let traceback_data = TVMFFIByteArray::from_str(traceback);
+ let mut error_handle: TVMFFIObjectHandle = std::ptr::null_mut();
+ let ret = TVMFFIErrorCreate(
+ &kind_data,
+ &message_data,
+ &traceback_data,
+ &mut error_handle,
+ );
+ assert_eq!(ret, 0, "Failed to create error object");
+ let error_obj = ObjectArc::from_raw(error_handle as *const ErrorObj);
+ Self { data: error_obj }
+ }
+ }
+
+ /// Create a new error by moving from raised error
+ ///
+ /// # Returns
+ /// The error from the raised error
+ pub fn from_raised() -> Self {
+ unsafe {
+ let mut error_handle: TVMFFIObjectHandle = std::ptr::null_mut();
+ TVMFFIErrorMoveFromRaised(&mut error_handle as *mut TVMFFIObjectHandle);
+ assert!(
+ !error_handle.is_null(),
+ "Calling Error::from_raised but no error was raised"
+ );
+ let error_obj = ObjectArc::from_raw(error_handle as *const ErrorObj);
+ Self { data: error_obj }
+ }
+ }
+
+ /// Set the error as raised
+ ///
+ /// # Arguments
+ /// * `error` - The error to set as raised
+ pub fn set_raised(error: &Self) {
+ unsafe {
+ TVMFFIErrorSetRaised(ObjectArc::as_raw(&error.data) as TVMFFIObjectHandle);
+ }
+ }
+
+ /// Get the kind of the error
+ ///
+ /// # Returns
+ /// The kind of the error
+ pub fn kind(&self) -> ErrorKind<'_> {
+ ErrorKind(&self.data.cell.kind.as_str())
+ }
+
+ /// Get the message of the error
+ ///
+ /// # Returns
+ /// The message of the error
+ pub fn message(&self) -> &str {
+ self.data.cell.message.as_str()
+ }
+
+ /// Get the backtrace of the error
+ ///
+ /// # Returns
+ /// The backtrace of the error
+ pub fn backtrace(&self) -> &str {
+ self.data.cell.backtrace.as_str()
+ }
+
+ /// Get the traceback of the error in the order of most recent call last
+ ///
+ /// # Returns
+ /// The traceback of the error
+ pub fn traceback_most_recent_call_last(&self) -> String {
+ let backtrace = self.backtrace();
+ let backtrace_lines = backtrace.split('\n');
+ let mut traceback = String::new();
+ for line in backtrace_lines.rev() {
+ traceback.push_str(line);
+ traceback.push('\n');
+ }
+ traceback
+ }
+
+ /// Append the backtrace to the error
+ ///
+ /// # Arguments
+ /// * `this` - The error to append the backtrace to
+ /// * `backtrace` - The backtrace to append
+ ///
+ /// # Returns
+ /// The error with the appended backtrace
+ pub fn with_appended_backtrace(this: Self, backtrace: &str) -> Self {
+ if ObjectArc::strong_count(&this.data) == 1 {
+ // this is the only reference to the error
+ // we can safely mutate the error
+ unsafe {
+ let backtrace_data = TVMFFIByteArray::from_str(backtrace);
+ (this.data.cell.update_backtrace)(
+ ObjectArc::as_raw(&this.data) as *mut ErrorObj as *mut c_void,
+ &backtrace_data,
+ kTVMFFIBacktraceUpdateModeAppend as i32,
+ );
+ this
+ }
+ } else {
+ // we need to create a new error because there is more than one unique reference
+ // to the error
+ let mut new_backtrace = String::new();
+ new_backtrace.push_str(this.backtrace());
+ new_backtrace.push_str(backtrace);
+ return Error::new(this.kind(), this.message(), &new_backtrace);
+ }
+ }
+}
+
+impl std::fmt::Display for Error {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "Traceback (most recent call last):\n{}{}: {}",
+ self.traceback_most_recent_call_last(),
+ self.kind().as_str(),
+ self.message()
+ )
+ }
+}
+
+impl std::fmt::Debug for Error {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ std::fmt::Display::fmt(self, f)
+ }
+}
+
+impl std::error::Error for Error {}
diff --git a/rust/tvm-ffi/src/extra/mod.rs b/rust/tvm-ffi/src/extra/mod.rs
new file mode 100644
index 0000000..2a4c01d
--- /dev/null
+++ b/rust/tvm-ffi/src/extra/mod.rs
@@ -0,0 +1,19 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+pub mod module;
diff --git a/rust/tvm-ffi/src/extra/module.rs b/rust/tvm-ffi/src/extra/module.rs
new file mode 100644
index 0000000..e3009e2
--- /dev/null
+++ b/rust/tvm-ffi/src/extra/module.rs
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::derive::{Object, ObjectRef};
+use crate::error::Result;
+use crate::function::Function;
+use crate::object::{Object, ObjectArc};
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+
+//-----------------------------------------------------
+// Module
+//-----------------------------------------------------
+
+/// A TVM FFI Module for loading dynamic libraries and retrieving functions.
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Module"]
+#[type_index(TypeIndex::kTVMFFIModule)]
+pub struct ModuleObj {
+ object: Object,
+}
+
+/// ABI-stable owned Module for FFI operations.
+#[repr(C)]
+#[derive(ObjectRef, Clone)]
+pub struct Module {
+ data: ObjectArc<ModuleObj>,
+}
+
+impl Module {
+ /// Load a module from a dynamic library file.
+ ///
+ /// # Arguments
+ /// * `file_name` - Path to the dynamic library file to load
+ ///
+ /// # Returns
+ /// * `Result<Module>` - A `Module` instance on success
+ pub fn load_from_file<Str: AsRef<str>>(file_name: Str) -> Result<Module> {
+ static API_FUNC: std::sync::LazyLock<Function> =
+ std::sync::LazyLock::new(|| Function::get_global("ffi.ModuleLoadFromFile").unwrap());
+ let file_name = crate::string::String::from(file_name);
+ (*API_FUNC)
+ .call_tuple_with_len::<1, _>((file_name,))?
+ .try_into()
+ }
+
+ /// Get a function from the module by name.
+ ///
+ /// # Arguments
+ /// * `name` - The name of the function to retrieve
+ ///
+ /// # Returns
+ /// * `Result<Function>` - A `Function` instance on success
+ pub fn get_function<Str: AsRef<str>>(&self, name: Str) -> Result<Function> {
+ static API_FUNC: std::sync::LazyLock<Function> =
+ std::sync::LazyLock::new(|| Function::get_global("ffi.ModuleGetFunction").unwrap());
+ let name = crate::string::String::from(name);
+ (*API_FUNC)
+ .call_tuple_with_len::<3, _>((self, name, true))?
+ .try_into()
+ }
+}
diff --git a/rust/tvm-ffi/src/function.rs b/rust/tvm-ffi/src/function.rs
new file mode 100644
index 0000000..e488983
--- /dev/null
+++ b/rust/tvm-ffi/src/function.rs
@@ -0,0 +1,272 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::any::{Any, AnyView};
+use crate::derive::{Object, ObjectRef};
+use crate::error::{Error, Result};
+use crate::function_internal::{AsPackedCallable, TupleAsPackedArgs};
+use crate::object::{Object, ObjectArc, ObjectCore};
+use tvm_ffi_sys::{
+ TVMFFIAny, TVMFFIByteArray, TVMFFIFunctionCell, TVMFFIFunctionCreate, TVMFFIFunctionGetGlobal,
+ TVMFFIFunctionSetGlobal, TVMFFIObjectHandle, TVMFFISafeCallType, TVMFFITypeIndex,
+};
+
+/// function object
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Function"]
+#[type_index(TVMFFITypeIndex::kTVMFFIFunction)]
+pub struct FunctionObj {
+ object: Object,
+ cell: TVMFFIFunctionCell,
+}
+
+/// Error reference class
+#[derive(Clone, ObjectRef)]
+pub struct Function {
+ data: ObjectArc<FunctionObj>,
+}
+
+//------------------------------------------------------------------------
+// CallbackFunctionObjImpl
+//------------------------------------------------------------------------
+/// Special helper class to hold a generic callback state as Object
+/// Logically this Impl can be viewed as a FunctionObj
+/// We can create an ObjectArc<CallbackFunctionObjImpl<F>> so the deleter
+/// can correctly delete the entire object including callback part
+/// then we will convert to ObjectArc<FunctionObj> to be used as function
+#[repr(C)]
+struct CallbackFunctionObjImpl<F: Fn(&[AnyView]) -> Result<Any> + 'static> {
+ function: FunctionObj,
+ callback: F,
+}
+
+impl<F: Fn(&[AnyView]) -> Result<Any> + 'static> CallbackFunctionObjImpl<F> {
+ pub fn from_callback(callback: F) -> Self {
+ Self {
+ function: FunctionObj {
+ object: Object::new(),
+ cell: TVMFFIFunctionCell {
+ // specfic callback for F
+ safe_call: Self::invoke_callback,
+ cxx_call: std::ptr::null_mut(),
+ },
+ },
+ callback,
+ }
+ }
+
+ unsafe extern "C" fn invoke_callback(
+ handle: *mut std::ffi::c_void,
+ args: *const TVMFFIAny,
+ num_args: i32,
+ result: *mut TVMFFIAny,
+ ) -> i32 {
+ let this = &*(handle as *mut Self);
+ let packed_args = std::slice::from_raw_parts(args as *const AnyView, num_args as usize);
+ let ret_value = (this.callback)(packed_args);
+ match ret_value {
+ Ok(value) => {
+ *result = Any::into_raw_ffi_any(value);
+ 0
+ }
+ Err(error) => {
+ Error::set_raised(&error);
+ -1
+ }
+ }
+ }
+}
+
+unsafe impl<F: Fn(&[AnyView]) -> Result<Any> + 'static> ObjectCore for CallbackFunctionObjImpl<F> {
+ const TYPE_KEY: &'static str = FunctionObj::TYPE_KEY;
+ fn type_index() -> i32 {
+ FunctionObj::type_index()
+ }
+ unsafe fn object_header_mut(this: &mut Self) -> &mut tvm_ffi_sys::TVMFFIObject {
+ FunctionObj::object_header_mut(&mut this.function)
+ }
+}
+
+impl Function {
+ /// Call the function in packed format.
+ pub fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any> {
+ unsafe {
+ let packed_args_ptr = packed_args.as_ptr() as *const TVMFFIAny;
+ let mut result = Any::new();
+ let ret_code = (self.data.cell.safe_call)(
+ ObjectArc::as_raw(&self.data) as *mut FunctionObj as *mut std::ffi::c_void,
+ packed_args_ptr,
+ packed_args.len() as i32,
+ Any::as_data_ptr(&mut result),
+ );
+ if ret_code == 0 {
+ Ok(result)
+ } else {
+ Err(Error::from_raised())
+ }
+ }
+ }
+
+ pub fn call_tuple<TupleType>(&self, tuple_args: TupleType) -> Result<Any>
+ where
+ TupleType: TupleAsPackedArgs,
+ {
+ // This is a workaround for Rust's requirement that stack allocation size
+ // must be known at compile time for generic types.
+ // While we know args_len is a constant, Rust doesn't allow us to directly
+ // declare [AnyView::new(); args_len] in generic contexts.
+ //
+ // We use a small vector optimization pattern:
+ // 1. First allocate a small stack buffer (stack_args)
+ // 2. If args_len exceeds STACK_LEN, allocate a heap buffer (heap_args)
+ // 3. Use the appropriate buffer based on size
+ //
+ // Since args_len is a compile-time constant, the compiler should optimize
+ // away the unused branch, making this approach efficient.
+ const STACK_LEN: usize = 4;
+ let mut stack_args = [AnyView::new(); STACK_LEN];
+ let mut heap_args = Vec::<AnyView>::new();
+ let args_len = <TupleType as TupleAsPackedArgs>::LEN;
+ // get packed arguments
+ let packed_args: &mut [AnyView] = if args_len <= STACK_LEN {
+ &mut stack_args[..args_len]
+ } else {
+ heap_args.resize(args_len, AnyView::new());
+ &mut heap_args[..args_len]
+ };
+ (&tuple_args).fill_any_view(packed_args);
+ self.call_packed(packed_args)
+ }
+ /// Call function with compile-time known argument count
+ /// This is an optimized version of call_tuple for when the argument count
+ /// is known at compile time, avoiding the small vector optimization overhead.
+ ///
+ /// # Arguments
+ /// * `tuple_args` - The tuple arguments
+ ///
+ /// # Returns
+ /// * `Any` - The result
+ pub fn call_tuple_with_len<const LEN: usize, TupleType>(
+ &self,
+ tuple_args: TupleType,
+ ) -> Result<Any>
+ where
+ TupleType: TupleAsPackedArgs,
+ {
+ let mut packed_args = [AnyView::new(); LEN];
+ (&tuple_args).fill_any_view(&mut packed_args);
+ self.call_packed(&packed_args)
+ }
+ /// Get global function by name
+ /// This function will throw an error if the function is not found.
+ ///
+ /// # Arguments
+ /// * `name` - The name of the function
+ ///
+ /// # Returns
+ /// * `Function` - The global function
+ pub fn get_global(name: &str) -> Result<Function> {
+ unsafe {
+ let name_arg = TVMFFIByteArray::from_str(name);
+ let mut result: TVMFFIObjectHandle = ::std::ptr::null_mut();
+ crate::check_safe_call!(TVMFFIFunctionGetGlobal(&name_arg, &mut result))?;
+ if result.is_null() {
+ crate::bail!(crate::error::RUNTIME_ERROR, "Function {} not found", name);
+ }
+ Ok(Self {
+ data: ObjectArc::<FunctionObj>::from_raw(result as *mut FunctionObj),
+ })
+ }
+ }
+
+ /// Register a function as a global function
+ /// # Arguments
+ /// * `name` - The name of the function
+ /// * `func` - The function to register
+ ///
+ /// # Returns
+ /// * `Result<()>` - The result of the registration
+ pub fn register_global(name: &str, func: Function) -> Result<()> {
+ unsafe {
+ let name_arg = TVMFFIByteArray::from_str(name);
+ let can_override = 0;
+ crate::check_safe_call!(TVMFFIFunctionSetGlobal(
+ &name_arg,
+ ObjectArc::as_raw(&func.data) as *mut FunctionObj as TVMFFIObjectHandle,
+ can_override
+ ))?;
+ Ok(())
+ }
+ }
+ /// Construct a function from a packed function
+ /// # Arguments
+ /// * `func` - The packed function in signature of `Fn(&[AnyView]) -> Result<Any>`
+ ///
+ /// # Returns
+ /// * `Function` - The function
+ pub fn from_packed<F>(func: F) -> Self
+ where
+ F: Fn(&[AnyView]) -> Result<Any> + 'static,
+ {
+ unsafe {
+ let callback_arc = ObjectArc::new(CallbackFunctionObjImpl::from_callback(func));
+ let func_arc = ObjectArc::<FunctionObj>::from_raw(
+ ObjectArc::into_raw(callback_arc) as *mut FunctionObj
+ );
+ Self { data: func_arc }
+ }
+ }
+
+ /// Construct a function from a typed function
+ /// # Arguments
+ /// * `func` - The typed function with function signature of `F(T0, T1, ...) -> Result<O>`
+ ///
+ /// # Returns
+ /// * `Function` - The function
+ pub fn from_typed<F, I, O>(func: F) -> Self
+ where
+ F: AsPackedCallable<I, O> + 'static,
+ {
+ let closure = move |packed_args: &[AnyView]| -> Result<Any> {
+ let ret_value = func.call_packed(packed_args)?;
+ Ok(ret_value)
+ };
+ Self::from_packed(closure)
+ }
+
+ pub fn from_extern_c(
+ handle: *mut std::ffi::c_void,
+ safe_call: TVMFFISafeCallType,
+ deleter: Option<unsafe extern "C" fn(*mut std::ffi::c_void)>,
+ ) -> Self {
+ unsafe {
+ let mut out_handle: TVMFFIObjectHandle = std::ptr::null_mut();
+ crate::check_safe_call!(TVMFFIFunctionCreate(
+ handle,
+ safe_call,
+ deleter,
+ &mut out_handle
+ ))
+ .unwrap();
+ Self {
+ data: ObjectArc::<FunctionObj>::from_raw(out_handle as *mut FunctionObj),
+ }
+ }
+ }
+}
diff --git a/rust/tvm-ffi/src/function_internal.rs b/rust/tvm-ffi/src/function_internal.rs
new file mode 100644
index 0000000..e059051
--- /dev/null
+++ b/rust/tvm-ffi/src/function_internal.rs
@@ -0,0 +1,193 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::any::{Any, AnyView, ArgTryFromAnyView};
+use crate::error::Result;
+use crate::string::{Bytes, String};
+use crate::type_traits::AnyCompatible;
+
+//------------------------------------------------------------------------
+// PackedCallable
+//------------------------------------------------------------------------
+pub trait AsPackedCallable<I, O> {
+ // Call the function in packed convention
+ fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>;
+}
+
+#[inline]
+pub fn call_packed_callable<Fun, I, O>(func: Fun, packed_args: &[AnyView]) -> Result<Any>
+where
+ Fun: AsPackedCallable<I, O>,
+{
+ func.call_packed(packed_args)
+}
+
+macro_rules! impl_as_packed_callable {
+ ($len:literal; $($t:ident),*) => {
+ impl<Fun, $($t,)* Out> AsPackedCallable<($($t,)*), Out> for Fun
+ where
+ Fun: Fn($($t,)*) -> Result<Out> + 'static,
+ Any: From<Out>,
+ $($t: ArgTryFromAnyView),*
+ {
+ fn call_packed(&self, packed_args: &[AnyView]) -> Result<Any>
+ {
+ crate::ensure!(
+ packed_args.len() == $len, crate::error::VALUE_ERROR,
+ "Expected {} arguments, got {}", $len, packed_args.len()
+ );
+ // Expand the function call, consuming the iterator.
+ let mut _arg_iter = packed_args.iter().enumerate();
+ let ret_value = self(
+ $({
+ // unwrap is safe due to the length check above
+ let (i, view) = _arg_iter.next().unwrap();
+ $t::try_from_any_view(view, i)?
+ }),*
+ )?;
+ Ok(Any::from(ret_value))
+ }
+ }
+ }
+}
+
+impl_as_packed_callable!(0;);
+impl_as_packed_callable!(1; T0);
+impl_as_packed_callable!(2; T0, T1);
+impl_as_packed_callable!(3; T0, T1, T2);
+impl_as_packed_callable!(4; T0, T1, T2, T3);
+impl_as_packed_callable!(5; T0, T1, T2, T3, T4);
+impl_as_packed_callable!(6; T0, T1, T2, T3, T4, T5);
+impl_as_packed_callable!(7; T0, T1, T2, T3, T4, T5, T6);
+impl_as_packed_callable!(8; T0, T1, T2, T3, T4, T5, T6, T7);
+
+//--------------------------------------------------------------
+// IntoArgHolder, helper to convert to canonical holding type
+//
+// This is needed sometimes for reference types that may need to
+// be converted to value types.
+//--------------------------------------------------------------
+pub trait IntoArgHolder {
+ type Target;
+ fn into_arg_holder(self) -> Self::Target;
+}
+
+crate::impl_into_arg_holder_default!(
+ bool, i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64, String, Bytes
+);
+
+// string will be converted to String for argument passing
+impl IntoArgHolder for &str {
+ type Target = String;
+ fn into_arg_holder(self) -> Self::Target {
+ String::from(self)
+ }
+}
+
+// string will be converted to String for argument passing
+impl IntoArgHolder for &[u8] {
+ type Target = Bytes;
+ fn into_arg_holder(self) -> Self::Target {
+ Bytes::from(self)
+ }
+}
+
+// helper trait to implement IntoArgHolderTuple to apply into_arg_holder to each element
+pub trait IntoArgHolderTuple {
+ type Target;
+ fn into_arg_holder_tuple(self) -> Self::Target;
+}
+
+macro_rules! impl_into_arg_holder_tuple {
+ ( $($T:ident),* ; $($idx:tt),* ) => {
+ impl<$($T),*> $crate::function_internal::IntoArgHolderTuple for ($($T,)*)
+ where
+ $($T: IntoArgHolder),* {
+ type Target = ($($T::Target,)*);
+
+ fn into_arg_holder_tuple(self) -> Self::Target {
+ ($(self.$idx.into_arg_holder(),)*)
+ }
+ }
+ };
+}
+
+impl_into_arg_holder_tuple!(;);
+impl_into_arg_holder_tuple!(T0; 0);
+impl_into_arg_holder_tuple!(T0, T1; 0, 1);
+impl_into_arg_holder_tuple!(T0, T1, T2; 0, 1, 2);
+impl_into_arg_holder_tuple!(T0, T1, T2, T3; 0, 1, 2, 3);
+impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4; 0, 1, 2, 3, 4);
+impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5);
+impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6);
+impl_into_arg_holder_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7);
+
+//------------------------------------------------------------
+// ArgIntoRef
+//
+// Helper to turn argument type to reference type
+// This is effectively AsRef<T> but removes the need of T
+//-----------------------------------------------------------
+pub trait ArgIntoRef {
+ type Target;
+ fn to_ref(&self) -> &Self::Target;
+}
+
+crate::impl_arg_into_ref!(
+ bool, i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64, String, Bytes
+);
+
+//-----------------------------------------------------------
+// TupleAsPackedArgs
+//
+// Helper to turn tuple type to packed arguments
+//-----------------------------------------------------------
+pub trait TupleAsPackedArgs {
+ const LEN: usize;
+ fn fill_any_view<'a>(&'a self, any_view: &mut [AnyView<'a>]);
+}
+
+macro_rules! impl_tuple_as_packed_args {
+ ( $len:expr; $($T:ident),* ; $($idx:tt),* ) => {
+ impl<$($T),*> TupleAsPackedArgs for ($($T,)*)
+ where
+ $(
+ $T: ArgIntoRef,
+ $T::Target: AnyCompatible,
+ )*
+ {
+ const LEN: usize = $len;
+
+ fn fill_any_view<'a>(&'a self, _any_view: &mut [AnyView<'a>]) {
+ $(
+ _any_view[$idx] = AnyView::from(self.$idx.to_ref());
+ )*
+ }
+ }
+ };
+}
+
+impl_tuple_as_packed_args!(0;;);
+impl_tuple_as_packed_args!(1; T0; 0);
+impl_tuple_as_packed_args!(2; T0, T1; 0, 1);
+impl_tuple_as_packed_args!(3; T0, T1, T2; 0, 1, 2);
+impl_tuple_as_packed_args!(4; T0, T1, T2, T3; 0, 1, 2, 3);
+impl_tuple_as_packed_args!(5; T0, T1, T2, T3, T4; 0, 1, 2, 3, 4);
+impl_tuple_as_packed_args!(6; T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5);
+impl_tuple_as_packed_args!(7; T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6);
+impl_tuple_as_packed_args!(8; T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7);
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
new file mode 100644
index 0000000..94e87b0
--- /dev/null
+++ b/rust/tvm-ffi/src/lib.rs
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+pub mod any;
+pub mod collections;
+pub mod derive;
+pub mod device;
+pub mod dtype;
+pub mod error;
+pub mod extra;
+pub mod function;
+pub mod function_internal;
+pub mod macros;
+pub mod object;
+pub mod string;
+pub mod type_traits;
+pub use tvm_ffi_sys;
+
+pub use crate::any::{Any, AnyView};
+pub use crate::collections::shape::Shape;
+pub use crate::collections::tensor::{CPUNDAlloc, NDAllocator, Tensor};
+pub use crate::device::{current_stream, with_stream};
+pub use crate::dtype::DLDataTypeExt;
+pub use crate::error::{Error, ErrorKind, Result};
+pub use crate::error::{
+ ATTRIBUTE_ERROR, INDEX_ERROR, KEY_ERROR, RUNTIME_ERROR, TYPE_ERROR, VALUE_ERROR,
+};
+pub use crate::extra::module::Module;
+pub use crate::function::Function;
+pub use crate::object::{Object, ObjectArc, ObjectCore, ObjectCoreWithExtraItems, ObjectRefCore};
+pub use crate::string::{Bytes, String};
+pub use crate::type_traits::AnyCompatible;
+
+pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+pub use tvm_ffi_sys::{
+ DLDataType, DLDataTypeCode, DLDevice, DLDeviceType, TVMFFIAny, TVMFFIObject, TVMFFIStreamHandle,
+};
diff --git a/rust/tvm-ffi/src/macros.rs b/rust/tvm-ffi/src/macros.rs
new file mode 100644
index 0000000..95c2b5b
--- /dev/null
+++ b/rust/tvm-ffi/src/macros.rs
@@ -0,0 +1,398 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+// rexport paste under macro namespace so downstream do not need to specify dep
+pub use paste;
+// ----------------------------------------------------------------------------
+// Macros for error handling
+// ----------------------------------------------------------------------------
+
+/// Macro gto get the name of the function
+///
+/// # Usage
+/// Usage: function_name!()
+#[macro_export]
+macro_rules! function_name {
+ () => {{
+ // dummy function to get the name of the function
+ fn f() {}
+ fn type_name_of<T>(_: T) -> &'static str {
+ std::any::type_name::<T>()
+ }
+ let name = type_name_of(f);
+ // remove the f() from the name
+ &name[..name.len() - 3]
+ }};
+}
+
+/// Check the return code of the safe call
+///
+/// # Arguments
+/// * `ret_code` - The return code of the safe call
+///
+/// # Returns
+/// * `Result<(), Error>` - The result of the safe call
+/// Macro to check safe calls and automatically update traceback with file/line info
+///
+/// Usage: check_safe_call!(function(args))?;
+#[macro_export]
+macro_rules! check_safe_call {
+ ($expr:expr) => {{
+ let ret_code = $expr;
+ if ret_code == 0 {
+ Ok(())
+ } else {
+ let error = $crate::error::Error::from_raised();
+ Err(error)
+ }
+ }};
+}
+
+/// Create a new error with file/line info attached
+///
+/// This macro automatically appends file/line info to the traceback
+///
+/// # Arguments
+/// * `error_kind` - The kind of the error
+/// * `msg` - The message of the error
+/// * `args` - The posisble format arguments
+///
+/// # Returns
+/// * `Result<(), Error>` - The result of the safe call
+#[macro_export]
+macro_rules! bail {
+ ($error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
+ let context = format!(
+ " File \"{}\", line {}, in {}\n",
+ file!(),
+ line!(),
+ $crate::function_name!()
+ );
+ return Err($crate::error::Error::new($error_kind, &format!($fmt $(, $args)*), &context));
+ }};
+}
+
+/// Create a new error with file/line info attached
+///
+/// This macro automatically appends file/line info to the traceback
+///
+/// # Arguments
+/// * `kind` - The kind of the error
+/// * `msg` - The message of the error
+/// * `args` - The posisble format arguments
+///
+/// # Returns
+/// * `Result<(), Error>` - The result of the safe call
+#[macro_export]
+macro_rules! ensure {
+ ($cond:expr, $error_kind:expr, $fmt:expr $(, $args:expr)* $(,)?) => {{
+ if !$cond {
+ crate::bail!($error_kind, $fmt $(, $args)*);
+ }
+ }};
+}
+
+/// Attach a context to a result if it is error
+///
+/// This macro automatically appends file/line info to the traceback
+///
+/// # Arguments
+/// * `error` - The error to attach the context to
+/// * `msg` - The message of the error
+///
+/// # Returns
+/// * `Result<(), Error>` - The result of the safe call
+#[macro_export]
+macro_rules! attach_context {
+ ($error:expr) => {{
+ match $error {
+ Ok(value) => Ok(value),
+ Err(error) => {
+ let context = format!(
+ " File \"{}\", line {}, in {}\n",
+ file!(),
+ line!(),
+ $crate::function_name!()
+ );
+ Err(Error::with_appended_backtrace(error, &context))
+ }
+ }
+ }};
+}
+
+// ----------------------------------------------------------------------------
+// Macros for any definitions
+// ----------------------------------------------------------------------------
+
+// implements try from any for all integer types
+/// Macro to implement `TryFrom<AnyView>` and `TryFrom<Any>` for a list of types
+#[macro_export]
+macro_rules! impl_try_from_any {
+ ($($t:ty),* $(,)?) => {
+ $(
+ impl<'a> TryFrom<$crate::any::AnyView<'a>> for $t {
+ type Error = $crate::error::Error;
+ #[inline(always)]
+ fn try_from(
+ value: $crate::any::AnyView<'a>
+ ) -> Result<Self, Self::Error> {
+ type TryFromTemp = $crate::any::TryFromTemp<$t>;
+ return TryFromTemp::try_from(value).map(TryFromTemp::into_value);
+ }
+ }
+
+ impl TryFrom<$crate::any::Any> for $t {
+ type Error = $crate::error::Error;
+ #[inline(always)]
+ fn try_from(
+ value: $crate::any::Any
+ ) -> Result<Self, Self::Error> {
+ type TryFromTemp = $crate::any::TryFromTemp<$t>;
+ return TryFromTemp::try_from(value).map(TryFromTemp::into_value);
+ }
+ }
+ )*
+ };
+}
+
+/// Macro to implement `TryFrom<AnyView>` and `TryFrom<Any>` for generic types like `Option<T>`
+#[macro_export]
+macro_rules! impl_try_from_any_for_parametric {
+ ($generic_type:ident<$param:ident>) => {
+ impl<'a, $param: AnyCompatible> TryFrom<$crate::any::AnyView<'a>>
+ for $generic_type<$param>
+ {
+ type Error = $crate::error::Error;
+ #[inline(always)]
+ fn try_from(value: $crate::any::AnyView<'a>) -> Result<Self, Self::Error> {
+ type TryFromTemp<T> = $crate::any::TryFromTemp<$generic_type<$param>>;
+ return TryFromTemp::<T>::try_from(value).map(TryFromTemp::<T>::into_value);
+ }
+ }
+
+ impl<$param: AnyCompatible> TryFrom<$crate::any::Any> for $generic_type<$param> {
+ type Error = $crate::error::Error;
+ #[inline(always)]
+ fn try_from(value: $crate::any::Any) -> Result<Self, Self::Error> {
+ type TryFromTemp<T> = $crate::any::TryFromTemp<$generic_type<$param>>;
+ return TryFromTemp::<T>::try_from(value).map(TryFromTemp::<T>::into_value);
+ }
+ }
+ };
+}
+
+/// Macro to implement IntoArgHolder for a list of types
+#[macro_export]
+macro_rules! impl_into_arg_holder_default {
+ ($($t:ty),*) => {
+ $(
+ impl $crate::function_internal::IntoArgHolder for $t {
+ type Target = $t;
+ fn into_arg_holder(self) -> Self::Target {
+ self
+ }
+ }
+ impl<'a> $crate::function_internal::IntoArgHolder for &'a $t {
+ type Target = &'a $t;
+ fn into_arg_holder(self) -> Self::Target {
+ self
+ }
+ }
+ )*
+ };
+}
+
+/// Macro to implement ArgIntoRef for a list of types
+#[macro_export]
+macro_rules! impl_arg_into_ref {
+ ($($t:ty),*) => {
+ $(
+ impl $crate::function_internal::ArgIntoRef for $t {
+ type Target = $t;
+ fn to_ref(&self) -> &Self::Target {
+ &self
+ }
+ }
+ impl<'a> $crate::function_internal::ArgIntoRef for &'a $t {
+ type Target = $t;
+ fn to_ref(&self) -> &Self::Target {
+ &self
+ }
+ }
+ )*
+ }
+}
+
+// ----------------------------------------------------------------------------
+// Macros for function definitions
+// ----------------------------------------------------------------------------
+
+/// Macro to export a typed function as a C symbol that follows the tvm-ffi ABI
+///
+/// # Arguments
+/// * `$name` - The name of the function
+/// * `$func` - The function to export
+///
+/// # Example
+/// ```rust
+/// use tvm_ffi::*;
+///
+/// fn add_one(x: i32) -> Result<i32> { Ok(x + 1) }
+///
+/// tvm_ffi_dll_export_typed_func!(add_one, add_one);
+/// ```
+#[macro_export]
+macro_rules! tvm_ffi_dll_export_typed_func {
+ ($name:ident, $func:expr) => {
+ $crate::macros::paste::paste! {
+ pub unsafe extern "C" fn [<__tvm_ffi_ $name>](
+ _handle: *mut std::ffi::c_void,
+ args: *const tvm_ffi_sys::TVMFFIAny,
+ num_args: i32,
+ result: *mut tvm_ffi_sys::TVMFFIAny,
+ ) -> i32 {
+ let packed_args =
+ std::slice::from_raw_parts(args as *const $crate::any::AnyView, num_args as usize);
+ let ret_value = $crate::function_internal::call_packed_callable($func, packed_args);
+ match ret_value {
+ Ok(value) => {
+ *result = $crate::any::Any::into_raw_ffi_any(value);
+ 0
+ }
+ Err(error) => {
+ $crate::error::Error::set_raised(&error);
+ -1
+ }
+ }
+ }
+ }
+ };
+}
+
+///-----------------------------------------------------------
+/// into_typed_fn
+///
+/// Converts a generic `Function` into a typed function with compile-time
+/// argument count and type checking. This macro provides a convenient way
+/// to create type-safe wrappers around TVM functions.
+///
+/// # Arguments
+/// * `$f` - The function identifier to convert
+/// * `$trait` - The trait type (typically `Fn`)
+/// * `($t0, $t1, ...)` - The argument types
+/// * `$ret_ty` - The return type
+///
+/// # Example
+/// ```rust
+/// use tvm_ffi::*;
+///
+/// let func = Function::from_typed(|x: i32, y: i32| -> Result<i32> { Ok(x + y) });
+/// let typed_func = into_typed_fn!(func, Fn(i32, &i32) -> Result<i32>);
+/// let result = typed_func(10, &20).unwrap(); // Returns 30
+/// assert_eq!(result, 30);
+/// ```
+/// Note that the `into_typed_fn!` macro can specify arguments to be passed either
+/// by reference or by value in the argument list.
+/// We recommend passing by reference for ObjectRef types such as Tensor.
+/// Since the ffi mechanism requires us to pass arguments by reference.
+///
+/// # Supported Argument Counts
+/// This macro supports functions with 0 to 8 arguments.
+///-----------------------------------------------------------
+#[macro_export]
+macro_rules! into_typed_fn {
+ // Case for 0 arguments
+ ($f:expr, $trait:ident() -> $ret_ty:ty) => {{
+ let _f = $f;
+ move || -> $ret_ty { Ok(_f.call_tuple_with_len::<0, _>(())?.try_into()?) }
+ }};
+ // Case for 1 argument
+ ($f:expr, $trait:ident($t0:ty) -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0,).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<1, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 2 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty) -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<2, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 3 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty) -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1, a2: $t2| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1, a2).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<3, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 4 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty) -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1, a2: $t2, a3: $t3| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1, a2, a3).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<4, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 5 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty) -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1, a2, a3, a4).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<5, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 6 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty) -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1, a2, a3, a4, a5).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<6, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 7 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty)
+ -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1, a2, a3, a4, a5, a6).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<7, _>(tuple_args)?.try_into()?)
+ }
+ }};
+ // Case for 8 arguments
+ ($f:expr, $trait:ident($t0:ty, $t1:ty, $t2:ty, $t3:ty, $t4:ty, $t5:ty, $t6:ty, $t7:ty)
+ -> $ret_ty:ty) => {{
+ let _f = $f;
+ move |a0: $t0, a1: $t1, a2: $t2, a3: $t3, a4: $t4, a5: $t5, a6: $t6, a7: $t7| -> $ret_ty {
+ use $crate::function_internal::IntoArgHolderTuple;
+ let tuple_args = (a0, a1, a2, a3, a4, a5, a6, a7).into_arg_holder_tuple();
+ Ok(_f.call_tuple_with_len::<8, _>(tuple_args)?.try_into()?)
+ }
+ }};
+}
diff --git a/rust/tvm-ffi/src/object.rs b/rust/tvm-ffi/src/object.rs
new file mode 100644
index 0000000..dc1970a
--- /dev/null
+++ b/rust/tvm-ffi/src/object.rs
@@ -0,0 +1,473 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use std::ops::{Deref, DerefMut};
+use std::sync::atomic::AtomicU64;
+
+use crate::derive::ObjectRef;
+pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+/// Object related ABI handling
+use tvm_ffi_sys::{TVMFFIObject, COMBINED_REF_COUNT_BOTH_ONE};
+
+/// Object type is by default the TVMFFIObject
+#[repr(C)]
+pub struct Object {
+ /// example implementation of the object
+ header: TVMFFIObject,
+}
+
+/// Arc-like wrapper for Object that allows shared ownership
+///
+/// \tparam T The type of the object to be wrapped
+#[repr(C)]
+pub struct ObjectArc<T: ObjectCore> {
+ ptr: std::ptr::NonNull<T>,
+ _phantom: std::marker::PhantomData<T>,
+}
+
+unsafe impl<T: Send + Sync + ObjectCore> Send for ObjectArc<T> {}
+unsafe impl<T: Send + Sync + ObjectCore> Sync for ObjectArc<T> {}
+
+/// Traits that can be used to check if a type is an object
+///
+/// This trait is unsafe because it is used to access the object header
+/// and the object header is unsafe to access
+pub unsafe trait ObjectCore: Sized + 'static {
+ /// the type key of the object
+ const TYPE_KEY: &'static str;
+ // return the type index of the object
+ fn type_index() -> i32;
+ /// Return the object header
+ /// This function is implemented as a static function so
+ ///
+ /// # Arguments
+ /// * `this` - The object to get the header
+ ///
+ /// # Returns
+ /// * `&mut TVMFFIObject` - The object header
+ /// \return The object header
+ unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject;
+}
+
+/// Traits for objects with extra items that follows the object
+///
+/// This extra trait can be helpful to implement array types and string types
+pub unsafe trait ObjectCoreWithExtraItems: ObjectCore {
+ /// type of extra items storage that follows the object
+ type ExtraItem;
+ /// Return the number of extra items
+ fn extra_items_count(this: &Self) -> usize;
+ /// Return the extra items data pointer
+ unsafe fn extra_items(this: &Self) -> &[Self::ExtraItem] {
+ let extra_items_ptr = (this as *const Self as *const u8).add(std::mem::size_of::<Self>());
+ std::slice::from_raw_parts(
+ extra_items_ptr as *const Self::ExtraItem,
+ Self::extra_items_count(this),
+ )
+ }
+ /// Return the extra items data pointer
+ unsafe fn extra_items_mut(this: &mut Self) -> &mut [Self::ExtraItem] {
+ let extra_items_ptr = (this as *mut Self as *mut u8).add(std::mem::size_of::<Self>());
+ std::slice::from_raw_parts_mut(
+ extra_items_ptr as *mut Self::ExtraItem,
+ Self::extra_items_count(this),
+ )
+ }
+}
+
+/// Traits to specify core operations of ObjectRef
+///
+/// used by the ffi Any system and not user facing
+///
+/// We mark as unsafe since it moves out the internal of the ObjectRef
+pub unsafe trait ObjectRefCore: Sized + Clone {
+ type ContainerType: ObjectCore;
+ fn data(this: &Self) -> &ObjectArc<Self::ContainerType>;
+ fn into_data(this: Self) -> ObjectArc<Self::ContainerType>;
+ fn from_data(data: ObjectArc<Self::ContainerType>) -> Self;
+}
+
+/// Base class for ObjectRef
+///
+/// This class is used to store the data of the ObjectRef
+#[repr(C)]
+#[derive(ObjectRef, Clone)]
+pub struct ObjectRef {
+ data: ObjectArc<Object>,
+}
+
+/// Unsafe operations on object
+pub(crate) mod unsafe_ {
+ use tvm_ffi_sys::{
+ COMBINED_REF_COUNT_BOTH_ONE, COMBINED_REF_COUNT_MASK_U32, COMBINED_REF_COUNT_STRONG_ONE,
+ COMBINED_REF_COUNT_WEAK_ONE,
+ };
+
+ use std::ffi::c_void;
+ use std::sync::atomic::{fence, Ordering};
+ use tvm_ffi_sys::TVMFFIObject;
+ use tvm_ffi_sys::TVMFFIObjectDeleterFlagBitMask::{
+ kTVMFFIObjectDeleterFlagBitMaskBoth, kTVMFFIObjectDeleterFlagBitMaskStrong,
+ kTVMFFIObjectDeleterFlagBitMaskWeak,
+ };
+
+ /// Increase the strong reference count of the object
+ ///
+ /// This function is same as TVMFFIObjectIncRef but implemented natively in Rust
+ ///
+ /// # Arguments
+ /// * `obj` - The object to increase the reference count
+ #[inline]
+ pub unsafe fn inc_ref(handle: *mut TVMFFIObject) {
+ let obj = &mut *handle;
+ obj.combined_ref_count.fetch_add(1, Ordering::Relaxed);
+ }
+
+ /// Decrease the strong reference count of the object
+ ///
+ /// This function is same as TVMFFIObjectDecRef but implemented natively in Rust
+ ///
+ /// # Arguments
+ /// * `obj` - The object to decrease the reference count
+ #[inline]
+ pub unsafe fn dec_ref(handle: *mut TVMFFIObject) {
+ let obj = &mut *handle;
+ let old_combined_count = obj
+ .combined_ref_count
+ .fetch_sub(COMBINED_REF_COUNT_STRONG_ONE, Ordering::Relaxed);
+ if old_combined_count == COMBINED_REF_COUNT_BOTH_ONE {
+ if let Some(deleter) = obj.deleter {
+ fence(Ordering::Acquire);
+ deleter(
+ obj as *mut TVMFFIObject as *mut c_void,
+ kTVMFFIObjectDeleterFlagBitMaskBoth as i32,
+ );
+ }
+ } else if (old_combined_count & COMBINED_REF_COUNT_MASK_U32)
+ == COMBINED_REF_COUNT_STRONG_ONE
+ {
+ // slow path, there is still a weak reference left
+ // need to run two phase decrement
+ fence(Ordering::Acquire);
+ if let Some(deleter) = obj.deleter {
+ deleter(
+ obj as *mut TVMFFIObject as *mut c_void,
+ kTVMFFIObjectDeleterFlagBitMaskStrong as i32,
+ );
+ }
+ let old_weak_count = obj
+ .combined_ref_count
+ .fetch_sub(COMBINED_REF_COUNT_WEAK_ONE, Ordering::Release);
+ if old_weak_count == COMBINED_REF_COUNT_WEAK_ONE {
+ fence(Ordering::Acquire);
+ if let Some(deleter) = obj.deleter {
+ deleter(
+ obj as *mut TVMFFIObject as *mut c_void,
+ kTVMFFIObjectDeleterFlagBitMaskWeak as i32,
+ );
+ }
+ }
+ }
+ }
+
+ #[inline]
+ pub unsafe fn strong_count(handle: *mut TVMFFIObject) -> usize {
+ let obj = &mut *handle;
+ (obj.combined_ref_count.load(Ordering::Relaxed) & COMBINED_REF_COUNT_MASK_U32) as usize
+ }
+
+ #[inline]
+ pub unsafe fn weak_count(handle: *mut TVMFFIObject) -> usize {
+ let obj = &mut *handle;
+ (obj.combined_ref_count.load(Ordering::Relaxed) >> 32) as usize
+ }
+
+ /// Generic object deleter that works for object allocated from Box then into_raw
+ pub unsafe extern "C" fn object_deleter_for_new<T>(ptr: *mut c_void, flags: i32)
+ where
+ T: super::ObjectCore,
+ {
+ let obj = ptr as *mut T;
+ if flags & kTVMFFIObjectDeleterFlagBitMaskStrong as i32 != 0 {
+ // calling destructor of the object, does not free the memory
+ std::ptr::drop_in_place(obj);
+ }
+ if flags & kTVMFFIObjectDeleterFlagBitMaskWeak as i32 != 0 {
+ // free the memory
+ std::alloc::dealloc(ptr as *mut u8, std::alloc::Layout::new::<T>());
+ }
+ }
+
+ pub unsafe extern "C" fn object_deleter_for_new_with_extra_items<T, U>(
+ ptr: *mut c_void,
+ flags: i32,
+ ) where
+ T: super::ObjectCoreWithExtraItems<ExtraItem = U>,
+ {
+ let obj = ptr as *mut T;
+ if flags == kTVMFFIObjectDeleterFlagBitMaskBoth as i32 {
+ // must get extra items count before dropping the object
+ let extra_items_count = T::extra_items_count(&(*obj));
+ std::ptr::drop_in_place(obj);
+ let layout = std::alloc::Layout::from_size_align(
+ std::mem::size_of::<T>() + extra_items_count * std::mem::size_of::<U>(),
+ std::mem::align_of::<T>(),
+ )
+ .unwrap();
+ // free the memory
+ std::alloc::dealloc(ptr as *mut u8, layout);
+ } else {
+ assert_eq!(std::mem::size_of::<T>() % std::mem::size_of::<u64>(), 0);
+ if flags & kTVMFFIObjectDeleterFlagBitMaskStrong as i32 != 0 {
+ // must get extra items count before dropping the object
+ let extra_items_count = T::extra_items_count(&(*obj));
+ // calling destructor of the object, does not free the memory
+ std::ptr::drop_in_place(obj);
+ // record extra count in the original memory
+ std::ptr::write(obj as *mut u64, extra_items_count as u64);
+ }
+ if flags & kTVMFFIObjectDeleterFlagBitMaskWeak as i32 != 0 {
+ // read extra items count from the original memory
+ // note we can no longer read it by calling T::extra_items_count(&(*obj))
+ // because the object is already dropped
+ let extra_items_count = std::ptr::read(obj as *mut u64) as usize;
+ let layout = std::alloc::Layout::from_size_align(
+ std::mem::size_of::<T>() + extra_items_count * std::mem::size_of::<U>(),
+ std::mem::align_of::<T>(),
+ )
+ .unwrap();
+ // free the memory
+ std::alloc::dealloc(ptr as *mut u8, layout);
+ }
+ }
+ }
+}
+
+//---------------------
+// Object
+//---------------------
+
+impl Object {
+ pub fn new() -> Self {
+ Self {
+ header: TVMFFIObject::new(),
+ }
+ }
+}
+
+unsafe impl ObjectCore for Object {
+ const TYPE_KEY: &'static str = "ffi.Object";
+ #[inline]
+ fn type_index() -> i32 {
+ TypeIndex::kTVMFFIStaticObjectBegin as i32
+ }
+ #[inline]
+ unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject {
+ &mut this.header
+ }
+}
+
+//---------------------
+// ObjectArc
+//---------------------
+impl<T: ObjectCore> ObjectArc<T> {
+ pub fn new(data: T) -> Self {
+ unsafe {
+ let layout = std::alloc::Layout::new::<T>();
+ let raw_data_ptr = std::alloc::alloc(layout);
+ if raw_data_ptr.is_null() {
+ std::alloc::handle_alloc_error(layout);
+ }
+ let ptr = raw_data_ptr as *mut T;
+ std::ptr::write(ptr, data);
+ // now override the header directly
+ std::ptr::write(
+ ptr as *mut TVMFFIObject,
+ TVMFFIObject {
+ combined_ref_count: AtomicU64::new(COMBINED_REF_COUNT_BOTH_ONE),
+ type_index: T::type_index(),
+ __padding: 0,
+ deleter: Some(unsafe_::object_deleter_for_new::<T>),
+ },
+ );
+ // move into the object arc ptr
+ Self {
+ ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T),
+ _phantom: std::marker::PhantomData,
+ }
+ }
+ }
+ pub fn new_with_extra_items<U>(data: T) -> Self
+ where
+ T: ObjectCoreWithExtraItems<ExtraItem = U>,
+ {
+ unsafe {
+ // ensure strict alignment requirements
+ // so we can have { T, U*extra_items } layout
+ assert_eq!(std::mem::align_of::<T>() % std::mem::align_of::<U>(), 0);
+ assert_eq!(std::mem::size_of::<T>() % std::mem::align_of::<U>(), 0);
+ let extra_items_count = T::extra_items_count(&data);
+ let layout = std::alloc::Layout::from_size_align(
+ std::mem::size_of::<T>() + extra_items_count * std::mem::size_of::<U>(),
+ std::mem::align_of::<T>(),
+ )
+ .unwrap();
+ let raw_data_ptr = std::alloc::alloc(layout);
+ if raw_data_ptr.is_null() {
+ std::alloc::handle_alloc_error(layout);
+ }
+ let ptr = raw_data_ptr as *mut T;
+ std::ptr::write(ptr, data);
+ // now override the header directly
+ std::ptr::write(
+ ptr as *mut TVMFFIObject,
+ TVMFFIObject {
+ combined_ref_count: AtomicU64::new(COMBINED_REF_COUNT_BOTH_ONE),
+ type_index: T::type_index(),
+ __padding: 0,
+ deleter: Some(unsafe_::object_deleter_for_new_with_extra_items::<T, U>),
+ },
+ );
+ // move into the object arc ptr
+ Self {
+ ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T),
+ _phantom: std::marker::PhantomData,
+ }
+ }
+ }
+
+ /// Move a previously allocated object into the ObjectArc
+ ///
+ /// # Arguments
+ /// * `ptr` - The raw pointer to move into the ObjectArc
+ ///
+ /// # Returns
+ /// * `ObjectArc<T>` - The ObjectArc
+ /// \return The ObjectArc
+ #[inline]
+ pub unsafe fn from_raw(ptr: *const T) -> Self {
+ Self {
+ ptr: std::ptr::NonNull::new_unchecked(ptr as *mut T),
+ _phantom: std::marker::PhantomData,
+ }
+ }
+
+ /// Move the ObjectArc into a raw pointer
+ ///
+ /// # Arguments
+ /// * `this` - The ObjectArc to move into a raw pointer
+ ///
+ /// # Returns
+ /// * `*const T` - The raw pointer
+ #[inline]
+ pub unsafe fn into_raw(this: Self) -> *const T {
+ let droped_this = std::mem::ManuallyDrop::new(this);
+ droped_this.ptr.as_ptr() as *const T
+ }
+
+ /// Get the raw pointer from the ObjectArc
+ ///
+ /// Caller should view this as a non-owning reference
+ ///
+ /// # Arguments
+ /// * `this` - The ObjectArc to get the raw pointer
+ ///
+ /// # Returns
+ /// * `*const T` - The raw pointer
+ /// \return The raw pointer
+ #[inline]
+ pub unsafe fn as_raw(this: &Self) -> *const T {
+ this.ptr.as_ptr() as *const T
+ }
+
+ /// Get the raw mutable pointer from the ObjectArc
+ ///
+ /// Caller should view this as a non-owning reference
+ ///
+ /// # Arguments
+ /// * `this` - The ObjectArc to get the raw pointer
+ ///
+ /// # Returns
+ /// * `*mut T` - The raw pointer
+ #[inline]
+ pub unsafe fn as_raw_mut(this: &mut Self) -> *mut T {
+ this.ptr.as_mut()
+ }
+
+ /// Get the strong reference count of the ObjectArc
+ ///
+ /// # Arguments
+ /// * `this` - The ObjectArc to get the strong reference count
+ ///
+ /// # Returns
+ /// * `usize` - The strong reference count
+ #[inline]
+ pub fn strong_count(this: &Self) -> usize {
+ unsafe {
+ unsafe_::strong_count(this.ptr.as_ref() as *const T as *mut T as *mut TVMFFIObject)
+ }
+ }
+
+ /// Get the weak reference count of the ObjectArc
+ ///
+ /// # Arguments
+ /// * `this` - The ObjectArc to get the weak reference count
+ ///
+ /// # Returns
+ /// * `usize` - The weak reference count
+ #[inline]
+ pub fn weak_count(this: &Self) -> usize {
+ unsafe { unsafe_::weak_count(this.ptr.as_ref() as *const T as *mut T as *mut TVMFFIObject) }
+ }
+}
+
+// implement Deref for ObjectArc
+impl<T: ObjectCore> Deref for ObjectArc<T> {
+ type Target = T;
+ #[inline]
+ fn deref(&self) -> &Self::Target {
+ unsafe { self.ptr.as_ref() }
+ }
+}
+
+// implement DerefMut for ObjectArc
+impl<T: ObjectCore> DerefMut for ObjectArc<T> {
+ #[inline]
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ unsafe { self.ptr.as_mut() }
+ }
+}
+
+// implement Drop for ObjectArc
+impl<T: ObjectCore> Drop for ObjectArc<T> {
+ fn drop(&mut self) {
+ unsafe { unsafe_::dec_ref(self.ptr.as_mut() as *mut T as *mut TVMFFIObject) }
+ }
+}
+
+// implement Clone for ObjectArc
+impl<T: ObjectCore> Clone for ObjectArc<T> {
+ #[inline]
+ fn clone(&self) -> Self {
+ unsafe { unsafe_::inc_ref(self.ptr.as_ref() as *const T as *mut T as *mut TVMFFIObject) }
+ Self {
+ ptr: self.ptr,
+ _phantom: std::marker::PhantomData,
+ }
+ }
+}
diff --git a/rust/tvm-ffi/src/string.rs b/rust/tvm-ffi/src/string.rs
new file mode 100644
index 0000000..94739a4
--- /dev/null
+++ b/rust/tvm-ffi/src/string.rs
@@ -0,0 +1,533 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use crate::derive::Object;
+use crate::object::{unsafe_, Object, ObjectArc, ObjectCoreWithExtraItems};
+use crate::type_traits::AnyCompatible;
+use std::cmp::Ordering;
+use std::fmt::{Debug, Display};
+use std::hash::{Hash, Hasher};
+use std::ops::Deref;
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIAnyDataUnion, TVMFFIByteArray, TVMFFIObject};
+
+//-----------------------------------------------------
+// Bytes
+//-----------------------------------------------------
+/// ABI stable Bytes container for ffi
+#[repr(C)]
+pub struct Bytes {
+ data: TVMFFIAny,
+}
+
+// BytesObj for heap-allocated bytes
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.Bytes"]
+#[type_index(TypeIndex::kTVMFFIBytes)]
+pub(crate) struct BytesObj {
+ object: Object,
+ data: TVMFFIByteArray,
+}
+
+impl Bytes {
+ /// Create a new empty Bytes container
+ pub fn new() -> Self {
+ Self {
+ data: TVMFFIAny {
+ type_index: TypeIndex::kTVMFFISmallBytes as i32,
+ small_str_len: 0,
+ data_union: TVMFFIAnyDataUnion { v_int64: 0 },
+ },
+ }
+ }
+
+ /// Get the length of the bytes
+ pub fn len(&self) -> usize {
+ self.as_slice().len()
+ }
+
+ /// Get the bytes as a slice
+ pub fn as_slice(&self) -> &[u8] {
+ unsafe {
+ if self.data.type_index == TypeIndex::kTVMFFISmallBytes as i32 {
+ std::slice::from_raw_parts(
+ self.data.data_union.v_bytes.as_ptr(),
+ self.data.small_str_len as usize,
+ )
+ } else {
+ let str_obj: &BytesObj = &*(self.data.data_union.v_obj as *const BytesObj);
+ std::slice::from_raw_parts(str_obj.data.data, str_obj.data.size)
+ }
+ }
+ }
+}
+
+impl Default for Bytes {
+ #[inline]
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+unsafe impl ObjectCoreWithExtraItems for BytesObj {
+ type ExtraItem = u8;
+ // extra item is the trailing \0 for ffi compatibility
+ #[inline]
+ /// Get the count of extra items (trailing null byte for FFI compatibility)
+ fn extra_items_count(this: &Self) -> usize {
+ return this.data.size + 1;
+ }
+}
+
+impl<T> From<T> for Bytes
+where
+ T: AsRef<[u8]>,
+{
+ #[inline]
+ /// Create Bytes from any type that can be converted to a byte slice
+ fn from(src: T) -> Self {
+ let value: &[u8] = src.as_ref();
+ // to be compatible with normal c++
+ const MAX_SMALL_BYTES_LEN: usize = 7;
+ unsafe {
+ if value.len() <= MAX_SMALL_BYTES_LEN {
+ let mut data_union = TVMFFIAnyDataUnion { v_int64: 0 };
+ data_union.v_bytes[..value.len()].copy_from_slice(value);
+ // small bytes
+ Self {
+ data: TVMFFIAny {
+ type_index: TypeIndex::kTVMFFISmallBytes as i32,
+ small_str_len: value.len() as u32,
+ data_union: data_union,
+ },
+ }
+ } else {
+ // large bytes
+ let mut obj_arc = ObjectArc::new_with_extra_items(BytesObj {
+ object: Object::new(),
+ data: TVMFFIByteArray {
+ data: std::ptr::null(),
+ size: value.len(),
+ },
+ });
+ // reset the data ptr correctly after Arc is created
+ obj_arc.data.data = BytesObj::extra_items(&obj_arc).as_ptr();
+ let extra_items = BytesObj::extra_items_mut(&mut obj_arc);
+ extra_items[..value.len()].copy_from_slice(value);
+ // write the trailing \0 for ffi compatibility
+ extra_items[value.len()] = 0;
+ Self {
+ data: TVMFFIAny {
+ type_index: TypeIndex::kTVMFFIBytes as i32,
+ small_str_len: 0,
+ data_union: TVMFFIAnyDataUnion {
+ v_obj: ObjectArc::into_raw(obj_arc) as *mut BytesObj
+ as *mut TVMFFIObject,
+ },
+ },
+ }
+ }
+ }
+ }
+}
+
+impl Deref for Bytes {
+ type Target = [u8];
+ #[inline]
+ fn deref(&self) -> &[u8] {
+ self.as_slice()
+ }
+}
+
+impl Clone for Bytes {
+ #[inline]
+ fn clone(&self) -> Self {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { unsafe_::inc_ref(self.data.data_union.v_obj) }
+ }
+ Self { data: self.data }
+ }
+}
+
+impl Drop for Bytes {
+ #[inline]
+ fn drop(&mut self) {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { unsafe_::dec_ref(self.data.data_union.v_obj) }
+ }
+ }
+}
+
+impl PartialEq for Bytes {
+ #[inline]
+ fn eq(&self, other: &Self) -> bool {
+ self.as_slice() == other.as_slice()
+ }
+}
+
+impl Eq for Bytes {}
+
+impl PartialOrd for Bytes {
+ #[inline]
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ self.as_slice().partial_cmp(other.as_slice())
+ }
+}
+
+impl Ord for Bytes {
+ #[inline]
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.as_slice().cmp(other.as_slice())
+ }
+}
+
+impl Hash for Bytes {
+ #[inline]
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.as_slice().hash(state);
+ }
+}
+
+impl Display for Bytes {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ Debug::fmt(&self.as_slice(), f)
+ }
+}
+
+impl Debug for Bytes {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("ffi.Bytes")
+ .field("data", &self.as_slice())
+ .finish()
+ }
+}
+
+//-----------------------------------------------------
+// String
+//-----------------------------------------------------
+
+/// ABI stable String container for ffi
+#[repr(C)]
+pub struct String {
+ data: TVMFFIAny,
+}
+
+// StringObj for heap-allocated string
+#[repr(C)]
+#[derive(Object)]
+#[type_key = "ffi.String"]
+#[type_index(TypeIndex::kTVMFFIStr)]
+pub(crate) struct StringObj {
+ object: Object,
+ data: TVMFFIByteArray,
+}
+
+unsafe impl ObjectCoreWithExtraItems for StringObj {
+ type ExtraItem = u8;
+ #[inline]
+ /// Get the count of extra items (trailing null byte for FFI compatibility)
+ fn extra_items_count(this: &Self) -> usize {
+ // extra item is the trailing \0 for ffi compatibility
+ return this.data.size + 1;
+ }
+}
+
+impl String {
+ /// Create a new empty String container
+ pub fn new() -> Self {
+ Self {
+ data: TVMFFIAny {
+ type_index: TypeIndex::kTVMFFISmallStr as i32,
+ small_str_len: 0,
+ data_union: TVMFFIAnyDataUnion { v_int64: 0 },
+ },
+ }
+ }
+
+ /// Get the length of the string in bytes
+ pub fn len(&self) -> usize {
+ self.as_bytes().len()
+ }
+
+ /// Get the string as a byte slice
+ pub fn as_bytes(&self) -> &[u8] {
+ unsafe {
+ if self.data.type_index == TypeIndex::kTVMFFISmallStr as i32 {
+ std::slice::from_raw_parts(
+ self.data.data_union.v_bytes.as_ptr(),
+ self.data.small_str_len as usize,
+ )
+ } else {
+ let str_obj: &StringObj = &*(self.data.data_union.v_obj as *const StringObj);
+ std::slice::from_raw_parts(str_obj.data.data, str_obj.data.size)
+ }
+ }
+ }
+
+ /// Get the string as a str slice
+ pub fn as_str(&self) -> &str {
+ unsafe { std::str::from_utf8_unchecked(self.as_bytes()) }
+ }
+}
+
+impl<T> From<T> for String
+where
+ T: AsRef<str>,
+{
+ #[inline]
+ /// Create String from any type that can be converted to a string slice
+ fn from(src: T) -> Self {
+ unsafe {
+ let value: &str = src.as_ref();
+ let bytes = value.as_bytes();
+ const MAX_SMALL_BYTES_LEN: usize = 7;
+ if bytes.len() <= MAX_SMALL_BYTES_LEN {
+ let mut data_union = TVMFFIAnyDataUnion { v_int64: 0 };
+ data_union.v_bytes[..bytes.len()].copy_from_slice(bytes);
+ Self {
+ data: TVMFFIAny {
+ type_index: TypeIndex::kTVMFFISmallStr as i32,
+ small_str_len: bytes.len() as u32,
+ data_union: data_union,
+ },
+ }
+ } else {
+ let mut obj_arc = ObjectArc::new_with_extra_items(StringObj {
+ object: Object::new(),
+ data: TVMFFIByteArray {
+ data: std::ptr::null(),
+ size: bytes.len(),
+ },
+ });
+ obj_arc.data.data = StringObj::extra_items(&obj_arc).as_ptr();
+ let extra_items = StringObj::extra_items_mut(&mut obj_arc);
+ extra_items[..bytes.len()].copy_from_slice(bytes);
+ // write the trailing \0 for ffi compatibility
+ extra_items[bytes.len()] = 0;
+ Self {
+ data: TVMFFIAny {
+ type_index: TypeIndex::kTVMFFIStr as i32,
+ small_str_len: 0,
+ data_union: TVMFFIAnyDataUnion {
+ v_obj: ObjectArc::into_raw(obj_arc) as *mut StringObj
+ as *mut TVMFFIObject,
+ },
+ },
+ }
+ }
+ }
+ }
+}
+
+impl Default for String {
+ #[inline]
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl Deref for String {
+ type Target = str;
+ #[inline]
+ fn deref(&self) -> &str {
+ self.as_str()
+ }
+}
+
+impl Clone for String {
+ #[inline]
+ fn clone(&self) -> Self {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { unsafe_::inc_ref(self.data.data_union.v_obj) }
+ }
+ Self { data: self.data }
+ }
+}
+
+impl Drop for String {
+ #[inline]
+ fn drop(&mut self) {
+ if self.data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { unsafe_::dec_ref(self.data.data_union.v_obj) }
+ }
+ }
+}
+
+impl PartialEq for String {
+ #[inline]
+ fn eq(&self, other: &Self) -> bool {
+ self.as_bytes() == other.as_bytes()
+ }
+}
+
+// Allows `my_string == "hello"`
+impl<T> PartialEq<T> for String
+where
+ T: AsRef<str>,
+{
+ #[inline]
+ fn eq(&self, other: &T) -> bool {
+ self.as_str() == other.as_ref()
+ }
+}
+
+impl Eq for String {}
+
+impl<T> PartialEq<T> for Bytes
+where
+ T: AsRef<[u8]>,
+{
+ #[inline]
+ fn eq(&self, other: &T) -> bool {
+ self.as_slice() == other.as_ref()
+ }
+}
+
+impl PartialOrd for String {
+ #[inline]
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ self.as_str().partial_cmp(other.as_str())
+ }
+}
+
+impl Ord for String {
+ #[inline]
+ fn cmp(&self, other: &Self) -> Ordering {
+ self.as_str().cmp(other.as_str())
+ }
+}
+
+impl Hash for String {
+ #[inline]
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.as_str().hash(state);
+ }
+}
+
+impl Display for String {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ Debug::fmt(&self.as_str(), f)
+ }
+}
+
+impl Debug for String {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("ffi.String")
+ .field("data", &self.as_str())
+ .finish()
+ }
+}
+
+//-----------------------------------------------------
+// AnyCompatible implementation for Bytes and String
+//-----------------------------------------------------
+unsafe impl AnyCompatible for Bytes {
+ fn type_str() -> std::string::String {
+ "ffi.Bytes".to_string()
+ }
+
+ unsafe fn copy_to_any_view(this: &Self, data: &mut TVMFFIAny) {
+ *data = this.data;
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ *data = src.data;
+ std::mem::forget(src);
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ return data.type_index == TypeIndex::kTVMFFISmallBytes as i32
+ || data.type_index == TypeIndex::kTVMFFIBytes as i32;
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ if data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { unsafe_::inc_ref(data.data_union.v_obj) }
+ }
+ Self { data: *data }
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ Self { data: *data }
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIByteArrayPtr as i32 {
+ // deep copy from bytearray ptr
+ let bytes = &*(data.data_union.v_ptr as *const TVMFFIByteArray);
+ Ok(Self::from(std::slice::from_raw_parts(
+ bytes.data, bytes.size,
+ )))
+ } else if data.type_index == TypeIndex::kTVMFFISmallBytes as i32 {
+ Ok(Self { data: *data })
+ } else if data.type_index == TypeIndex::kTVMFFIBytes as i32 {
+ unsafe { unsafe_::inc_ref(data.data_union.v_obj) }
+ Ok(Self { data: *data })
+ } else {
+ Err(())
+ }
+ }
+}
+
+unsafe impl AnyCompatible for String {
+ fn type_str() -> std::string::String {
+ "ffi.String".to_string()
+ }
+
+ unsafe fn copy_to_any_view(this: &Self, data: &mut TVMFFIAny) {
+ *data = this.data;
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ *data = src.data;
+ std::mem::forget(src);
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ return data.type_index == TypeIndex::kTVMFFISmallStr as i32
+ || data.type_index == TypeIndex::kTVMFFIStr as i32;
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ if data.type_index >= TypeIndex::kTVMFFIStaticObjectBegin as i32 {
+ unsafe { unsafe_::inc_ref(data.data_union.v_obj) }
+ }
+ Self { data: *data }
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ Self { data: *data }
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIRawStr as i32 {
+ // 1. Create a CStr wrapper from the raw pointer.
+ let c_str =
+ std::ffi::CStr::from_ptr(data.data_union.v_c_str as *const std::os::raw::c_char);
+ Ok(Self::from(c_str.to_str().expect("Invalid UTF-8")))
+ } else if data.type_index == TypeIndex::kTVMFFISmallStr as i32 {
+ Ok(Self { data: *data })
+ } else if data.type_index == TypeIndex::kTVMFFIStr as i32 {
+ unsafe { unsafe_::inc_ref(data.data_union.v_obj) }
+ Ok(Self { data: *data })
+ } else {
+ Err(())
+ }
+ }
+}
diff --git a/rust/tvm-ffi/src/type_traits.rs b/rust/tvm-ffi/src/type_traits.rs
new file mode 100644
index 0000000..d39da4b
--- /dev/null
+++ b/rust/tvm-ffi/src/type_traits.rs
@@ -0,0 +1,333 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIGetTypeInfo};
+
+//-----------------------------------------------------
+// AnyCompatible
+//-----------------------------------------------------
+/// Trait to enable a value to be compatible with Any
+/// Enables TryFrom/Into AnyView/Any
+pub unsafe trait AnyCompatible: Sized {
+ /// the value to copy to TVMFFIAny
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny);
+ /// consume the value to move to Any
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny);
+ // check if the value is compatible with the type
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool;
+
+ /// Copy value from TVMFFIAny after checking
+ /// caller must ensure that the value is compatible with the type
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self;
+ /// the value to move from TVMFFIAny
+ /// NOTE: pay very careful attention to avoid memory leak!
+ /// - When calling from managed Any, remember to use std::mem::ManuallyDrop
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self;
+ /// try to cast the value from AnyView
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()>;
+ /// Get the type key of a type when TryCastFromAnyView fails.
+ fn get_mismatch_type_info(data: &TVMFFIAny) -> String {
+ unsafe {
+ let info = TVMFFIGetTypeInfo(data.type_index);
+ assert!(!info.is_null(), "TVMFFIGetTypeInfo returned null");
+ (*info).type_key.as_str().to_string()
+ }
+ }
+ /// the type string of the type
+ fn type_str() -> String;
+}
+
+/// AnyCompatible for bool
+unsafe impl AnyCompatible for bool {
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIBool as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = *src as i64;
+ }
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ data.type_index == TypeIndex::kTVMFFIBool as i32
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ data.data_union.v_int64 != 0
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIBool as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = src as i64;
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ data.data_union.v_int64 != 0
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIBool as i32
+ || data.type_index == TypeIndex::kTVMFFIInt as i32
+ {
+ unsafe { Ok(data.data_union.v_int64 != 0) }
+ } else {
+ Err(())
+ }
+ }
+
+ fn type_str() -> String {
+ "bool".to_string()
+ }
+}
+
+/// Macro to implement AnyCompatible for integer types
+macro_rules! impl_any_compatible_for_int {
+ ($($int_type:ty),* $(,)?) => {
+ $(
+ unsafe impl AnyCompatible for $int_type {
+ fn type_str() -> String {
+ "int".to_string()
+ }
+
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIInt as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = *src as i64;
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ data.type_index == TypeIndex::kTVMFFIInt as i32
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ data.data_union.v_int64 as $int_type
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIInt as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = src as i64;
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ data.data_union.v_int64 as $int_type
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIInt as i32 ||
+ data.type_index == TypeIndex::kTVMFFIBool as i32 {
+ Ok(data.data_union.v_int64 as $int_type)
+ } else {
+ Err(())
+ }
+ }
+ }
+ )*
+ };
+}
+
+// Implement AnyCompatible for all integer types
+impl_any_compatible_for_int!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize);
+
+// Implement AnyCompatible for `Option<T>`
+unsafe impl<T: AnyCompatible> AnyCompatible for Option<T> {
+ fn type_str() -> String {
+ // make it consistent with c++ representation
+ "Optional<".to_string() + T::type_str().as_str() + ">"
+ }
+
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ if let Some(ref value) = src {
+ T::copy_to_any_view(value, data);
+ } else {
+ data.type_index = TypeIndex::kTVMFFINone as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = 0;
+ }
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ if let Some(value) = src {
+ T::move_to_any(value, data);
+ } else {
+ data.type_index = TypeIndex::kTVMFFINone as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = 0;
+ }
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ return T::check_any_strict(data) || data.type_index == TypeIndex::kTVMFFINone as i32;
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ if data.type_index == TypeIndex::kTVMFFINone as i32 {
+ None
+ } else {
+ Some(T::copy_from_any_view_after_check(data))
+ }
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ if data.type_index == TypeIndex::kTVMFFINone as i32 {
+ None
+ } else {
+ Some(T::move_from_any_after_check(data))
+ }
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFINone as i32 {
+ Ok(None)
+ } else {
+ T::try_cast_from_any_view(data).map(Some)
+ }
+ }
+}
+
+/// AnyCompatible for void*
+unsafe impl AnyCompatible for *mut core::ffi::c_void {
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIOpaquePtr as i32;
+ data.small_str_len = 0;
+ data.data_union.v_ptr = *src;
+ }
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ data.type_index == TypeIndex::kTVMFFIOpaquePtr as i32
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ data.data_union.v_ptr
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIOpaquePtr as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = src as i64;
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ data.data_union.v_ptr
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFINone as i32 {
+ Ok(std::ptr::null_mut())
+ } else if data.type_index == TypeIndex::kTVMFFIOpaquePtr as i32 {
+ unsafe { Ok(data.data_union.v_ptr) }
+ } else {
+ Err(())
+ }
+ }
+
+ fn type_str() -> String {
+ "void*".to_string()
+ }
+}
+
+/// Macro to implement AnyCompatible for integer types
+macro_rules! impl_any_compatible_for_float {
+ ($($float_type:ty),* $(,)?) => {
+ $(
+ unsafe impl AnyCompatible for $float_type {
+ fn type_str() -> String {
+ "float".to_string()
+ }
+
+ unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIFloat as i32;
+ data.small_str_len = 0;
+ data.data_union.v_float64 = *src as f64;
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ data.type_index == TypeIndex::kTVMFFIFloat as i32
+ }
+
+ unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self {
+ data.data_union.v_float64 as $float_type
+ }
+
+ unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFIFloat as i32;
+ data.small_str_len = 0;
+ data.data_union.v_float64 = src as f64;
+ }
+
+ unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self {
+ data.data_union.v_float64 as $float_type
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFIFloat as i32 {
+ Ok(data.data_union.v_float64 as $float_type)
+ } else if data.type_index == TypeIndex::kTVMFFIInt as i32 ||
+ data.type_index == TypeIndex::kTVMFFIBool as i32 {
+ Ok(data.data_union.v_int64 as $float_type)
+ } else {
+ Err(())
+ }
+ }
+ }
+ )*
+ };
+}
+
+// Implement AnyCompatible for all float types
+impl_any_compatible_for_float!(f32, f64);
+
+// Special rule: we convert () to None
+// This allows us to effectively pass () around as null value
+// note that this is indeed a bit relaxation of the type
+// but it is necessary for us to enable void/none interoperability
+unsafe impl AnyCompatible for () {
+ fn type_str() -> String {
+ // make it consistent with c++ representation
+ "None".to_string()
+ }
+
+ unsafe fn copy_to_any_view(_src: &Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFINone as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = 0;
+ }
+
+ unsafe fn move_to_any(_src: Self, data: &mut TVMFFIAny) {
+ data.type_index = TypeIndex::kTVMFFINone as i32;
+ data.small_str_len = 0;
+ data.data_union.v_int64 = 0;
+ }
+
+ unsafe fn check_any_strict(data: &TVMFFIAny) -> bool {
+ return data.type_index == TypeIndex::kTVMFFINone as i32;
+ }
+
+ unsafe fn copy_from_any_view_after_check(_data: &TVMFFIAny) -> Self {
+ ()
+ }
+
+ unsafe fn move_from_any_after_check(_data: &mut TVMFFIAny) -> Self {
+ ()
+ }
+
+ unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> {
+ if data.type_index == TypeIndex::kTVMFFINone as i32 {
+ Ok(())
+ } else {
+ Err(())
+ }
+ }
+}
diff --git a/rust/tvm-ffi/tests/test_any.rs b/rust/tvm-ffi/tests/test_any.rs
new file mode 100644
index 0000000..b950419
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_any.rs
@@ -0,0 +1,349 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+/// Macro to test integer types with both Any and AnyView
+macro_rules! test_any_int {
+ ($($type:ty: $value:expr, $any_name:ident, $any_view_name:ident),* $(,)?) => {
+ $(
+ #[test]
+ fn $any_name() {
+ let any = Any::from($value);
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIInt as i32);
+ assert_eq!(<$type>::try_from(any).unwrap(), $value);
+ }
+
+ #[test]
+ fn $any_view_name() {
+ let value = $value;
+ let any_view = AnyView::from(&value);
+ assert_eq!(any_view.type_index(), TypeIndex::kTVMFFIInt as i32);
+ assert_eq!(<$type>::try_from(any_view).unwrap(), $value);
+ }
+ )*
+ };
+}
+
+// Test all integer types with both Any and AnyView
+test_any_int!(
+ i8: 127i8, test_any_int_i8, test_any_view_int_i8,
+ i16: 32767i16, test_any_int_i16, test_any_view_int_i16,
+ i32: 42i32, test_any_int_i32, test_any_view_int_i32,
+ i64: 123456789i64, test_any_int_i64, test_any_view_int_i64,
+ isize: 42isize, test_any_int_isize, test_any_view_int_isize,
+ u8: 255u8, test_any_int_u8, test_any_view_int_u8,
+ u16: 65535u16, test_any_int_u16, test_any_view_int_u16,
+ u32: 4294967295u32, test_any_int_u32, test_any_view_int_u32,
+ u64: 18446744073709551615u64, test_any_int_u64, test_any_view_int_u64,
+ usize: 42usize, test_any_int_usize, test_any_view_int_usize
+);
+
+/// Macro to test float types with both Any and AnyView
+macro_rules! test_any_float {
+ ($($type:ty: $value:expr, $any_name:ident, $any_view_name:ident),* $(,)?) => {
+ $(
+ #[test]
+ fn $any_name() {
+ let any = Any::from($value);
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIFloat as i32);
+ assert_eq!(<$type>::try_from(any).unwrap(), $value);
+ }
+ )*
+ };
+}
+
+// Test all float types with both Any and AnyView
+test_any_float!(
+ f32: 3.14f32, test_any_float_f32, test_any_view_float_f32,
+ f64: 3.14f64, test_any_float_f64, test_any_view_float_f64
+);
+
+#[test]
+fn test_any_bool() {
+ // Test both Any and AnyView with both true and false
+ let any_true = Any::from(true);
+ let any_false = Any::from(false);
+
+ let value_true = true;
+ let value_false = false;
+ let any_view_true = AnyView::from(&value_true);
+ let any_view_false = AnyView::from(&value_false);
+
+ // Test Any
+ assert_eq!(any_true.type_index(), TypeIndex::kTVMFFIBool as i32);
+ assert_eq!(any_false.type_index(), TypeIndex::kTVMFFIBool as i32);
+ assert_eq!(bool::try_from(any_true).unwrap(), true);
+ assert_eq!(bool::try_from(any_false).unwrap(), false);
+
+ // Test AnyView
+ assert_eq!(any_view_true.type_index(), TypeIndex::kTVMFFIBool as i32);
+ assert_eq!(any_view_false.type_index(), TypeIndex::kTVMFFIBool as i32);
+ assert_eq!(bool::try_from(any_view_true).unwrap(), true);
+ assert_eq!(bool::try_from(any_view_false).unwrap(), false);
+}
+
+#[test]
+fn test_type_mismatch_error() {
+ let any = Any::from(42i32);
+ // try from will try casting when possible
+ // Try to convert int to bool - should be ok
+ let result: Result<bool, _> = bool::try_from(any);
+ assert!(result.is_ok());
+ // try_as is more strict and do not allow conversion from int to bool
+ let any = Any::from(42i32);
+ let opt_try_as = any.try_as::<bool>();
+ assert!(opt_try_as.is_none());
+}
+
+#[test]
+fn test_type_mismatch_error_message() {
+ let any_none = Any::from(Option::<i32>::from(None));
+ // try from will try casting when possible
+ // Try to convert int to bool - should be ok
+ let err = i32::try_from(any_none).unwrap_err();
+ assert_eq!(err.kind(), crate::error::TYPE_ERROR);
+ assert!(err.message().contains("`None` to `int`"));
+
+ let any_float = Any::from(1.2f32);
+ let err = Option::<i32>::try_from(any_float).unwrap_err();
+ assert_eq!(err.kind(), crate::error::TYPE_ERROR);
+ assert!(err.message().contains("`float` to `Optional<int>`"));
+}
+
+#[test]
+fn test_cross_type_conversion_int_to_bool() {
+ // Test that integers can be converted to bool (0 = false, non-zero = true)
+ let any_true = Any::from(42i32);
+ let any_false = Any::from(0i32);
+
+ // This should work because our macro allows conversion from int to bool
+ let bool_true: Result<bool, _> = bool::try_from(any_true);
+ let bool_false: Result<bool, _> = bool::try_from(any_false);
+
+ // Note: This test might fail depending on how try_cast_from_any_view is implemented
+ // If it fails, it means the conversion isn't implemented yet
+ assert_eq!(bool_true.unwrap(), true);
+ assert_eq!(bool_false.unwrap(), false);
+}
+
+#[test]
+fn test_max_min_values() {
+ // Test maximum and minimum values for i64
+ let any_i64_max = Any::from(i64::MAX);
+ let any_i64_min = Any::from(i64::MIN);
+
+ assert_eq!(i64::try_from(any_i64_max).unwrap(), i64::MAX);
+ assert_eq!(i64::try_from(any_i64_min).unwrap(), i64::MIN);
+}
+
+#[test]
+fn test_any_void_ptr() {
+ let any = Any::from(std::ptr::null_mut());
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIOpaquePtr as i32);
+ assert_eq!(
+ <*mut core::ffi::c_void>::try_from(any).unwrap(),
+ std::ptr::null_mut()
+ );
+}
+
+#[test]
+fn test_any_option_i32() {
+ // Test Option<i32> with Some value
+ let some_value = Some(42i32);
+ let any_some = Any::from(some_value);
+ let any_view_some = AnyView::from(&some_value);
+
+ // Test Any with Some value
+ assert_eq!(any_some.type_index(), TypeIndex::kTVMFFIInt as i32);
+ assert_eq!(Option::<i32>::try_from(any_some).unwrap(), Some(42i32));
+
+ // Test AnyView with Some value
+ assert_eq!(any_view_some.type_index(), TypeIndex::kTVMFFIInt as i32);
+ assert_eq!(Option::<i32>::try_from(any_view_some).unwrap(), Some(42i32));
+
+ // Test Option<i32> with None value
+ let none_value: Option<i32> = None;
+ let any_none = Any::from(none_value);
+ let any_view_none = AnyView::from(&none_value);
+
+ // Test Any with None value
+ assert_eq!(any_none.type_index(), TypeIndex::kTVMFFINone as i32);
+ assert_eq!(Option::<i32>::try_from(any_none).unwrap(), None::<i32>);
+
+ // Test AnyView with None value
+ assert_eq!(any_view_none.type_index(), TypeIndex::kTVMFFINone as i32);
+ assert_eq!(Option::<i32>::try_from(any_view_none).unwrap(), None::<i32>);
+
+ let any_float = Any::from(1.2f32);
+ assert_eq!(any_float.type_index(), TypeIndex::kTVMFFIFloat as i32);
+ assert_eq!(
+ Option::<i32>::try_from(any_float).unwrap_err().kind(),
+ TYPE_ERROR
+ );
+}
+
+#[test]
+fn test_any_string() {
+ let any = Any::from(String::from("hello"));
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFISmallStr as i32);
+
+ // Check reference count - small strings are not ref counted
+ assert_eq!(any.debug_strong_count(), None);
+
+ let out_string = String::try_from(any).unwrap();
+ assert_eq!(out_string, "hello");
+}
+
+#[test]
+fn test_any_bytes() {
+ let data_arr: &[u8; 3] = &[1, 2, 3];
+ let any = Any::from(Bytes::from(data_arr));
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFISmallBytes as i32);
+
+ // Check reference count - small bytes are not ref counted
+ assert_eq!(any.debug_strong_count(), None);
+
+ let out_bytes = Bytes::try_from(any).unwrap();
+ assert_eq!(out_bytes, data_arr);
+}
+
+#[test]
+fn test_any_big_string() {
+ // Use a string longer than 7 characters to trigger big string allocation
+ let long_string = "hello world this is a long string";
+ let input_str = String::from(long_string);
+ let any = Any::from(input_str.clone());
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIStr as i32);
+
+ // Check reference count - big strings are ref counted
+ assert_eq!(any.debug_strong_count(), Some(2));
+
+ let out_string = String::try_from(any).unwrap();
+ assert_eq!(out_string, long_string);
+
+ drop(out_string);
+ assert_eq!(AnyView::from(&input_str).debug_strong_count(), Some(1));
+}
+
+#[test]
+fn test_any_big_bytes() {
+ // Use bytes longer than 7 bytes to trigger big bytes allocation
+ let data_arr: &[u8; 10] = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+ let any = Any::from(Bytes::from(data_arr));
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIBytes as i32);
+
+ // Check reference count - big bytes are ref counted
+ assert_eq!(any.debug_strong_count(), Some(1));
+
+ let out_bytes = Bytes::try_from(any).unwrap();
+ assert_eq!(out_bytes, data_arr);
+
+ assert_eq!(AnyView::from(&out_bytes).debug_strong_count(), Some(1));
+}
+
+#[test]
+fn test_anyview_any_conversion_and_cloning() {
+ let input_str = String::from("hello world this is a long string");
+
+ // Test AnyView behavior - should not increment reference count
+ let any_view = AnyView::from(&input_str);
+ assert_eq!(any_view.debug_strong_count(), Some(1));
+
+ // Test AnyView cloning - should not affect reference count
+ let any_view_clone = any_view.clone();
+ assert_eq!(any_view.debug_strong_count(), Some(1));
+ assert_eq!(any_view_clone.debug_strong_count(), Some(1));
+
+ // Test AnyView to Any conversion - should increment reference count
+ let any1 = Any::from(any_view);
+ assert_eq!(any1.debug_strong_count(), Some(2));
+
+ // Test Any cloning - should increment reference count
+ let any2 = any1.clone();
+ assert_eq!(any1.debug_strong_count(), Some(3));
+ assert_eq!(any2.debug_strong_count(), Some(3));
+
+ // Test Any to AnyView conversion - should not change reference count
+ let any_view_from_any = AnyView::from(&any1);
+ assert_eq!(any_view_from_any.debug_strong_count(), Some(3));
+
+ // Test AnyView to Any conversion again
+ let any3 = Any::from(any_view_clone);
+ assert_eq!(any3.debug_strong_count(), Some(4));
+
+ // Test data integrity through conversions
+ let out_string1 = String::try_from(any1.clone()).unwrap();
+ let out_string2 = String::try_from(any2.clone()).unwrap();
+ let out_string3 = String::try_from(any3.clone()).unwrap();
+ let out_string_view = String::try_from(any_view_from_any).unwrap();
+
+ assert_eq!(out_string1, input_str);
+ assert_eq!(out_string2, input_str);
+ assert_eq!(out_string3, input_str);
+ assert_eq!(out_string_view, input_str);
+
+ // Test reference count cleanup
+ drop(any1);
+ assert_eq!(any3.debug_strong_count(), Some(7));
+
+ drop(any2);
+ assert_eq!(any3.debug_strong_count(), Some(6));
+
+ drop(any3);
+ assert_eq!(AnyView::from(&input_str).debug_strong_count(), Some(5));
+}
+
+#[test]
+fn test_any_dl_device() {
+ use tvm_ffi::{DLDevice, DLDeviceType};
+
+ // Test DLDevice with CPU device
+ let cpu_device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let any_cpu = Any::from(cpu_device);
+ let any_view_cpu = AnyView::from(&cpu_device);
+
+ // Test Any with CPU device
+ assert_eq!(any_cpu.type_index(), TypeIndex::kTVMFFIDevice as i32);
+ let converted_cpu = DLDevice::try_from(any_cpu).unwrap();
+ assert_eq!(converted_cpu.device_type, DLDeviceType::kDLCPU);
+ assert_eq!(converted_cpu.device_id, 0);
+
+ // Test AnyView with CPU device
+ assert_eq!(any_view_cpu.type_index(), TypeIndex::kTVMFFIDevice as i32);
+ let converted_cpu_view = DLDevice::try_from(any_view_cpu).unwrap();
+ assert_eq!(converted_cpu_view.device_type, DLDeviceType::kDLCPU);
+ assert_eq!(converted_cpu_view.device_id, 0);
+
+ // Test DLDevice with CUDA device
+ let cuda_device = DLDevice::new(DLDeviceType::kDLCUDA, 1);
+ let any_cuda = Any::from(cuda_device);
+ let any_view_cuda = AnyView::from(&cuda_device);
+
+ // Test Any with CUDA device
+ assert_eq!(any_cuda.type_index(), TypeIndex::kTVMFFIDevice as i32);
+ let converted_cuda = DLDevice::try_from(any_cuda).unwrap();
+ assert_eq!(converted_cuda.device_type, DLDeviceType::kDLCUDA);
+ assert_eq!(converted_cuda.device_id, 1);
+
+ // Test AnyView with CUDA device
+ assert_eq!(any_view_cuda.type_index(), TypeIndex::kTVMFFIDevice as i32);
+ let converted_cuda_view = DLDevice::try_from(any_view_cuda).unwrap();
+ assert_eq!(converted_cuda_view.device_type, DLDeviceType::kDLCUDA);
+ assert_eq!(converted_cuda_view.device_id, 1);
+}
diff --git a/rust/tvm-ffi/tests/test_device.rs b/rust/tvm-ffi/tests/test_device.rs
new file mode 100644
index 0000000..cc41310
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_device.rs
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+#[test]
+fn test_device_stream() {
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let dummy_stream = 2 as TVMFFIStreamHandle;
+ with_stream(&device, dummy_stream, || {
+ assert_eq!(current_stream(&device), dummy_stream);
+ Ok(())
+ })
+ .unwrap();
+ assert_eq!(current_stream(&device), std::ptr::null_mut());
+}
+
+#[test]
+fn test_device_any_conversion() {
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let any_device: Any = device.into();
+ assert_eq!(any_device.type_index(), TypeIndex::kTVMFFIDevice as i32);
+ let converted_device: DLDevice = any_device.try_into().unwrap();
+ assert_eq!(converted_device.device_type, device.device_type);
+ assert_eq!(converted_device.device_id, device.device_id);
+}
diff --git a/rust/tvm-ffi/tests/test_dtype.rs b/rust/tvm-ffi/tests/test_dtype.rs
new file mode 100644
index 0000000..c0a28fb
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_dtype.rs
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+#[test]
+fn test_dl_datatype_from_string() {
+ // Test valid dtype strings
+ let test_cases = vec![
+ (
+ "int32",
+ DLDataType {
+ code: DLDataTypeCode::kDLInt as u8,
+ bits: 32,
+ lanes: 1,
+ },
+ ),
+ (
+ "float32",
+ DLDataType {
+ code: DLDataTypeCode::kDLFloat as u8,
+ bits: 32,
+ lanes: 1,
+ },
+ ),
+ (
+ "int8",
+ DLDataType {
+ code: DLDataTypeCode::kDLInt as u8,
+ bits: 8,
+ lanes: 1,
+ },
+ ),
+ (
+ "float64",
+ DLDataType {
+ code: DLDataTypeCode::kDLFloat as u8,
+ bits: 64,
+ lanes: 1,
+ },
+ ),
+ (
+ "uint32",
+ DLDataType {
+ code: DLDataTypeCode::kDLUInt as u8,
+ bits: 32,
+ lanes: 1,
+ },
+ ),
+ ];
+
+ for item in test_cases {
+ let dtype = DLDataType::try_from_str(item.0).unwrap();
+ assert_eq!(dtype, item.1);
+ }
+}
+
+#[test]
+fn test_dl_datatype_any_conversion() {
+ let original_dtype = DLDataType {
+ code: DLDataTypeCode::kDLInt as u8,
+ bits: 32,
+ lanes: 1,
+ };
+
+ // Test conversion to Any
+ let any_dtype: Any = original_dtype.into();
+
+ // Test conversion back from Any
+ let converted_dtype: DLDataType = any_dtype.try_into().unwrap();
+ assert_eq!(converted_dtype.code, original_dtype.code);
+ assert_eq!(converted_dtype.bits, original_dtype.bits);
+ assert_eq!(converted_dtype.lanes, original_dtype.lanes);
+}
+
+#[test]
+fn test_dl_datatype_any_with_string_conversion() {
+ // Test that DLDataType can be converted from string via Any
+ let dtype_str = "int32";
+
+ let dtype = DLDataType::try_from_str(dtype_str).unwrap();
+
+ // Convert to Any
+ let any_dtype: Any = dtype.into();
+
+ // Convert back from Any
+ let converted_dtype: DLDataType = any_dtype.try_into().unwrap();
+ assert_eq!(converted_dtype.code, dtype.code);
+ assert_eq!(converted_dtype.bits, dtype.bits);
+ assert_eq!(converted_dtype.lanes, dtype.lanes);
+}
diff --git a/rust/tvm-ffi/tests/test_error.rs b/rust/tvm-ffi/tests/test_error.rs
new file mode 100644
index 0000000..ba48dc9
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_error.rs
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+#[test]
+fn test_error_new() {
+ let error = Error::new(RUNTIME_ERROR, "test error", "test traceback");
+ assert_eq!(error.kind(), RUNTIME_ERROR);
+ assert_eq!(error.message(), "test error");
+ assert_eq!(error.backtrace(), "test traceback");
+}
+
+#[test]
+fn test_error_from_raised() {
+ let error0 = Error::new(RUNTIME_ERROR, "test error", "test traceback");
+ Error::set_raised(&error0);
+ let error1 = Error::from_raised();
+ assert_eq!(error1.kind(), RUNTIME_ERROR);
+ assert_eq!(error1.message(), "test error");
+ assert_eq!(error1.backtrace(), "test traceback");
+ // we have two references because one in error, and another one in the function call
+ assert_eq!(ObjectArc::strong_count(Error::data(&error1)), 2);
+}
+
+fn error_fn0(flag: bool) -> Result<i32> {
+ // if flag is true, throw an error
+ ensure!(!flag, RUNTIME_ERROR, "test error {flag}");
+ Ok(1)
+}
+
+fn error_fn1(flag: bool) -> Result<i32> {
+ attach_context!(error_fn0(flag))
+}
+
+#[test]
+fn test_error_with_context() {
+ let error0 = attach_context!(error_fn1(true)).unwrap_err();
+ assert_eq!(error0.kind(), RUNTIME_ERROR);
+ assert_eq!(error0.message(), "test error true");
+ assert!(error0.backtrace().contains("test_error.rs"));
+ assert!(error0.backtrace().contains("error_fn0"));
+ assert!(error0.backtrace().contains("error_fn1"));
+ assert!(error0.backtrace().contains("test_error_with_context"));
+ let result = error_fn1(false).unwrap();
+ assert_eq!(result, 1);
+}
+
+#[test]
+fn test_error_any_convert() {
+ let error = Error::new(RUNTIME_ERROR, "test error", "test traceback");
+ let any = Any::from(error.clone());
+ let error2 = any.clone().try_as::<Error>().unwrap();
+ assert_eq!(error2.kind(), RUNTIME_ERROR);
+ assert_eq!(error2.message(), "test error");
+ assert_eq!(error2.backtrace(), "test traceback");
+
+ let res = i32::try_from(any);
+ assert!(res.is_err());
+ assert!(res.unwrap_err().message().contains("`ffi.Error` to `int`"));
+
+ let anyview = AnyView::from(&error);
+ let error3 = Error::try_from(anyview).unwrap();
+ assert_eq!(error3.kind(), RUNTIME_ERROR);
+ assert_eq!(error3.message(), "test error");
+ assert_eq!(error3.backtrace(), "test traceback");
+}
diff --git a/rust/tvm-ffi/tests/test_function.rs b/rust/tvm-ffi/tests/test_function.rs
new file mode 100644
index 0000000..fef5e34
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_function.rs
@@ -0,0 +1,194 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+#[test]
+fn test_function_get_global_required() {
+ let fecho = Function::get_global("testing.echo").unwrap();
+ let a = 1;
+ let args = [AnyView::from(&a)];
+ let result = fecho.call_packed(&args).unwrap();
+ assert_eq!(i32::try_from(result).unwrap(), 1);
+}
+
+#[test]
+fn test_function_from_packed() {
+ let value = 2;
+ let v2 = 4;
+ let check_and_add_value = Function::from_packed(move |args: &[AnyView]| -> Result<Any> {
+ ensure!(
+ args.len() == 1,
+ VALUE_ERROR,
+ "Expected 1 argument, got {}",
+ args.len()
+ );
+ let v0 = i32::try_from(args[0])?;
+ ensure!(v0 == value, VALUE_ERROR, "Expected {}, got {}", value, v0);
+ Ok(Any::from(v0 + v2))
+ });
+ let args = [AnyView::from(&value)];
+ let result = check_and_add_value.call_packed(&args).unwrap();
+ assert_eq!(i32::try_from(result).unwrap(), 6);
+}
+
+#[test]
+fn test_function_from_typed() {
+ let offset = 2;
+ // test one argument
+ let sum1 = Function::from_typed(move |x: i32| -> Result<i32> { Ok(x + offset) });
+ let result = sum1.call_packed(&[AnyView::from(&1)]).unwrap();
+ assert_eq!(i32::try_from(result).unwrap(), 1 + offset);
+ // test two arguments
+ let sum2 = Function::from_typed(move |x: i32, y: i32| -> Result<i32> { Ok(x + y) });
+ let result = sum2
+ .call_packed(&[AnyView::from(&1), AnyView::from(&2)])
+ .unwrap();
+ assert_eq!(i32::try_from(result).unwrap(), 3);
+ // test three arguments
+ let sum3f = Function::from_typed(move |x: i32, y: i32, z: f32| -> Result<f32> {
+ Ok((x + y) as f32 + z)
+ });
+ let result = sum3f
+ .call_packed(&[AnyView::from(&1), AnyView::from(&2), AnyView::from(&3)])
+ .unwrap();
+ assert_eq!(f32::try_from(result).unwrap(), 6.0);
+}
+
+#[test]
+fn test_function_call_tuple() {
+ let offset = 2;
+ // test one argument
+ let sum1 = Function::from_typed(move |x: i32| -> Result<i32> { Ok(x + offset) });
+ let result = sum1.call_tuple((1,)).unwrap();
+ assert_eq!(i32::try_from(result).unwrap(), 1 + offset);
+ // test pass by reference
+ let result = sum1.call_tuple_with_len::<1, _>((&1,)).unwrap();
+ assert_eq!(i32::try_from(result).unwrap(), 1 + offset);
+ let typed_fn = |x: &i32| -> Result<i32> { Ok(sum1.call_tuple((x,))?.try_into()?) };
+ let result = typed_fn(&1);
+ assert_eq!(result.unwrap(), 1 + offset);
+}
+
+#[test]
+fn test_function_into_typed_fn() {
+ let offset = 2;
+ let typed_sum1 = into_typed_fn!(
+ Function::from_typed(move |x: i32| -> Result<i32> { Ok(x + offset) }),
+ Fn(&i32) -> Result<i32>);
+ assert_eq!(typed_sum1(&1).unwrap(), 1 + offset);
+ // try to box the resulting function
+ let sum2 = Function::from_typed(move |x: i32, y: i32| -> Result<i32> { Ok(x + y) });
+ let typed_sum2 = Box::new(into_typed_fn!(sum2, Fn(&i32, i32) -> Result<i32>));
+ assert_eq!(typed_sum2(&1, 2).unwrap(), 3);
+
+ // test three arguments
+ let sum3 = Function::from_typed(move |x: i32, y: i32, z: f32| -> Result<f32> {
+ Ok((x + y) as f32 + z)
+ });
+ let typed_sum3 = Box::new(into_typed_fn!(sum3, Fn(&i32, i32, f32) -> Result<f32>));
+ assert_eq!(typed_sum3(&1, 2, 3.0).unwrap(), 6.0);
+}
+
+#[test]
+fn test_function_echo_tensor_typed() {
+ let echo = into_typed_fn!(
+ Function::get_global("testing.echo").unwrap(),
+ Fn(&Tensor) -> Result<Tensor>
+ );
+ let data: &[f32] = &[1.0, 2.0, 3.0, 4.0];
+ let tensor = Tensor::from_slice(data, &[1, 2, 2]).unwrap();
+ // write tensor content here
+ let result = echo(&tensor).unwrap();
+ assert_eq!(result.data_ptr(), tensor.data_ptr());
+ assert_eq!(AnyView::from(&result).debug_strong_count(), Some(2));
+ let result_data = result.data_as_slice::<f32>().unwrap();
+ assert_eq!(result_data.len(), 4);
+ assert_eq!(result_data[0], 1.0);
+ assert_eq!(result_data[1], 2.0);
+ assert_eq!(result_data[2], 3.0);
+ assert_eq!(result_data[3], 4.0);
+}
+
+fn testing_add_one(x: i32) -> Result<i32> {
+ Ok(x + 1)
+}
+tvm_ffi_dll_export_typed_func!(testing_add_one, testing_add_one);
+
+#[test]
+fn test_function_from_extern_c() {
+ let add_one = Function::from_extern_c(std::ptr::null_mut(), __tvm_ffi_testing_add_one, None);
+ let typed_add_one = into_typed_fn!(add_one, Fn(i32) -> Result<i32>);
+ assert_eq!(typed_add_one(1).unwrap(), 2);
+}
+
+#[test]
+fn test_function_echo_string_bytes() {
+ let echo = Function::get_global("testing.echo").unwrap();
+ let echo_str = into_typed_fn!(
+ echo.clone(),
+ Fn(&str) -> Result<String>
+ );
+ let result = echo_str("hello").unwrap();
+ assert_eq!(result, "hello");
+ let echo_bytes = into_typed_fn!(
+ echo.clone(),
+ Fn(&[u8]) -> Result<Bytes>
+ );
+ let result = echo_bytes(b"hello").unwrap();
+ assert_eq!(result, b"hello");
+}
+
+#[test]
+fn test_function_apply() {
+ let add_one = Function::from_typed(|x: i32| -> Result<i32> { Ok(x + 1) });
+ let fapply = into_typed_fn!(
+ Function::get_global("testing.apply").unwrap(),
+ Fn(Function, i32) -> Result<i32>
+ );
+ let result = fapply(add_one, 3).unwrap();
+ assert_eq!(result, 4);
+}
+
+fn test_add_one_tensor(x: tvm_ffi::Tensor, y: tvm_ffi::Tensor) -> Result<()> {
+ let x_data = x.data_as_slice::<f32>()?;
+ let y_data = y.data_as_slice_mut::<f32>()?;
+ for i in 0..x_data.len() {
+ y_data[i] = x_data[i] + 1.0;
+ }
+ Ok(())
+}
+
+tvm_ffi_dll_export_typed_func!(test_add_one_tensor, test_add_one_tensor);
+
+#[test]
+fn test_function_call_tensor_fn() {
+ let add_one =
+ Function::from_extern_c(std::ptr::null_mut(), __tvm_ffi_test_add_one_tensor, None);
+ let typed_add_one = into_typed_fn!(add_one, Fn(&Tensor, &Tensor) -> Result<()>);
+ let x_data: &[f32] = &[0.0, 1.0, 2.0, 3.0];
+ let x = Tensor::from_slice(x_data, &[2, 2]).unwrap();
+ let y_data: &[f32] = &[0.0, 0.0, 0.0, 0.0];
+ let y = Tensor::from_slice(y_data, &[2, 2]).unwrap();
+ typed_add_one(&x, &y).unwrap();
+ let y_data = y.data_as_slice::<f32>().unwrap();
+ assert_eq!(y_data[0], 1.0);
+ assert_eq!(y_data[1], 2.0);
+ assert_eq!(y_data[2], 3.0);
+ assert_eq!(y_data[3], 4.0);
+}
diff --git a/rust/tvm-ffi/tests/test_object.rs b/rust/tvm-ffi/tests/test_object.rs
new file mode 100644
index 0000000..60378c2
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_object.rs
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::sync::Arc;
+use tvm_ffi::*;
+
+// must have repr(C) for the object header stays in the same position
+#[repr(C)]
+struct TestIntObj {
+ object: Object,
+ pub value: i64,
+ // counter for recording the number of times the object is deleted
+ delete_counter: Arc<AtomicU32>,
+ pub extra_item_count: u64,
+}
+
+impl TestIntObj {
+ pub fn new(value: i64, delete_counter: Arc<AtomicU32>, extra_item_count: u64) -> Self {
+ Self {
+ object: Object::new(),
+ value,
+ delete_counter,
+ extra_item_count,
+ }
+ }
+}
+
+impl Drop for TestIntObj {
+ fn drop(&mut self) {
+ self.delete_counter.fetch_add(1, Ordering::Relaxed);
+ }
+}
+
+unsafe impl ObjectCore for TestIntObj {
+ const TYPE_KEY: &'static str = Object::TYPE_KEY;
+ #[inline]
+ fn type_index() -> i32 {
+ Object::type_index()
+ }
+ #[inline]
+ unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject {
+ Object::object_header_mut(&mut this.object)
+ }
+}
+
+unsafe impl ObjectCoreWithExtraItems for TestIntObj {
+ type ExtraItem = u64;
+ #[inline]
+ fn extra_items_count(this: &Self) -> usize {
+ this.extra_item_count as usize
+ }
+}
+
+#[test]
+fn test_object_arc() {
+ let delete_counter = Arc::new(AtomicU32::new(0));
+ let obj_arc = ObjectArc::new(TestIntObj::new(11, delete_counter.clone(), 0));
+ assert_eq!(obj_arc.value, 11);
+ assert_eq!(ObjectArc::strong_count(&obj_arc), 1);
+ assert_eq!(ObjectArc::weak_count(&obj_arc), 1);
+
+ let ref1 = obj_arc.clone();
+ assert_eq!(ObjectArc::strong_count(&obj_arc), 2);
+ assert_eq!(ObjectArc::weak_count(&obj_arc), 1);
+
+ let ref2 = obj_arc.clone();
+ assert_eq!(ObjectArc::strong_count(&obj_arc), 3);
+ assert_eq!(ObjectArc::weak_count(&obj_arc), 1);
+ assert_eq!(ref1.value, 11);
+ // drop obj_arc
+ drop(obj_arc);
+ assert_eq!(ObjectArc::strong_count(&ref1), 2);
+ assert_eq!(ObjectArc::weak_count(&ref1), 1);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 0);
+ // drop ref1
+ drop(ref1);
+ assert_eq!(ObjectArc::strong_count(&ref2), 1);
+ assert_eq!(ObjectArc::weak_count(&ref2), 1);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 0);
+ // drop ref2
+ drop(ref2);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 1);
+}
+
+#[test]
+fn test_object_arc_with_extra_items() {
+ let delete_counter = Arc::new(AtomicU32::new(0));
+ let mut obj_arc =
+ ObjectArc::new_with_extra_items(TestIntObj::new(12, delete_counter.clone(), 10));
+ assert_eq!(obj_arc.value, 12);
+ assert_eq!(ObjectArc::strong_count(&obj_arc), 1);
+ assert_eq!(ObjectArc::weak_count(&obj_arc), 1);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 0);
+ unsafe {
+ // layout check of extra items
+ assert_eq!(TestIntObj::extra_items_count(&obj_arc), 10);
+ assert_eq!(TestIntObj::extra_items(&obj_arc).len(), 10);
+ assert_eq!(TestIntObj::extra_items_mut(&mut obj_arc).len(), 10);
+ assert_eq!(
+ TestIntObj::extra_items_mut(&mut obj_arc).as_ptr() as *mut u8,
+ (ObjectArc::as_raw_mut(&mut obj_arc) as *mut u8).add(std::mem::size_of::<TestIntObj>())
+ );
+ }
+ drop(obj_arc);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 1);
+}
+
+#[test]
+fn test_object_arc_from_raw() {
+ unsafe {
+ let delete_counter = Arc::new(AtomicU32::new(0));
+ let obj_arc = ObjectArc::new(TestIntObj::new(11, delete_counter.clone(), 0));
+ let raw_ptr = ObjectArc::into_raw(obj_arc);
+ let obj_arc2 = ObjectArc::from_raw(raw_ptr);
+ assert_eq!(obj_arc2.value, 11);
+ assert_eq!(ObjectArc::strong_count(&obj_arc2), 1);
+ assert_eq!(ObjectArc::weak_count(&obj_arc2), 1);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 0);
+ // drop obj_arc2
+ drop(obj_arc2);
+ assert_eq!(delete_counter.load(Ordering::Relaxed), 1);
+ }
+}
+
+#[test]
+fn test_object_arc_option_size() {
+ assert_eq!(
+ std::mem::size_of::<Option<ObjectArc<TestIntObj>>>(),
+ std::mem::size_of::<ObjectArc<TestIntObj>>()
+ );
+}
diff --git a/rust/tvm-ffi/tests/test_shape.rs b/rust/tvm-ffi/tests/test_shape.rs
new file mode 100644
index 0000000..dcbb53d
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_shape.rs
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+// ============================================================================
+// Shape Tests
+// ============================================================================
+
+#[test]
+fn test_any_shape() {
+ let shape = Shape::from(vec![1, 2, 3, 4]);
+ let any = Any::from(shape.clone());
+ let any_view = AnyView::from(&shape);
+
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFIShape as i32);
+ let converted = Shape::try_from(any).unwrap();
+ assert_eq!(converted.as_slice(), &[1, 2, 3, 4]);
+
+ assert_eq!(any_view.type_index(), TypeIndex::kTVMFFIShape as i32);
+ let converted_view = Shape::try_from(any_view).unwrap();
+ assert_eq!(converted_view.as_slice(), &[1, 2, 3, 4]);
+}
diff --git a/rust/tvm-ffi/tests/test_string.rs b/rust/tvm-ffi/tests/test_string.rs
new file mode 100644
index 0000000..8d8384f
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_string.rs
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
+use tvm_ffi::*;
+
+// ============================================================================
+// Bytes Tests
+// ============================================================================
+
+#[test]
+fn test_bytes_new() {
+ let bytes = Bytes::new();
+ assert_eq!(bytes.len(), 0);
+ assert_eq!(bytes.as_slice(), &[]);
+}
+
+#[test]
+fn test_bytes_default() {
+ let bytes = Bytes::default();
+ assert_eq!(bytes.len(), 0);
+ assert_eq!(bytes.as_slice(), &[]);
+}
+
+#[test]
+fn test_bytes_from_small() {
+ let data = b"hello";
+ let bytes = Bytes::from(data.as_slice());
+ assert_eq!(bytes.len(), 5);
+ assert_eq!(bytes.as_slice(), data);
+}
+
+#[test]
+fn test_bytes_from_large() {
+ let data = b"this is a very long string that exceeds the small bytes limit";
+ let bytes = Bytes::from(data.as_slice());
+ assert_eq!(
+ AnyView::from(&bytes).type_index(),
+ TypeIndex::kTVMFFIBytes as i32
+ );
+ assert_eq!(bytes.len(), data.len());
+ assert_eq!(bytes.as_slice(), data);
+}
+
+#[test]
+fn test_bytes_deref() {
+ let data = b"test";
+ let bytes = Bytes::from(data.as_slice());
+ assert_eq!(&*bytes, data);
+}
+
+#[test]
+fn test_bytes_clone_basic() {
+ let data = b"clone test";
+ let bytes1 = Bytes::from(data.as_slice());
+ let bytes2 = bytes1.clone();
+ assert_eq!(bytes1, bytes2);
+ assert_eq!(bytes1.as_slice(), bytes2.as_slice());
+}
+
+#[test]
+fn test_bytes_eq() {
+ let data = b"equal";
+ let bytes1 = Bytes::from(data.as_slice());
+ let bytes2 = Bytes::from(data.as_slice());
+ assert_eq!(bytes1, bytes2);
+
+ let bytes3 = Bytes::from(b"different");
+ assert_ne!(bytes1, bytes3);
+}
+
+#[test]
+fn test_bytes_ord() {
+ let bytes1 = Bytes::from(b"abc");
+ let bytes2 = Bytes::from(b"def");
+ let bytes3 = Bytes::from(b"abc");
+
+ assert!(bytes1 < bytes2);
+ assert!(bytes1 <= bytes2);
+ assert!(bytes2 > bytes1);
+ assert!(bytes2 >= bytes1);
+ assert_eq!(bytes1.cmp(&bytes3), std::cmp::Ordering::Equal);
+}
+
+#[test]
+fn test_bytes_hash() {
+ let data = b"hash test";
+ let bytes1 = Bytes::from(data.as_slice());
+ let bytes2 = Bytes::from(data.as_slice());
+
+ let mut hasher1 = DefaultHasher::new();
+ let mut hasher2 = DefaultHasher::new();
+ bytes1.hash(&mut hasher1);
+ bytes2.hash(&mut hasher2);
+
+ assert_eq!(hasher1.finish(), hasher2.finish());
+}
+
+// ============================================================================
+// String Tests
+// ============================================================================
+
+#[test]
+fn test_string_new() {
+ let s = String::new();
+ assert_eq!(s.len(), 0);
+ assert_eq!(s.as_str(), "");
+ assert_eq!(s.as_bytes(), &[]);
+}
+
+#[test]
+fn test_string_default() {
+ let s = String::default();
+ assert_eq!(s.len(), 0);
+ assert_eq!(s.as_str(), "");
+}
+
+#[test]
+fn test_string_from_small() {
+ let input = "hello";
+ let s = String::from(input);
+ assert_eq!(s.len(), 5);
+ assert_eq!(s.as_str(), input);
+ assert_eq!(s.as_bytes(), input.as_bytes());
+}
+
+#[test]
+fn test_string_from_large() {
+ let input = "this is a very long string that exceeds the small string limit";
+ let s = String::from(input);
+ assert_eq!(s.len(), input.len());
+ assert_eq!(s.as_str(), input);
+ assert_eq!(s.as_bytes(), input.as_bytes());
+}
+
+#[test]
+fn test_string_deref() {
+ let input = "deref test";
+ let s = String::from(input);
+ assert_eq!(&*s, input);
+}
+
+#[test]
+fn test_string_clone_basic() {
+ let input = "clone test";
+ let s1 = String::from(input);
+ let s2 = s1.clone();
+ assert_eq!(s1, s2);
+ assert_eq!(s1.as_str(), s2.as_str());
+}
+
+#[test]
+fn test_string_eq() {
+ let input = "equal";
+ let s1 = String::from(input);
+ let s2 = String::from(input);
+ assert_eq!(s1, s2);
+
+ let s3 = String::from("different");
+ assert_ne!(s1, s3);
+}
+
+#[test]
+fn test_string_ord() {
+ let s1 = String::from("abc");
+ let s2 = String::from("def");
+ let s3 = String::from("abc");
+
+ assert!(s1 < s2);
+ assert!(s1 <= s2);
+ assert!(s2 > s1);
+ assert!(s2 >= s1);
+ assert_eq!(s1.cmp(&s3), std::cmp::Ordering::Equal);
+}
+
+#[test]
+fn test_string_hash() {
+ let input = "hash test";
+ let s1 = String::from(input);
+ let s2 = String::from(input);
+
+ let mut hasher1 = DefaultHasher::new();
+ let mut hasher2 = DefaultHasher::new();
+ s1.hash(&mut hasher1);
+ s2.hash(&mut hasher2);
+
+ assert_eq!(hasher1.finish(), hasher2.finish());
+}
+
+// ============================================================================
+// Edge Cases
+// ============================================================================
+
+#[test]
+fn test_empty_strings() {
+ let empty_str = String::from("");
+ assert_eq!(empty_str.len(), 0);
+ assert_eq!(empty_str.as_str(), "");
+
+ let empty_bytes = Bytes::from(&[]);
+ assert_eq!(empty_bytes.len(), 0);
+ assert_eq!(empty_bytes.as_slice(), &[]);
+}
+
+#[test]
+fn test_unicode_strings() {
+ let unicode = "\u{1F680} Hello \u{4E16}\u{754C} \u{1F30D}";
+ let s = String::from(unicode);
+ assert_eq!(s.as_str(), unicode);
+ assert_eq!(s.len(), unicode.len());
+}
+
+#[test]
+fn test_small_vs_large_boundary() {
+ // Test exactly at the boundary (7 bytes)
+ let boundary = "1234567"; // 7 characters = 7 bytes
+ let s = String::from(boundary);
+ assert_eq!(s.as_str(), boundary);
+
+ // Test just over the boundary (8 bytes)
+ let over_boundary = "12345678"; // 8 characters = 8 bytes
+ let s2 = String::from(over_boundary);
+ assert_eq!(s2.as_str(), over_boundary);
+}
+
+#[test]
+fn test_bytes_small_vs_large_boundary() {
+ // Test exactly at the boundary (7 bytes)
+ let boundary = b"1234567"; // 7 bytes
+ let bytes = Bytes::from(boundary.as_slice());
+ assert_eq!(bytes.as_slice(), boundary);
+
+ // Test just over the boundary (8 bytes)
+ let over_boundary = b"12345678"; // 8 bytes
+ let bytes2 = Bytes::from(over_boundary.as_slice());
+ assert_eq!(bytes2.as_slice(), over_boundary);
+}
+
+#[test]
+fn test_mixed_comparisons() {
+ let s1 = String::from("test");
+ let s2 = String::from("test");
+ let s3 = String::from("different");
+
+ // Test all comparison operators
+ assert!(s1 == s2);
+ assert!(s1 != s3);
+ assert!(s1 <= s2);
+ assert!(s1 >= s2);
+ assert!(s1 > s3);
+ assert!(s3 < s1);
+}
+
+#[test]
+fn test_bytes_mixed_comparisons() {
+ let b1 = Bytes::from(b"test");
+ let b2 = Bytes::from(b"test");
+ let b3 = Bytes::from(b"different");
+
+ // Test all comparison operators
+ assert!(b1 == b2);
+ assert!(b1 != b3);
+ assert!(b1 <= b2);
+ assert!(b1 >= b2);
+ assert!(b1 > b3);
+ assert!(b3 < b1);
+}
+
+#[test]
+fn test_bytes_debug() {
+ let data = b"hello";
+ let bytes = Bytes::from(data.as_slice());
+ let debug_str = format!("{:?}", bytes);
+ assert!(debug_str.contains("Bytes"));
+}
+
+#[test]
+fn test_string_debug() {
+ let input = "world";
+ let s = String::from(input);
+ let debug_str = format!("{:?}", s);
+ assert!(debug_str.contains("String"));
+ assert!(debug_str.contains("world"));
+}
+
+#[test]
+fn test_bytes_clone_small() {
+ let data = b"small"; // 5 bytes - small case
+ let bytes1 = Bytes::from(data.as_slice());
+ let bytes2 = bytes1.clone();
+
+ assert_eq!(bytes1, bytes2);
+ assert_eq!(bytes1.as_slice(), bytes2.as_slice());
+ assert_eq!(bytes1.len(), bytes2.len());
+
+ // For small bytes, the underlying data should be different (copied)
+ // We can verify this by checking that the slices have different pointers
+ let slice1 = bytes1.as_slice();
+ let slice2 = bytes2.as_slice();
+ assert_ne!(slice1.as_ptr(), slice2.as_ptr());
+}
+
+#[test]
+fn test_bytes_clone_large() {
+ let data = b"this is a very long string that exceeds the small bytes limit"; // large case
+ let bytes1 = Bytes::from(data.as_slice());
+ let bytes2 = bytes1.clone();
+
+ assert_eq!(bytes1, bytes2);
+ assert_eq!(bytes1.as_slice(), bytes2.as_slice());
+ assert_eq!(bytes1.len(), bytes2.len());
+
+ // For large bytes, the underlying data should be the same (shared)
+ // We can verify this by checking that the slices have the same pointer
+ let slice1 = bytes1.as_slice();
+ let slice2 = bytes2.as_slice();
+ assert_eq!(slice1.as_ptr(), slice2.as_ptr());
+}
+
+#[test]
+fn test_string_clone_small() {
+ let input = "small"; // 5 bytes - small case
+ let s1 = String::from(input);
+ assert!(AnyView::from(&s1).debug_strong_count().is_none());
+ let s2 = s1.clone();
+
+ assert_eq!(s1, s2);
+ assert_eq!(s1.as_str(), s2.as_str());
+ assert_eq!(s1.len(), s2.len());
+
+ // For small strings, the underlying data should be different (copied)
+ // We can verify this by checking that the byte slices have different pointers
+ let bytes1 = s1.as_bytes();
+ let bytes2 = s2.as_bytes();
+ assert_ne!(bytes1.as_ptr(), bytes2.as_ptr());
+}
+
+#[test]
+fn test_string_clone_large() {
+ // large case
+ let input = "this is a very long string that exceeds the small string limit";
+ let s1 = String::from(input);
+ assert_eq!(AnyView::from(&s1).debug_strong_count().unwrap(), 1);
+ let s2 = s1.clone();
+ assert_eq!(AnyView::from(&s1).debug_strong_count().unwrap(), 2);
+ assert_eq!(AnyView::from(&s2).debug_strong_count().unwrap(), 2);
+
+ assert_eq!(s1, s2);
+ assert_eq!(s1.as_str(), s2.as_str());
+ assert_eq!(s1.len(), s2.len());
+
+ // For large strings, the underlying data should be the same (shared)
+ // We can verify this by checking that the byte slices have the same pointer
+ let bytes1 = s1.as_bytes();
+ let bytes2 = s2.as_bytes();
+ assert_eq!(bytes1.as_ptr(), bytes2.as_ptr());
+}
diff --git a/rust/tvm-ffi/tests/test_tensor.rs b/rust/tvm-ffi/tests/test_tensor.rs
new file mode 100644
index 0000000..d7e8f9b
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_tensor.rs
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+use tvm_ffi::*;
+
+// ============================================================================
+// Tensor Tests
+// ============================================================================
+
+#[test]
+fn test_tensor_basic() {
+ // Create a tensor using CPUNDAlloc
+ let shape = [2, 3, 4];
+ let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1);
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device);
+
+ // Test accessor methods
+ assert_eq!(tensor.shape(), &shape);
+ assert_eq!(tensor.ndim(), 3);
+ assert_eq!(tensor.dtype().code, DLDataTypeCode::kDLFloat as u8);
+ assert_eq!(tensor.dtype().bits, 32 as u8);
+ assert_eq!(tensor.device().device_type, DLDeviceType::kDLCPU);
+
+ // Test strides (should be calculated correctly for row-major layout)
+ let strides = tensor.strides();
+ assert_eq!(strides.len(), 3);
+ // For shape [2, 3, 4], strides should be [12, 4, 1] (row-major)
+ assert_eq!(strides[0], 12); // 3 * 4
+ assert_eq!(strides[1], 4); // 4
+ assert_eq!(strides[2], 1); // 1
+}
+
+#[test]
+fn test_tensor_data_as_slice_f32() {
+ // Create a tensor using CPUNDAlloc with f32 data type
+ let shape = [2, 3, 4];
+ let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1);
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device);
+
+ // Test data_as_slice for f32
+ let data_slice = tensor.data_as_slice::<f32>().unwrap();
+ assert_eq!(data_slice.len(), 24); // 2 * 3 * 4 = 24 elements
+
+ // Test data_as_slice_mut for f32
+ let data_slice_mut = tensor.data_as_slice_mut::<f32>().unwrap();
+ assert_eq!(data_slice_mut.len(), 24);
+
+ // Test that we can write to the tensor data
+ for i in 0..data_slice_mut.len() {
+ data_slice_mut[i] = i as f32;
+ }
+
+ // Test that we can read the written data
+ let data_slice_read = tensor.data_as_slice::<f32>().unwrap();
+ for i in 0..data_slice_read.len() {
+ assert_eq!(data_slice_read[i], i as f32);
+ }
+}
+
+#[test]
+fn test_tensor_data_as_slice_type_mismatch() {
+ // Create a tensor with f32 data type
+ let shape = [2, 3];
+ let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1);
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device);
+
+ // Test that trying to access as f64 (wrong type) fails
+ let result = tensor.data_as_slice::<f64>();
+ assert!(result.is_err());
+
+ // Test that trying to access as i32 (wrong type) fails
+ let result = tensor.data_as_slice::<i32>();
+ assert!(result.is_err());
+}
+
+#[test]
+fn test_any_tensor() {
+ let shape = [2, 3, 4];
+ let dtype = DLDataType::new(DLDataTypeCode::kDLFloat, 32, 1);
+ let device = DLDevice::new(DLDeviceType::kDLCPU, 0);
+ let tensor = Tensor::from_nd_alloc(CPUNDAlloc {}, &shape, dtype, device);
+ let any = Any::from(tensor.clone());
+ let any_view = AnyView::from(&tensor);
+
+ assert_eq!(any.type_index(), TypeIndex::kTVMFFITensor as i32);
+ let converted = Tensor::try_from(any).unwrap();
+ assert_eq!(converted.shape(), &shape);
+ assert_eq!(converted.dtype().code, DLDataTypeCode::kDLFloat as u8);
+
+ assert_eq!(any_view.type_index(), TypeIndex::kTVMFFITensor as i32);
+ let converted_view = Tensor::try_from(any_view).unwrap();
+ assert_eq!(converted_view.shape(), &shape);
+ assert_eq!(converted_view.dtype().code, DLDataTypeCode::kDLFloat as u8);
+}