Add Rust stub generator. #434
Conversation
Summary of ChangesHello @Blealtan, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Rust ecosystem within TVM by introducing a dedicated stub generator. This new tool automates the creation of Rust FFI bindings, allowing developers to seamlessly integrate TVM's C++ functions and types into Rust projects. By generating a compilable Rust crate that wraps TVM-FFI DLLs, it streamlines the process of calling TVM functions from Rust, improving developer experience and reducing manual FFI boilerplate. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This is a great addition, introducing a Rust stub generator which will improve the development experience for Rust users of tvm-ffi. The implementation is comprehensive, covering CLI parsing, FFI interaction, code generation, and robust integration testing. The code is well-structured and follows good Rust practices. I have a couple of suggestions for improvement, mainly around reducing code duplication by leveraging existing types from tvm-ffi and simplifying some of the string manipulation logic.
rust/tvm-ffi-stubgen/src/ffi.rs
Outdated
| #[repr(C)] | ||
| #[derive(Debug)] | ||
| pub(crate) struct Array { | ||
| handle: TVMFFIObjectHandle, | ||
| } | ||
|
|
||
| extern "C" { | ||
| fn TVMFFIObjectIncRef(handle: TVMFFIObjectHandle) -> i32; | ||
| fn TVMFFIObjectDecRef(handle: TVMFFIObjectHandle) -> i32; | ||
| } | ||
|
|
||
| impl Clone for Array { | ||
| fn clone(&self) -> Self { | ||
| unsafe { | ||
| TVMFFIObjectIncRef(self.handle); | ||
| } | ||
| Self { handle: self.handle } | ||
| } | ||
| } | ||
|
|
||
| impl Drop for Array { | ||
| fn drop(&mut self) { | ||
| unsafe { | ||
| TVMFFIObjectDecRef(self.handle); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| unsafe impl tvm_ffi::type_traits::AnyCompatible for Array { | ||
| fn type_str() -> String { | ||
| "ffi.Array".to_string() | ||
| } | ||
|
|
||
| unsafe fn copy_to_any_view(src: &Self, data: &mut TVMFFIAny) { | ||
| data.type_index = TVMFFITypeIndex::kTVMFFIArray as i32; | ||
| data.small_str_len = 0; | ||
| data.data_union.v_obj = src.handle as *mut tvm_ffi::tvm_ffi_sys::TVMFFIObject; | ||
| } | ||
|
|
||
| unsafe fn move_to_any(src: Self, data: &mut TVMFFIAny) { | ||
| data.type_index = TVMFFITypeIndex::kTVMFFIArray as i32; | ||
| data.small_str_len = 0; | ||
| data.data_union.v_obj = src.handle as *mut tvm_ffi::tvm_ffi_sys::TVMFFIObject; | ||
| } | ||
|
|
||
| unsafe fn check_any_strict(data: &TVMFFIAny) -> bool { | ||
| data.type_index == TVMFFITypeIndex::kTVMFFIArray as i32 | ||
| } | ||
|
|
||
| unsafe fn copy_from_any_view_after_check(data: &TVMFFIAny) -> Self { | ||
| let handle = data.data_union.v_obj as TVMFFIObjectHandle; | ||
| TVMFFIObjectIncRef(handle); | ||
| Self { handle } | ||
| } | ||
|
|
||
| unsafe fn move_from_any_after_check(data: &mut TVMFFIAny) -> Self { | ||
| let handle = data.data_union.v_obj as TVMFFIObjectHandle; | ||
| Self { handle } | ||
| } | ||
|
|
||
| unsafe fn try_cast_from_any_view(data: &TVMFFIAny) -> Result<Self, ()> { | ||
| if data.type_index == TVMFFITypeIndex::kTVMFFIArray as i32 { | ||
| Ok(Self::copy_from_any_view_after_check(data)) | ||
| } else { | ||
| Err(()) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl ArgIntoRef for Array { | ||
| type Target = Array; | ||
| fn to_ref(&self) -> &Self::Target { | ||
| self | ||
| } | ||
| } | ||
|
|
||
| impl<'a> ArgIntoRef for &'a Array { | ||
| type Target = Array; | ||
| fn to_ref(&self) -> &Self::Target { | ||
| self | ||
| } | ||
| } | ||
|
|
||
| impl IntoArgHolder for Array { | ||
| type Target = Array; | ||
| fn into_arg_holder(self) -> Self::Target { | ||
| self | ||
| } | ||
| } | ||
|
|
||
| impl<'a> IntoArgHolder for &'a Array { | ||
| type Target = &'a Array; | ||
| fn into_arg_holder(self) -> Self::Target { | ||
| self | ||
| } | ||
| } | ||
|
|
||
| impl TryFrom<Any> for Array { | ||
| type Error = Error; | ||
| fn try_from(value: Any) -> Result<Self, Self::Error> { | ||
| if let Some(ret) = value.try_as::<Array>() { | ||
| Ok(ret) | ||
| } else { | ||
| Err(Error::new(TYPE_ERROR, "Expected ffi.Array", "")) | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
The local Array struct and its implementations (Clone, Drop, AnyCompatible, TryFrom<Any>, etc.) appear to duplicate functionality that should already be present in the tvm-ffi crate's tvm_ffi::Array. Using the existing tvm_ffi::Array would significantly reduce code duplication, simplify this module, and improve maintainability.
The tvm-ffi crate is already a dependency, so you should be able to use tvm_ffi::Array directly. This would also allow you to remove the helper functions array_size and array_get_item, as tvm_ffi::Array likely provides an iterator or similar methods for element access.
For example, list_registered_type_keys could be simplified to something like:
pub(crate) fn list_registered_type_keys() -> FfiResult<Vec<String>> {
let get_keys = Function::get_global("ffi.GetRegisteredTypeKeys")?;
let keys_any = get_keys.call_tuple_with_len::<0, _>(())?;
let keys: tvm_ffi::Array<tvm_ffi::String> = keys_any.try_into()?;
Ok(keys.iter().map(|s| s.as_str().to_string()).collect())
}
rust/tvm-ffi-stubgen/src/generate.rs
Outdated
| let mut remainder = full_name; | ||
| if !prefix.is_empty() && remainder.starts_with(prefix) { | ||
| remainder = &remainder[prefix.len()..]; | ||
| } else if !prefix.is_empty() && remainder.starts_with(prefix.trim_end_matches('.')) { | ||
| remainder = &remainder[prefix.trim_end_matches('.').len()..]; | ||
| remainder = remainder.trim_start_matches('.'); | ||
| } |
There was a problem hiding this comment.
The logic for stripping the prefix in split_name is a bit complex. Since the prefix is normalized to always end with a . (if not empty) in lib.rs, the else if branch seems redundant. You can simplify this logic using strip_prefix, which also makes the intent clearer.
let remainder = if !prefix.is_empty() {
full_name.strip_prefix(prefix).unwrap_or(full_name)
} else {
full_name
};
2 participants
This pull request introduces a new Rust crate,
tvm-ffi-stubgen, to the workspace. The crate is a command-line tool for generating Rust stubs fromtvm-ffimetadata. The implementation includes argument parsing, FFI helpers, metadata schema parsing, code generation logic, and the main entry point. The workspace is updated to include this new crate.Notably, this PR can now automatically generate Rust binding for TVM and TileLang, enabling Rust code to manually assemble a legal TileLang AST.
Relation with existing Python stubgen
This Rust stubgen only supports a Rust equivalent to the "
STUB_INITis on" mode in the existing Python stubgen, which yields a compilable Rust crate that wraps a TVM-FFI DLL.Addition of the
tvm-ffi-stubgencratetvm-ffi-stubgento the workspace members inCargo.toml, ensuring it is built as part of the Rust workspace.tvm-ffi-stubgen/Cargo.tomlwith dependencies, package metadata, and bin target for the new stub generator.Core implementation of the stub generator
cli.rsusingclap, supporting options for output directory, DLLs, prefix, and other configuration.ffi.rsfor loading DLLs, querying global function/type names, and extracting type information from the loaded libraries.model.rs, supporting code generation from metadata.schema.rsto interpret type schemas from JSON metadata.lib.rs, orchestrating argument parsing, FFI calls, filtering, module construction, and writing generated files.main.rsto parse arguments and invoke the main logic.Integration test
It tests the functionality through loading, compiling and running a crate that invokes
add_onefromlibtvm_ffi_testing.so. The testing is yet to be fully automated; it now relies on human to supply the correct paths.Non-stubgen changes
High-level rationale
The non-stubgen changes make the Rust FFI runtime capable of:
File-by-file details
Expand...
rust/tvm-ffi-macros/src/object_macros.rsrust/tvm-ffi/src/any.rsrust/tvm-ffi/src/type_traits.rsrust/tvm-ffi/src/object_wrapper.rsrust/tvm-ffi/src/macros.rsrust/tvm-ffi/src/function_internal.rsrust/tvm-ffi/src/collections/map.rsrust/tvm-ffi/src/collections/mod.rsrust/tvm-ffi/src/error.rsrust/tvm-ffi/src/lib.rsrust/Cargo.tomlThis PR is vibed with GPT-5.2-Codex through Copilot and then reviewed by human.