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
};
cursor
bot
force-pushed
the
rust-stubgen
branch
from
4a7b28e to
addd67d
Compare
The stubgen test was expecting functions.rs at src/functions.rs, but the generator actually writes it to src/_tvm_ffi_stubgen_detail/functions.rs. This commit updates the test to check the correct path. Also includes formatting fixes from cargo fmt. Co-authored-by: Cursor <cursoragent@cursor.com>
… subtyping 实现 repr(C) 零开销 Rust FFI wrapper 生成，使用 Deref trait 建模继承， 生成 100% safe 的用户代码。 ## 核心变更 ### tvm-ffi 库增强 - **subtyping.rs (新文件)**: 子类型转换辅助函数 - `upcast<From, To>`: 安全的消费型向上转换 - `try_downcast<From, To>`: 安全的消费型向下转换（使用 runtime type check） - `is_instance_of`: 继承关系运行时检查 - **macros.rs**: 新增 `impl_object_hierarchy!` 宏 - 自动生成 Deref/From/TryFrom impls 用于继承链 - 零开销引用 upcast（通过 Deref auto-coercion） - 安全的消费型转换（使用标准库 trait） - **object_wrapper.rs**: 移除 `is_instance_type`（移至 subtyping.rs） ### tvm-ffi-macros 修复 - **derive(ObjectRef)**: 修复为委托给 ObjectRef 的 AnyCompatible 实现， 消除对 `pub(crate) unsafe_::inc_ref` 的直接调用 - **derive(Object)**: 将 `proc_macro_error::abort!` 替换为 `panic!`， 避免在外部 crate 编译时失败 ### tvm-ffi-stubgen 重构 - **repr_c.rs (新文件)**: C 兼容性检查逻辑 - `check_repr_c`: 验证类型是否可用 repr(C) 表示 - 根据 field size 映射精确的 Rust POD 类型（i8/i16/i32/i64, f32/f64） - 改进 alignment padding 验证 - **generate.rs**: 全新 repr(C) 代码生成路径 - 生成 `#[repr(C)] *Obj` struct（零开销字段访问） - 生成 `#[repr(C)] *Ref` wrapper（user-facing 类型） - 调用 `impl_object_hierarchy!` 自动生成子类型转换 - 只为直接字段生成 getter（继承字段通过 Deref 自动可用） - 过滤内建类型（ffi.*）避免重复定义 - 对 repr(C) 类型使用 `self as &ObjectRef`（deref coercion） - 对 fallback 类型使用 `self.as_object_ref()` - **model.rs**: 扩展 TypeGen - 新增 `ancestor_chain` 字段存储继承链 - 修改 `call_expr` 使用 `Into<ObjectRef>` 而非 `as_object_ref().clone()` ## 生成代码特性 ✅ **100% Safe Rust**: 无任何 `unsafe` 关键字出现在生成的 stub 代码中 ✅ **零开销字段访问**: repr(C) 类型直接访问字段无 FFI 调用 ✅ **类型安全子类型转换**: 使用标准库 Deref/From/TryFrom trait ✅ **符合 Rust 惯用法**: 自动 upcast、fallible downcast ## 测试验证 使用 tvm_ffi_testing.so 生成 stub crate，编译通过，无任何 unsafe 代码。 Closes: stubgen repr(C) redesign initiative Co-authored-by: Cursor <cursoragent@cursor.com>
- impl_object_hierarchy!: 补充 TryFrom<DirectParent> for Self，覆盖根类型 ObjectRef 下转 - subtyping: upcast/try_downcast 改为 #[doc(hidden)] pub，供宏展开使用，用户统一走 From/TryFrom - check_repr_c: 仅允许对齐 padding (field.offset == align_up(pos, alignment))，并补齐 Any 映射以支持 Map/Array<Any> - define_object_wrapper!: 增加 From<Wrapper> for ObjectRef，保证 .into() 兼容 - generate: is_builtin_type 增加 None；render_facade_module 跳过空模块，避免空 pub mod ffi Co-authored-by: Cursor <cursoragent@cursor.com>
check_repr_c previously used ancestor[0] as parent, which for multi-level inheritance returns the root type (ffi.Object). This caused parent_total_size=24 while direct field offset=56, falsely rejected as layout gap. Use ancestor[type_depth-1] for direct parent so TestObjectDerived passes repr(C). Co-authored-by: Cursor <cursoragent@cursor.com>
Move the Rust stubgen usage guide into docs/packaging so it is grouped with packaging workflows, and keep the crate README focused on generated interface/design details. Update the stubgen integration test to validate the documented usage flow with raw-string test source and stable APIs. Co-authored-by: Cursor <cursoragent@cursor.com>
Switch generated object method wrappers from global registry lookup to type reflection method resolution, and map constructor wrappers to Rust-style `new` while removing legacy `c_ffi_init`. Extend integration tests and docs to validate typed constructors and Cxx inheritance roundtrip behavior. Co-authored-by: Cursor <cursoragent@cursor.com>
Drop unused helper and stale repr-c struct fields to keep the Rust workspace warning-free after the method lookup and constructor naming refactor. Co-authored-by: Cursor <cursoragent@cursor.com>
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.