Seven-Streams opened a new pull request, #609:
URL: https://github.com/apache/tvm-ffi/pull/609
This PR is based on #608.
## Summary
This PR adds Rust support to the stub generator.
Rust bindings can now be generated through both CMake and the CLI:
* In CMake, set `STUB_TARGET=rust`.
* In the CLI, use `tvm-ffi-stubgen <generated_dir> --target rust ...`.
This PR also adds documentation, examples, and tests for Rust stub
generation.
## Key Changes
* Added a Rust backend in `python/tvm_ffi/stubgen/rust_generator`.
* Added unit tests and string-level tests for Rust stub generation in
`test_stubgen.py`.
* Added documentation in `docs/packaging/stubgen.rst`.
* Added Rust stubgen examples in `examples/rust_stubgen`.
* Extended the CMake integration to support Rust stub generation.
* Updated the Rust package to support generated bindings.
## Example
Given the following C++ definition:
```c++
/*! \brief Data object: a pair of 64-bit integers `a` and `b`. */
class IntPairObj : public ffi::Object {
public:
int64_t a;
int64_t b;
IntPairObj() = default;
explicit IntPairObj(int64_t a, int64_t b) : a(a), b(b) {}
/*! \brief Sum of the two components. */
int64_t sum() const { return a + b; }
// All fields are writable, so the generated Rust wrapper gets `DerefMut`.
static constexpr bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rust_stubgen.IntPair", IntPairObj,
ffi::Object);
};
/*! \brief Reference wrapper for `IntPairObj`. */
class IntPair : public ffi::ObjectRef {
public:
explicit IntPair(int64_t a, int64_t b) { data_ =
ffi::make_object<IntPairObj>(a, b); }
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IntPair, ffi::ObjectRef,
IntPairObj);
};
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IntPairObj>()
.def(refl::init<int64_t, int64_t>())
.def_rw("a", &IntPairObj::a, "first component")
.def_rw("b", &IntPairObj::b, "second component")
.def("sum", &IntPairObj::sum, "a + b");
// Lets an `AnyView` holding this object convert back into the `IntPair`
ref,
// which is what the generated Rust bindings rely on for object returns.
refl::TypeAttrDef<IntPairObj>().def(refl::type_attr::kConvert,
&refl::details::FFIConvertFromAnyViewToObjectRef<IntPair>);
}
```
The Rust stub generator produces Rust wrappers for the reflected objects and
methods. For example, the generated bindings include:
```rust
// tvm-ffi-stubgen(begin): helpers
fn lookup_type_index(type_key: &'static str) -> i32 {
static CACHE: std::sync::OnceLock<
std::sync::Mutex<std::collections::HashMap<&'static str, i32>>,
> = std::sync::OnceLock::new();
let cache = CACHE.get_or_init(||
std::sync::Mutex::new(std::collections::HashMap::new()));
if let Some(v) = cache.lock().unwrap().get(type_key) {
return *v;
}
let arg = unsafe {
tvm_ffi::tvm_ffi_sys::TVMFFIByteArray::from_str(type_key) };
let mut tindex = 0;
let ret = unsafe { tvm_ffi::tvm_ffi_sys::TVMFFITypeKeyToIndex(&arg, &mut
tindex) };
assert_eq!(ret, 0, "type key `{type_key}` is not registered");
cache.lock().unwrap().insert(type_key, tindex);
tindex
}
fn get_type_method(
type_key: &'static str,
method_name: &str,
) -> tvm_ffi::Result<tvm_ffi::Function> {
let type_index = lookup_type_index(type_key);
unsafe {
let info = tvm_ffi::tvm_ffi_sys::TVMFFIGetTypeInfo(type_index);
if info.is_null() {
return Err(tvm_ffi::Error::new(
tvm_ffi::TYPE_ERROR,
&format!("no type info for `{type_key}`"),
"",
));
}
let info = &*info;
for i in 0..info.num_methods {
let mi = &*info.methods.add(i as usize);
if mi.name.as_str() == method_name {
if !<tvm_ffi::Function as
tvm_ffi::type_traits::AnyCompatible>::check_any_strict(
&mi.method,
) {
return Err(tvm_ffi::Error::new(
tvm_ffi::TYPE_ERROR,
&format!("method `{method_name}` on `{type_key}` is
not a Function"),
"",
));
}
return Ok(<tvm_ffi::Function as
tvm_ffi::type_traits::AnyCompatible>::copy_from_any_view_after_check(&mi.method));
}
}
}
Err(tvm_ffi::Error::new(
tvm_ffi::TYPE_ERROR,
&format!("method `{method_name}` not found on `{type_key}`"),
"",
))
}
// tvm-ffi-stubgen(end)
// tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair
#[repr(C)]
pub struct IntPairObj {
base: Object,
pub a: i64,
pub b: i64,
}
unsafe impl ObjectCore for IntPairObj {
const TYPE_KEY: &'static str = "rust_stubgen.IntPair";
fn type_index() -> i32 {
lookup_type_index(Self::TYPE_KEY)
}
unsafe fn object_header_mut(this: &mut Self) -> &mut TVMFFIObject {
Object::object_header_mut(&mut this.base)
}
}
#[repr(C)]
#[derive(DeriveObjectRef, Clone)]
pub struct IntPair {
data: ObjectArc<IntPairObj>,
}
impl Deref for IntPair {
type Target = IntPairObj;
fn deref(&self) -> &IntPairObj {
&self.data
}
}
impl DerefMut for IntPair {
fn deref_mut(&mut self) -> &mut IntPairObj {
&mut self.data
}
}
impl IntPair {
pub fn new(_0: i64, _1: i64) -> Result<Self> {
let ctor = get_type_method(IntPairObj::TYPE_KEY, "__ffi_init__")?;
Ok(ctor.call_packed(&[AnyView::from(&_0),
AnyView::from(&_1)])?.try_into()?)
}
pub fn sum(&mut self) -> Result<i64> {
let f = get_type_method(IntPairObj::TYPE_KEY, "sum")?;
Ok(f.call_packed(&[AnyView::from(&*self)])?.try_into()?)
}
}
```
The generated helper functions are responsible for runtime type lookup and
method dispatch, allowing the generated object wrappers to invoke reflected
methods through the TVM FFI runtime.
## Testing
* Added unit tests and string-level tests in `test_stubgen.py`.
* End-to-end tests are attached in the PR comments.
[e2e_test.tar.gz](https://github.com/user-attachments/files/28654124/e2e_test.tar.gz)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
