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 ee81cc80 [Rust] Generated constructors take object references as `impl
IntoObjectRef<T>` (#956)
ee81cc80 is described below
commit ee81cc80f609c4fc83c5b69cd0120e462d987706
Author: Linzhang Li <[email protected]>
AuthorDate: Sat Nov 7 08:05:06 2026 -0500
[Rust] Generated constructors take object references as `impl
IntoObjectRef<T>` (#956)
## Summary
This PR makes the constructors that `tvm-ffi-stubgen` generates take an
object reference as
`impl IntoObjectRef<T>` instead of `impl Into<T>`. This covers the
wrapper's `new`, and
`from_fields` under `custom-new`. Values keep `impl Into<T>`, and
scalars stay their own type,
so the rule depends only on the type and is the same for generated and
hand-written APIs:
- An object reference is `impl IntoObjectRef<T>`.
- A value is `impl Into<T>`.
- A scalar is itself.
```rust
// Before
pub fn new(span: impl Into<Option<Span>>, ty: impl Into<PrimType>,
a: impl Into<PrimExpr>, b: impl Into<PrimExpr>) -> Self
// After
pub fn new(span: impl Into<Option<Span>>, ty: impl IntoObjectRef<PrimType>,
a: impl IntoObjectRef<PrimExpr>, b: impl
IntoObjectRef<PrimExpr>) -> Self
```
`IntoObjectRef<T>` is the conversion that knows about borrowing. `&x` is
a borrow that does not
change a reference count. `into_owned_object_ref()` moves an owned value
and clones only a
borrowed one. A stored argument therefore costs what it cost with
`Into`. Before this change,
a type with a generated `from_fields` and a hand-written `new` used two
traits for the same
arguments.
---
python/tvm_ffi/stub/rust_generator/codegen.py | 105 +++++++++++++++-------
python/tvm_ffi/stub/rust_generator/consts.py | 24 ++++-
rust/tvm-ffi-macros/src/rs_object.rs | 4 +-
rust/tvm-ffi/src/lib.rs | 6 +-
tests/python/test_stubgen_rust.py | 123 +++++++++++++++++++++++---
5 files changed, 217 insertions(+), 45 deletions(-)
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
index 9f5917a8..3d3b8b83 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -24,9 +24,10 @@ the verdict of :mod:`tvm_ffi.stub.layout`:
- *complete*: the struct mirrors every physical field at its real offset and
width, pinned by a ``const`` size/alignment assertion. ``<Leaf>Obj::new``
and the wrapper's ``new`` take every field root to leaf: the former each as
- its own type, the latter a scalar as its own type and anything else as
- ``impl Into``. ``custom-new`` leaves the wrapper's ``new`` to hand-written
- code and names the generated one ``from_fields``. A type that is not final
+ its own type, the latter a scalar as its own type, an object reference as
+ ``impl IntoObjectRef`` and anything else as ``impl Into``. ``custom-new``
+ leaves the wrapper's ``new`` to hand-written code and names the generated
+ one ``from_fields``. A type that is not final
can be the parent of a type defined in Rust (``tvm_ffi::rs_object``): its
``<Leaf>Obj::new`` is public, and its object implements ``ShallowCopy`` and
``UnsafeInit`` when every ancestor's is generated too (a ``ty-map``'d one is
@@ -59,7 +60,7 @@ module, a ``crate::`` path anywhere else.
from __future__ import annotations
import dataclasses
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple
from .. import consts as C
from ..layout import Verdict, classify
@@ -77,6 +78,8 @@ if TYPE_CHECKING:
from collections.abc import Iterable, Mapping
from pathlib import Path
+ from tvm_ffi.core import TypeSchema
+
from ..file_utils import CodeBlock
from ..utils import InitConfig, NamedTypeSchema, ObjectInfo, Options
from .directives import EnumSpec
@@ -103,14 +106,52 @@ def _call_lines(open_: str, items: list[str], close: str)
-> list[str]:
]
-def _param_type(rust_type: str) -> str:
- """Type a parameter of the wrapper's allocator: a scalar as itself,
anything else as ``impl Into``."""
- return rust_type if rust_type in C_RUST.RUST_EXACT_PARAMS else f"impl
Into<{rust_type}>"
+class _Param(NamedTuple):
+ """A parameter of an allocator: one physical field, root to leaf."""
+
+ name: str
+ #: The type ``<Leaf>Obj::new`` takes: the field's mirror, or a ``field``
directive's narrowing.
+ rust_type: str
+ #: Whether the reflected field is an object reference
(:func:`_is_object_ref`), not an enum.
+ of_object: bool
+
+ @property
+ def is_object(self) -> bool:
+ """Whether ``rust_type`` is an object reference; a ``nullable`` one is
an ``Option``."""
+ return self.of_object and not self.rust_type.startswith("Option<")
-def _forward(ident: str, rust_type: str) -> str:
- """Pass a parameter of the wrapper's allocator to ``<Leaf>Obj::new``,
which takes ``rust_type``."""
- return ident if rust_type in C_RUST.RUST_EXACT_PARAMS else
f"{ident}.into()"
+def _is_object_ref(schema: TypeSchema) -> bool:
+ """Whether the Rust type of ``schema`` is an object reference.
+
+ A builtin origin is one when listed in ``RUST_OBJECT_REF_ORIGINS``; a type
key
+ outside ``ffi.*`` always is, as a generated or ``ty-map``'d binding.
+ """
+ origin = schema.origin
+ if origin in C_RUST.RUST_OBJECT_REF_ORIGINS:
+ return True
+ return "." in origin and not origin.startswith("ctypes.") and not
is_crate_type_key(origin)
+
+
+def _param_type(param: _Param, imports: RustImports) -> str:
+ """Type a parameter of the wrapper's allocator.
+
+ A scalar as itself, an object reference as ``impl IntoObjectRef``
(recording its ``use``,
+ which also brings ``into_owned_object_ref`` into scope), anything else as
``impl Into``.
+ """
+ if param.rust_type in C_RUST.RUST_EXACT_PARAMS:
+ return param.rust_type
+ if param.is_object:
+ return f"impl
{imports.record('tvm_ffi::IntoObjectRef')}<{param.rust_type}>"
+ return f"impl Into<{param.rust_type}>"
+
+
+def _forward(param: _Param) -> str:
+ """Pass a parameter of the wrapper's allocator to ``<Leaf>Obj::new``,
which takes its mirror."""
+ ident = rust_ident(param.name)
+ if param.rust_type in C_RUST.RUST_EXACT_PARAMS:
+ return ident
+ return f"{ident}.into_owned_object_ref()" if param.is_object else
f"{ident}.into()"
def _check_width(target: str, field: NamedTypeSchema, rust_type: str, width:
int) -> None:
@@ -522,13 +563,13 @@ class _ObjectRenderer:
def _allocator_params(
self, key: str, info: ObjectInfo, narrowed: frozenset[str] =
frozenset()
- ) -> list[tuple[str, str]]:
- """``(field, type)`` of every physical field root to leaf, as
``<key>Obj::new`` takes.
+ ) -> list[_Param]:
+ """Every physical field root to leaf, as ``<key>Obj::new`` takes it.
``narrowed`` names the fields a subtype of ``key`` narrows with a
``field`` directive.
"""
parent = info.parent_type_key
- inherited: list[tuple[str, str]] = []
+ inherited: list[_Param] = []
if parent is not None and self._provider(parent) != "crate":
inherited = self._allocator_params(
parent, object_info_from_type_key(parent), narrowed |
self._narrowed_names(key)
@@ -539,9 +580,9 @@ class _ObjectRenderer:
self,
key: str,
info: ObjectInfo,
- inherited: list[tuple[str, str]],
+ inherited: list[_Param],
narrowed: frozenset[str] = frozenset(),
- ) -> list[tuple[str, str]]:
+ ) -> list[_Param]:
"""Extend the parent's parameters with ``key``'s own fields by offset.
A ``field`` directive on an inherited field narrows that parameter; the
@@ -553,13 +594,15 @@ class _ObjectRenderer:
return RustImports() if name in narrowed else self.imports
params = [
- (name, self._narrowed(key, name, rust_type, imports(name)))
- for name, rust_type in inherited
+ p._replace(rust_type=self._narrowed(key, p.name, p.rust_type,
imports(p.name)))
+ for p in inherited
]
+ enums = self.imports.directives.enums
for field in sorted(info.fields, key=lambda f: f.offset or 0):
mirror = self._field_mirror(key, field, imports(field.name))
assert mirror is not None # complete: every field along the chain
has a mirror
- params.append((field.name, mirror))
+ of_object = f"{key}.{field.name}" not in enums and
_is_object_ref(field)
+ params.append(_Param(field.name, mirror, of_object))
return params
def _narrowed_names(self, key: str) -> frozenset[str]:
@@ -577,7 +620,7 @@ class _ObjectRenderer:
def _fn_lines(
self,
head: str,
- params: list[tuple[str, str]],
+ params: list[_Param],
call: tuple[str, list[str]],
init: list[str],
*,
@@ -585,15 +628,16 @@ class _ObjectRenderer:
) -> list[str]:
"""Render ``<head>(<params>) -> Self { let <call>; Self { <init> } }``
in an ``impl`` block.
- ``generic`` takes a parameter that is not a scalar as ``impl Into``
(:func:`_param_type`).
+ ``generic`` takes a parameter that is not a scalar as ``impl
IntoObjectRef`` or
+ ``impl Into`` (:func:`_param_type`).
"""
plist = [
- f"{rust_ident(name)}: {_param_type(rust_type) if generic else
rust_type}"
- for name, rust_type in params
+ f"{rust_ident(p.name)}: {_param_type(p, self.imports) if generic
else p.rust_type}"
+ for p in params
]
binding, args = call
return [
- *[f" {line}" for line in _allow_non_snake_case(rust_ident(n)
for n, _ in params)],
+ *[f" {line}" for line in
_allow_non_snake_case(rust_ident(p.name) for p in params)],
*_call_lines(f" {head}(", plist, ") -> Self {"),
*_call_lines(f" let {binding}(", args, ");"),
*_call_lines(" Self { ", init, " }"),
@@ -605,13 +649,14 @@ class _ObjectRenderer:
``<Leaf>Obj::new`` takes every field as its own type, so the parent's,
generated
or hand-written, takes what it hands on. The wrapper's ``new`` takes a
scalar
- field as its own type and any other as ``impl Into``
(:func:`_param_type`), as
- ``#[rs_object]``'s ``new`` does, and converts it on the way to
``<Leaf>Obj::new``.
+ field as its own type, as ``#[rs_object]``'s ``new`` does, an object
reference as
+ ``impl IntoObjectRef`` and any other as ``impl Into``
(:func:`_param_type`), and
+ converts it on the way to ``<Leaf>Obj::new``.
``custom-new`` names the wrapper's allocator ``from_fields`` instead.
``<Leaf>Obj::new`` of a type that is not final is public: a subtype
defined in Rust
builds its parent's object with it.
"""
- inherited: list[tuple[str, str]] = []
+ inherited: list[_Param] = []
if has_parent:
parent = self.info.parent_type_key
assert parent is not None
@@ -620,11 +665,11 @@ class _ObjectRenderer:
)
params = self._level_params(self.type_key, self.info, inherited)
to_parent = [
- f"{rust_ident(name)}.into()" if rust_type != parent_type else
rust_ident(name)
- for (name, rust_type), (_, parent_type) in zip(params, inherited)
+ f"{rust_ident(p.name)}.into()" if p.rust_type != q.rust_type else
rust_ident(p.name)
+ for p, q in zip(params, inherited)
]
- own = [rust_ident(name) for name, _ in params[len(inherited) :]]
- forward = [_forward(rust_ident(name), rust_type) for name, rust_type
in params]
+ own = [rust_ident(p.name) for p in params[len(inherited) :]]
+ forward = [_forward(p) for p in params]
allow = self._allow_missing_docs()
sections = [
[
diff --git a/python/tvm_ffi/stub/rust_generator/consts.py
b/python/tvm_ffi/stub/rust_generator/consts.py
index 99631218..fd0b9fab 100644
--- a/python/tvm_ffi/stub/rust_generator/consts.py
+++ b/python/tvm_ffi/stub/rust_generator/consts.py
@@ -96,12 +96,32 @@ RUST_SCALAR_WIDTHS = {
}
#: Parameters of the wrapper's allocator taken as their own type, as
``#[rs_object]``'s
-#: ``new`` takes them: a literal infers only from its exact type. Any other
parameter is
-#: ``impl Into<T>``; ``<Leaf>Obj::new`` takes every one as its own type.
+#: ``new`` takes them: a literal infers only from its exact type. An object
reference is
+#: ``impl IntoObjectRef<T>`` and any other parameter ``impl Into<T>``;
``<Leaf>Obj::new``
+#: takes every one as its own type.
RUST_EXACT_PARAMS = frozenset(
{"bool", "i8", "i16", "i32", "i64", "isize", "u8", "u16", "u32", "u64",
"usize", "f32", "f64"}
)
+#: Builtin origins whose Rust type is an object reference, taken as ``impl
IntoObjectRef<T>``.
+#: A type key outside ``ffi.*`` is one too: a generated or ``ty-map``'d
binding. ``Shape``,
+#: like ``Array`` / ``List`` / ``Map``, is a container taken as ``impl
Into<T>``, which also
+#: accepts what it is built from (``From<impl AsRef<[i64]>>``).
+RUST_OBJECT_REF_ORIGINS = frozenset(
+ {
+ "Object",
+ "Tensor",
+ "Callable",
+ "ffi.Object",
+ "ffi.Tensor",
+ "ffi.Function",
+ "ffi.Module",
+ "ffi.Error",
+ "ffi.reflection.AccessPath",
+ "ffi.reflection.AccessStep",
+ }
+)
+
#: Size of an object reference field; ``nullable`` may only wrap those.
RUST_POINTER_SIZE = 8
diff --git a/rust/tvm-ffi-macros/src/rs_object.rs
b/rust/tvm-ffi-macros/src/rs_object.rs
index 506442de..da3e0ddd 100644
--- a/rust/tvm-ffi-macros/src/rs_object.rs
+++ b/rust/tvm-ffi-macros/src/rs_object.rs
@@ -156,7 +156,9 @@ fn expand(args: &RsObjectArgs, item: &DeriveInput) ->
syn::Result<TokenStream2>
let base_id = &base.ident;
let base_ty = &base.ty;
// `new` takes a scalar as its own type, so that a literal infers as it,
and any
- // other field as `impl Into`.
+ // other field as `impl Into`. Tokens alone do not tell an object
reference from a
+ // value, so this object-level allocator does not take `impl
IntoObjectRef` as a
+ // reference's constructor does.
let params = fields.iter().map(|field| {
let (name, ty) = (&field.ident, &field.ty);
if is_exact_param(ty) {
diff --git a/rust/tvm-ffi/src/lib.rs b/rust/tvm-ffi/src/lib.rs
index 2e027dd8..423bb96a 100644
--- a/rust/tvm-ffi/src/lib.rs
+++ b/rust/tvm-ffi/src/lib.rs
@@ -244,7 +244,11 @@ pub use tvm_ffi_macros::match_any;
/// fields of both.
/// * `MyPairObj::new(base_, fields...)`, which takes the fields in order: a
/// scalar (`bool`, an integer or a float) as its own type, so that a literal
-/// infers as it, and any other field from anything that converts into it.
+/// infers as it, and any other field as `impl Into<T>`. This is the
+/// object-level allocator, which sees only tokens and so cannot tell an
+/// object reference from a value; a constructor on the reference takes an
+/// object reference as [`impl IntoObjectRef<T>`](IntoObjectRef), as the
+/// ones `tvm-ffi-stubgen` generates do.
///
/// # Arguments
///
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index 14d48e5a..7c56f1d8 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -748,8 +748,8 @@ impl ExprObj {
impl Expr {
/// Lossless complete-field allocation.
- pub fn new(span: impl Into<Option<Span>>, ty: impl Into<Type>) -> Self {
- let obj = ExprObj::new(span.into(), ty.into());
+ pub fn new(span: impl Into<Option<Span>>, ty: impl IntoObjectRef<Type>) ->
Self {
+ let obj = ExprObj::new(span.into(), ty.into_owned_object_ref());
Self { data_: ObjectArc::new(obj) }
}
}
@@ -817,11 +817,16 @@ impl Add {
/// Lossless complete-field allocation.
pub fn new(
span: impl Into<Option<Span>>,
- ty: impl Into<PrimType>,
- a: impl Into<PrimExpr>,
- b: impl Into<PrimExpr>,
+ ty: impl IntoObjectRef<PrimType>,
+ a: impl IntoObjectRef<PrimExpr>,
+ b: impl IntoObjectRef<PrimExpr>,
) -> Self {
- let obj = AddObj::new(span.into(), ty.into(), a.into(), b.into());
+ let obj = AddObj::new(
+ span.into(),
+ ty.into_owned_object_ref(),
+ a.into_owned_object_ref(),
+ b.into_owned_object_ref(),
+ );
Self { data_: ObjectArc::new(obj) }
}
}
@@ -838,6 +843,7 @@ def test_render_complete_expr_golden() -> None:
assert _uses(imports) == {
"std::ptr::addr_of_mut",
"tvm_ffi::Object",
+ "tvm_ffi::IntoObjectRef",
"tvm_ffi::ObjectArc",
"tvm_ffi::reflection::ShallowCopy",
"tvm_ffi::reflection::UnsafeInit",
@@ -876,6 +882,7 @@ def test_render_complete_add_golden() -> None:
assert text == ADD_EXPECTED
assert _uses(imports) == {
"std::ops::Deref",
+ "tvm_ffi::IntoObjectRef",
"tvm_ffi::ObjectArc",
"crate::ir::ExprObj",
"crate::ir::Expr",
@@ -1216,18 +1223,112 @@ def test_custom_new_renames_the_wrapper_allocator() ->
None:
assert (
"impl Add {\n /// Lossless complete-field allocation.\n"
" pub fn from_fields(\n"
- " span: impl Into<Span>,\n"
- " ty: impl Into<Type>,\n"
- " a: impl Into<Expr>,\n"
- " b: impl Into<Expr>,\n"
+ " span: impl IntoObjectRef<Span>,\n"
+ " ty: impl IntoObjectRef<Type>,\n"
+ " a: impl IntoObjectRef<Expr>,\n"
+ " b: impl IntoObjectRef<Expr>,\n"
" ) -> Self {\n"
- " let obj = AddObj::new(span.into(), ty.into(), a.into(),
b.into());\n"
+ " let obj = AddObj::new(\n"
+ " span.into_owned_object_ref(),\n"
+ " ty.into_owned_object_ref(),\n"
+ " a.into_owned_object_ref(),\n"
+ " b.into_owned_object_ref(),\n"
+ " );\n"
" Self { data_: ObjectArc::new(obj) }\n"
" }\n}"
) in text
assert " pub fn new(" not in text
+def test_wrapper_allocator_takes_object_references_as_into_object_ref() ->
None:
+ """An object reference is `impl IntoObjectRef`, a value `impl Into`, a
scalar itself.
+
+ A `field` directive narrowing an object reference keeps it one; `nullable`
makes it an
+ `Option`, a value. `Shape` is a container, taken as `impl Into` like
`Array`.
+ """
+ optional_var = TypeSchema("Optional", (TypeSchema("demo.Var"),))
+ info = _info(
+ "demo.Node",
+ (
+ _field("tensor", "Tensor", 24, 8),
+ _field("owner", "Object", 32, 8),
+ _field("callback", "Callable", 40, 8),
+ _field("module", "ffi.Module", 48, 8),
+ _field("var", "demo.Var", 56, 8),
+ _field("buffer", "demo.Var", 64, 8),
+ _field("span", "demo.Span", 72, 8),
+ _field("name", "str", 80, 16),
+ _field("items", "Array", 96, 8),
+ _field("attrs", "Map", 104, 8),
+ _field("shape", "Shape", 112, 8),
+ _field("value", "Any", 120, 16),
+ _field("kind", "int", 136, 8),
+ _field("parent", optional_var, 144, 8),
+ _field("dtype", "dtype", 152, 4, 2),
+ _field("count", "int", 156, 4),
+ ),
+ total_size=160,
+ is_final=True,
+ )
+ imports = RUST.new_imports()
+ RUST.add_directive(imports, "field", "demo.Node.buffer ->
crate::typed::BufferVar", 1)
+ RUST.add_directive(imports, "nullable", "demo.Node.span", 2)
+ RUST.add_directive(imports, "enum", "demo.Node.kind -> Kind(i64) { A=0 }",
3)
+ text, imports = _render(info, imports)
+ assert (
+ " pub fn new(\n"
+ " tensor: impl IntoObjectRef<Tensor>,\n"
+ " owner: impl IntoObjectRef<ObjectRef>,\n"
+ " callback: impl IntoObjectRef<Function>,\n"
+ " module: impl IntoObjectRef<Module>,\n"
+ " var: impl IntoObjectRef<Var>,\n"
+ " buffer: impl IntoObjectRef<BufferVar>,\n"
+ " span: impl Into<Option<Span>>,\n"
+ " name: impl Into<String>,\n"
+ " items: impl Into<Array<Any>>,\n"
+ " attrs: impl Into<Map<Any, Any>>,\n"
+ " shape: impl Into<Shape>,\n"
+ " value: impl Into<Any>,\n"
+ " kind: impl Into<Kind>,\n"
+ " parent: impl Into<Option<Var>>,\n"
+ " dtype: impl Into<DLDataType>,\n"
+ " count: i32,\n"
+ " ) -> Self {\n"
+ " let obj = NodeObj::new(\n"
+ " tensor.into_owned_object_ref(),\n"
+ " owner.into_owned_object_ref(),\n"
+ " callback.into_owned_object_ref(),\n"
+ " module.into_owned_object_ref(),\n"
+ " var.into_owned_object_ref(),\n"
+ " buffer.into_owned_object_ref(),\n"
+ " span.into(),\n"
+ " name.into(),\n"
+ " items.into(),\n"
+ " attrs.into(),\n"
+ " shape.into(),\n"
+ " value.into(),\n"
+ " kind.into(),\n"
+ " parent.into(),\n"
+ " dtype.into(),\n"
+ " count,\n"
+ " );\n"
+ ) in text, text
+ assert "tvm_ffi::IntoObjectRef" in _uses(imports)
+
+
+def
test_wrapper_allocator_without_object_references_does_not_import_into_object_ref()
-> None:
+ info = _info(
+ "demo.Pair",
+ (_field("name", "str", 24, 16), _field("count", "int", 40, 8)),
+ total_size=48,
+ is_final=True,
+ )
+ text, imports = _render(info)
+ assert " pub fn new(name: impl Into<String>, count: i64) -> Self {" in
text
+ assert "IntoObjectRef" not in text
+ assert "tvm_ffi::IntoObjectRef" not in _uses(imports)
+
+
def test_allocator_body_wraps_like_rustfmt() -> None:
"""Past 100 columns a `let` first moves its call to the next line; the
rest go one per line."""
