This is an automated email from the ASF dual-hosted git repository.
tqchen 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 41e9b7ee [FEAT][Rust] Read function metadata and doc strings from a
Module (#833)
41e9b7ee is described below
commit 41e9b7eec7a424e382bd260385836b200f9de01f
Author: Gabriel Wu <[email protected]>
AuthorDate: Tue Oct 6 22:10:46 2026 +0800
[FEAT][Rust] Read function metadata and doc strings from a Module (#833)
## Why
C++ `ModuleObj` and Python `Module` can tell whether a module implements
a function. They can also read the metadata and the doc string that a
library exports for it (`__tvm_ffi__metadata_<name>` and
`__tvm_ffi__doc_<name>`), and optionally look in the modules it imports.
The Rust `Module` can only load a library and get a function, so a Rust
host cannot check a library's exports or read their type schemas.
## What
- `Module::implements_function(name, query_imports) -> Result<bool>`
- `Module::get_function_metadata(name, query_imports) ->
Result<Option<String>>` returns the JSON string the library wrote, such
as `{"type_schema":"..."}`.
- `Module::get_function_doc(name, query_imports) ->
Result<Option<String>>`
They call `ffi.ModuleImplementsFunction`,
`ffi.ModuleGetFunctionMetadata` and `ffi.ModuleGetFunctionDoc`, the same
functions that Python calls. As in the C++ methods, `query_imports`
controls the search of imported modules. `None` means that the library
exports no metadata or doc string for the function. The metadata stays a
JSON string, since the crate has no JSON dependency.
The Rust guide documents the accessors.
---
docs/guides/rust_lang_guide.md | 15 +++++
rust/tvm-ffi/src/extra/module.rs | 67 +++++++++++++++++++
rust/tvm-ffi/tests/test_module.rs | 135 ++++++++++++++++++++++++++++++++++++++
3 files changed, 217 insertions(+)
diff --git a/docs/guides/rust_lang_guide.md b/docs/guides/rust_lang_guide.md
index 9c869121..02171630 100644
--- a/docs/guides/rust_lang_guide.md
+++ b/docs/guides/rust_lang_guide.md
@@ -90,6 +90,21 @@ fn main() -> Result<()> {
}
```
+A module also reports whether it implements a function, and reads the
+metadata and doc string a function was exported with, as
+`Module.implements_function`, `Module.get_function_metadata` and
+`Module.get_function_doc` do in Python. The second argument says whether to
+also look in the modules it imports:
+
+```rust
+let module = Module::load_from_file("build/add_one_cpu.so")?;
+if module.implements_function("add_one_cpu", false)? {
+ // A JSON object such as {"type_schema": "..."}, if it was exported
+ let metadata = module.get_function_metadata("add_one_cpu", false)?;
+ let doc = module.get_function_doc("add_one_cpu", false)?;
+}
+```
+
### Working with Tensors
Create and manipulate tensors:
diff --git a/rust/tvm-ffi/src/extra/module.rs b/rust/tvm-ffi/src/extra/module.rs
index 84cc4780..f4296747 100644
--- a/rust/tvm-ffi/src/extra/module.rs
+++ b/rust/tvm-ffi/src/extra/module.rs
@@ -20,6 +20,7 @@ use crate::derive::{Object, ObjectRef};
use crate::error::Result;
use crate::function::Function;
use crate::object::{Object, ObjectArc};
+use crate::string::String;
use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
//-----------------------------------------------------
@@ -71,4 +72,70 @@ impl Module {
.call_tuple_with_len::<3, _>((self, name, true))?
.try_into()
}
+
+ /// Check whether the module implements a function.
+ ///
+ /// # Arguments
+ /// * `name` - The name of the function
+ /// * `query_imports` - Whether to also check the modules this module
imports
+ ///
+ /// # Returns
+ /// * `Result<bool>` - Whether the function exists
+ pub fn implements_function<Str: AsRef<str>>(
+ &self,
+ name: Str,
+ query_imports: bool,
+ ) -> Result<bool> {
+ let name = String::from(name);
+ crate::cached_global_func!("ffi.ModuleImplementsFunction")
+ .call_tuple_with_len::<3, _>((self, name, query_imports))?
+ .try_into()
+ }
+
+ /// Get the metadata of a function the module exports.
+ ///
+ /// The metadata is a JSON object whose `type_schema` entry is the
+ /// function's type schema, as C++ `TVM_FFI_DLL_EXPORT_TYPED_FUNC` exports
+ /// it with `TVM_FFI_DLL_EXPORT_INCLUDE_METADATA`.
+ ///
+ /// # Arguments
+ /// * `name` - The name of the function
+ /// * `query_imports` - Whether to also check the modules this module
imports
+ ///
+ /// # Returns
+ /// * `Result<Option<String>>` - The metadata as a JSON string, or `None`
+ /// if the function has no metadata
+ pub fn get_function_metadata<Str: AsRef<str>>(
+ &self,
+ name: Str,
+ query_imports: bool,
+ ) -> Result<Option<String>> {
+ let name = String::from(name);
+ crate::cached_global_func!("ffi.ModuleGetFunctionMetadata")
+ .call_tuple_with_len::<3, _>((self, name, query_imports))?
+ .try_into()
+ }
+
+ /// Get the doc string of a function the module exports.
+ ///
+ /// The doc string is the one C++ `TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`
+ /// exports.
+ ///
+ /// # Arguments
+ /// * `name` - The name of the function
+ /// * `query_imports` - Whether to also check the modules this module
imports
+ ///
+ /// # Returns
+ /// * `Result<Option<String>>` - The doc string, or `None` if the function
+ /// has none
+ pub fn get_function_doc<Str: AsRef<str>>(
+ &self,
+ name: Str,
+ query_imports: bool,
+ ) -> Result<Option<String>> {
+ let name = String::from(name);
+ crate::cached_global_func!("ffi.ModuleGetFunctionDoc")
+ .call_tuple_with_len::<3, _>((self, name, query_imports))?
+ .try_into()
+ }
}
diff --git a/rust/tvm-ffi/tests/test_module.rs
b/rust/tvm-ffi/tests/test_module.rs
new file mode 100644
index 00000000..9cf44962
--- /dev/null
+++ b/rust/tvm-ffi/tests/test_module.rs
@@ -0,0 +1,135 @@
+/*
+ * 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.
+ */
+//! Module accessors read the functions, metadata and doc strings that
+//! libtvm_ffi_testing exports with `TVM_FFI_DLL_EXPORT_TYPED_FUNC` and
+//! `TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC`.
+
+use std::env::consts::{DLL_PREFIX, DLL_SUFFIX};
+
+use tvm_ffi::tvm_ffi_sys::{TVMFFIAny, TVMFFISafeCallType};
+use tvm_ffi::*;
+
+// The metadata and doc getters libtvm_ffi_testing exports, read directly.
+extern "C" {
+ fn __tvm_ffi__metadata_testing_dll_schema_id_int(
+ handle: *mut std::ffi::c_void,
+ args: *const TVMFFIAny,
+ num_args: i32,
+ result: *mut TVMFFIAny,
+ ) -> i32;
+ fn __tvm_ffi__doc_testing_dll_test_add_with_docstring(
+ handle: *mut std::ffi::c_void,
+ args: *const TVMFFIAny,
+ num_args: i32,
+ result: *mut TVMFFIAny,
+ ) -> i32;
+}
+
+const SCHEMA_ID_INT: &str = "testing_dll_schema_id_int";
+const ADD_WITH_DOCSTRING: &str = "testing_dll_test_add_with_docstring";
+const MISSING: &str = "testing_dll_missing";
+
+/// libtvm_ffi_testing, which the test binary links, so that the loader finds
+/// it by name.
+fn testing_module() -> Module {
+
Module::load_from_file(format!("{DLL_PREFIX}tvm_ffi_testing{DLL_SUFFIX}")).unwrap()
+}
+
+/// A module without functions of its own that imports libtvm_ffi_testing.
+fn importing_module() -> Module {
+ let module: Module = Function::get_global("ffi.SystemLib")
+ .unwrap()
+ .call_tuple_with_len::<1,
_>((String::from("testing.rust.test_module."),))
+ .unwrap()
+ .try_into()
+ .unwrap();
+ Function::get_global("ffi.ModuleImportModule")
+ .unwrap()
+ .call_tuple_with_len::<2, _>((&module, &testing_module()))
+ .unwrap();
+ module
+}
+
+/// Calls an exported metadata or doc getter.
+fn call_getter(getter: TVMFFISafeCallType) -> String {
+ // SAFETY: the getters take no arguments and do not use the handle.
+ let getter = unsafe { Function::from_extern_c(std::ptr::null_mut(),
getter, None) };
+ getter.call_packed(&[]).unwrap().try_into().unwrap()
+}
+
+#[test]
+fn test_module_implements_function() {
+ let module = testing_module();
+ assert!(module.implements_function(SCHEMA_ID_INT, false).unwrap());
+ assert!(!module.implements_function(MISSING, false).unwrap());
+}
+
+#[test]
+fn test_module_get_function_metadata() {
+ let module = testing_module();
+ let metadata = module
+ .get_function_metadata(SCHEMA_ID_INT, false)
+ .unwrap()
+ .unwrap();
+ let exported = call_getter(__tvm_ffi__metadata_testing_dll_schema_id_int);
+ assert_eq!(metadata.as_str(), exported.as_str());
+ assert!(module
+ .get_function_metadata(MISSING, false)
+ .unwrap()
+ .is_none());
+}
+
+#[test]
+fn test_module_get_function_doc() {
+ let module = testing_module();
+ let doc = module
+ .get_function_doc(ADD_WITH_DOCSTRING, false)
+ .unwrap()
+ .unwrap();
+ let exported =
call_getter(__tvm_ffi__doc_testing_dll_test_add_with_docstring);
+ assert_eq!(doc.as_str(), exported.as_str());
+ // A function exported without a doc string has none.
+ assert!(module
+ .get_function_doc(SCHEMA_ID_INT, false)
+ .unwrap()
+ .is_none());
+}
+
+#[test]
+fn test_module_accessors_query_imports() {
+ let module = importing_module();
+ assert!(!module.implements_function(SCHEMA_ID_INT, false).unwrap());
+ assert!(module.implements_function(SCHEMA_ID_INT, true).unwrap());
+ assert!(module
+ .get_function_metadata(SCHEMA_ID_INT, false)
+ .unwrap()
+ .is_none());
+ assert!(module
+ .get_function_metadata(SCHEMA_ID_INT, true)
+ .unwrap()
+ .is_some());
+ assert!(module
+ .get_function_doc(ADD_WITH_DOCSTRING, false)
+ .unwrap()
+ .is_none());
+ assert!(module
+ .get_function_doc(ADD_WITH_DOCSTRING, true)
+ .unwrap()
+ .is_some());
+}
