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 17ea4d7b [STUBGEN] Carry native layout facts through the
language-neutral data. (#720)
17ea4d7b is described below
commit 17ea4d7ba2d4153aa991f1c581cecd2a06f17e26
Author: Linzhang Li <[email protected]>
AuthorDate: Wed Sep 2 10:32:11 2026 -0400
[STUBGEN] Carry native layout facts through the language-neutral data.
(#720)
`NamedTypeSchema` now records the byte facts the C++ registry already
publishes for every reflected field (size/alignment/offset, static
default or factory value, structural-eq flag, frozen-ness), and
`ObjectInfo` records the ancestor chain plus the type's own `sizeof`
(`None` for a type without its own metadata entry, whose inherited size
says nothing about the type). This gives non-Python backends the
memory-layout inputs they need without a second reflection pass.
Language-neutral infrastructure for such backends:
- `//` marker syntax for `.rs` files (the block parser was already
syntax-driven);
- directory scans visit only the active generator's `source_exts`, so
one tree can hold stub files for several languages without one target
rewriting another's files;
- `--target` choices follow the generator registry instead of a
hard-coded list;
- drop the unused `generate_helpers_block` protocol method.
No generator other than `python` is registered yet.
---
python/tvm_ffi/stub/cli.py | 11 +-
python/tvm_ffi/stub/consts.py | 13 ++-
python/tvm_ffi/stub/file_utils.py | 14 ++-
python/tvm_ffi/stub/generator.py | 10 ++
python/tvm_ffi/stub/python_generator/generator.py | 1 +
python/tvm_ffi/stub/utils.py | 124 +++++++++++++++-----
tests/python/test_stubgen.py | 133 +++++++++++++++++++++-
7 files changed, 259 insertions(+), 47 deletions(-)
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index 66688de9..2b054c72 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -28,7 +28,7 @@ from typing import TYPE_CHECKING
from . import consts as C
from .file_utils import FileInfo, collect_files, syntax_for
-from .generator import get_generator
+from .generator import generator_names, get_generator
from .lib_state import (
collect_global_funcs,
collect_type_keys,
@@ -53,7 +53,7 @@ def __main__() -> int:
for imp in opt.imports or []:
importlib.import_module(imp)
dlls = [ctypes.CDLL(lib) for lib in opt.dlls]
- files: list[FileInfo] = collect_files([Path(f) for f in opt.files])
+ files: list[FileInfo] = collect_files([Path(f) for f in opt.files],
generator.source_exts)
global_funcs: dict[str, list[FuncInfo]] = collect_global_funcs()
init_path: Path | None = None
if opt.files:
@@ -347,15 +347,16 @@ def _parse_args() -> Options:
metavar="PATH",
help=(
"Files or directories to process. Directories are scanned
recursively; "
- "only .py and .pyi files are modified. Use tvm-ffi-stubgen
directives to "
- "select where stubs are generated."
+ "only files with the target's source extensions (.py and .pyi for
python) "
+ "are modified. Use tvm-ffi-stubgen directives to select where
stubs are "
+ "generated."
),
)
parser.add_argument(
"--target",
type=str,
default="python",
- choices=["python"],
+ choices=generator_names(),
help="Code generator target.",
)
parser.add_argument(
diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py
index 46ba9709..53115386 100644
--- a/python/tvm_ffi/stub/consts.py
+++ b/python/tvm_ffi/stub/consts.py
@@ -29,9 +29,9 @@ class MarkerSyntax:
"""Comment-syntax-specific stub directive markers.
All stub directives are embedded inside single-line comments. The comment
- token (currently ``#`` for Python sources) parameterizes the marker set,
- while the directive grammar (``tvm-ffi-stubgen(begin): ...`` etc.) stays
- uniform.
+ token (``#`` for Python sources, ``//`` for Rust) parameterizes the marker
+ set, while the directive grammar (``tvm-ffi-stubgen(begin): ...`` etc.)
+ stays uniform.
"""
comment: str
@@ -69,12 +69,15 @@ class MarkerSyntax:
PYTHON_SYNTAX = MarkerSyntax(comment="#")
+RUST_SYNTAX = MarkerSyntax(comment="//")
#: Map a source-file extension to the marker syntax used inside it. The block
-#: parser selects the syntax per file.
+#: parser selects the syntax per file. Which extensions a run actually visits
is
+#: decided by the active generator (``Generator.source_exts``), not by this
map.
SYNTAX_BY_EXT: dict[str, MarkerSyntax] = {
".py": PYTHON_SYNTAX,
".pyi": PYTHON_SYNTAX,
+ ".rs": RUST_SYNTAX,
}
STUB_BLOCK_KINDS: TypeAlias = Literal[
@@ -100,8 +103,6 @@ TERM_CYAN = "\033[36m"
TERM_WHITE = "\033[37m"
DOC_URL = "https://tvm.apache.org/ffi/packaging/stubgen.html";
-DEFAULT_SOURCE_EXTS = set(SYNTAX_BY_EXT)
-
# Language-neutral metadata transform applied while building `ObjectInfo` from
# the FFI reflection registry (see `utils.ObjectInfo.from_type_info`).
FN_NAME_MAP: dict[str, str] = {}
diff --git a/python/tvm_ffi/stub/file_utils.py
b/python/tvm_ffi/stub/file_utils.py
index d6c75f56..94d0769c 100644
--- a/python/tvm_ffi/stub/file_utils.py
+++ b/python/tvm_ffi/stub/file_utils.py
@@ -22,7 +22,7 @@ import dataclasses
import difflib
import os
import traceback
-from collections.abc import Generator, Iterable
+from collections.abc import Collection, Generator, Iterable
from pathlib import Path
from typing import Callable
@@ -239,8 +239,14 @@ class FileInfo:
self.code_blocks = source.code_blocks
-def collect_files(paths: list[Path]) -> list[FileInfo]:
- """Collect all files from the given paths and parse them into FileInfo
objects."""
+def collect_files(paths: list[Path], source_exts: Collection[str]) ->
list[FileInfo]:
+ """Collect all files from the given paths and parse them into FileInfo
objects.
+
+ A path given explicitly is always visited; a directory is walked
recursively
+ and only files whose (lower-cased) extension is in ``source_exts`` -- the
+ extensions the active generator owns -- are considered.
+ """
+ exts = {ext.lower() for ext in source_exts}
def _on_error(e: Exception) -> None:
print(
@@ -257,7 +263,7 @@ def collect_files(paths: list[Path]) -> list[FileInfo]:
for root, _dirs, files in path_walk(p, follow_symlinks=False,
on_error=_on_error):
for file in files:
f = Path(root) / file
- if f.suffix.lower() not in C.DEFAULT_SOURCE_EXTS:
+ if f.suffix.lower() not in exts:
continue
yield f
diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py
index 1651dd97..99c3aa0d 100644
--- a/python/tvm_ffi/stub/generator.py
+++ b/python/tvm_ffi/stub/generator.py
@@ -66,6 +66,11 @@ class Generator(Protocol):
#: Short identifier, e.g. ``"python"``.
name: str
+ #: Source-file extensions this generator owns (lower-cased, with the dot).
+ #: Directory scans only visit these, so one tree can hold stub files for
+ #: several languages without one target rewriting another's files.
+ source_exts: frozenset[str]
+
#: Comment-marker syntax for the files this generator emits.
syntax: C.MarkerSyntax
@@ -182,6 +187,11 @@ _GENERATORS: dict[str, Generator] = {
}
+def generator_names() -> list[str]:
+ """Return the registered target names, sorted (the ``--target``
choices)."""
+ return sorted(_GENERATORS)
+
+
def get_generator(target: str) -> Generator:
"""Resolve a generator by target name."""
if target not in _GENERATORS:
diff --git a/python/tvm_ffi/stub/python_generator/generator.py
b/python/tvm_ffi/stub/python_generator/generator.py
index da1fc16f..d05b2f2a 100644
--- a/python/tvm_ffi/stub/python_generator/generator.py
+++ b/python/tvm_ffi/stub/python_generator/generator.py
@@ -43,6 +43,7 @@ class PythonGenerator:
name = "python"
syntax = C.PYTHON_SYNTAX
+ source_exts = frozenset({".py", ".pyi"})
def default_ty_map(self) -> dict[str, str]:
"""Return the default FFI-origin -> Python-type name map."""
diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py
index 07f40f8d..c16eb657 100644
--- a/python/tvm_ffi/stub/utils.py
+++ b/python/tvm_ffi/stub/utils.py
@@ -25,12 +25,15 @@ generator (e.g.
:mod:`tvm_ffi.stub.python_generator.codegen`).
from __future__ import annotations
import dataclasses
-from typing import Any
+from typing import TYPE_CHECKING, Any
-from tvm_ffi.core import TypeInfo, TypeSchema, _lookup_type_attr
+from tvm_ffi.core import MISSING, TypeInfo, TypeSchema, _lookup_type_attr
from . import consts as C
+if TYPE_CHECKING:
+ from tvm_ffi.core import TypeField
+
def _parse_type_schema(raw: str | dict[str, Any]) -> TypeSchema:
"""Parse a type schema from either a JSON string or an already-parsed
dict."""
@@ -86,14 +89,72 @@ class Options:
@dataclasses.dataclass(init=False)
class NamedTypeSchema(TypeSchema):
- """A type schema with an associated name."""
+ """A type schema with an associated name.
+
+ For a reflected object field, the schema also carries the facts the C++
+ registry recorded about the native field, so a generator can reason about
+ memory layout and defaults without a second reflection pass:
+
+ - ``size`` / ``alignment`` / ``offset``: the byte facts of the native
field.
+ ``offset`` is measured from the start of the ``TVMFFIObject`` header, the
+ same base the C ABI field getters use. All three are ``None`` for
function
+ parameters and synthetic schemas.
+ - ``default``: the registered static default value (:data:`MISSING` when
+ none). ``default_is_factory`` marks a ``default_factory`` registration,
+ whose value only exists by calling the factory through the FFI.
+ - ``structural_eq``: the decoded structural-equality flag
+ (``"ignore"``, ``"def-recursive"``, ``"def-non-recursive"`` or ``None``).
+ - ``frozen``: ``True`` for read-only (``def_ro``) fields.
+ """
name: str
-
- def __init__(self, name: str, schema: TypeSchema) -> None:
- """Initialize a `NamedTypeSchema` with the given name and schema."""
+ size: int | None = None
+ alignment: int | None = None
+ offset: int | None = None
+ default: Any = MISSING
+ default_is_factory: bool = False
+ structural_eq: str | None = None
+ frozen: bool = False
+
+ def __init__(
+ self,
+ name: str,
+ schema: TypeSchema,
+ *,
+ size: int | None = None,
+ alignment: int | None = None,
+ offset: int | None = None,
+ default: Any = MISSING,
+ default_is_factory: bool = False,
+ structural_eq: str | None = None,
+ frozen: bool = False,
+ ) -> None:
+ """Initialize a `NamedTypeSchema` with the given name, schema and
field facts."""
super().__init__(origin=schema.origin, args=schema.args)
self.name = name
+ self.size = size
+ self.alignment = alignment
+ self.offset = offset
+ self.default = default
+ self.default_is_factory = default_is_factory
+ self.structural_eq = structural_eq
+ self.frozen = frozen
+
+ @staticmethod
+ def from_type_field(field: TypeField) -> NamedTypeSchema:
+ """Construct a `NamedTypeSchema` from a reflected
:class:`~tvm_ffi.core.TypeField`."""
+ is_factory = field.c_default_factory is not MISSING
+ return NamedTypeSchema(
+ name=field.name,
+ schema=_parse_type_schema(field.metadata["type_schema"]),
+ size=field.size,
+ alignment=field.alignment,
+ offset=field.offset,
+ default=field.c_default_factory if is_factory else field.c_default,
+ default_is_factory=is_factory,
+ structural_eq=field.c_structural_eq,
+ frozen=field.frozen,
+ )
@dataclasses.dataclass
@@ -121,7 +182,15 @@ class InitFieldInfo:
@dataclasses.dataclass
class ObjectInfo:
- """Information of an object type, including its fields and methods."""
+ """Information of an object type, including its fields and methods.
+
+ ``fields`` lists only the fields declared by this type; inherited fields
are
+ reached through ``parent_type_key`` / ``ancestors``. ``total_size`` is the
+ native ``sizeof`` of the object (header included) when the type registered
+ its own metadata, and ``None`` otherwise: a type without its own
+ ``ObjectDef`` inherits its parent's metadata entry, whose size says nothing
+ about the type itself.
+ """
fields: list[NamedTypeSchema]
methods: list[FuncInfo]
@@ -129,6 +198,10 @@ class ObjectInfo:
parent_type_key: str | None = None
init_fields: list[InitFieldInfo] = dataclasses.field(default_factory=list)
has_init: bool = False
+ ancestors: list[str] = dataclasses.field(default_factory=list)
+ """Type keys of every ancestor, root first (``["ffi.Object", "ir.Expr"]``
for ``tirx.Add``)."""
+ total_size: int | None = None
+ """Native ``sizeof`` in bytes, or ``None`` when the type has no metadata
of its own."""
def has_overloaded_methods(self) -> bool:
"""Return whether reflection exposed multiple signatures for a
method."""
@@ -143,47 +216,38 @@ class ObjectInfo:
@staticmethod
def from_type_info(type_info: TypeInfo) -> ObjectInfo:
"""Construct an `ObjectInfo` from a `TypeInfo` instance."""
- parent_type_key: str | None = None
- if type_info.parent_type_info is not None:
- parent_type_key = type_info.parent_type_info.type_key
+ # Ancestor chain, root first (`ancestor_infos[-1]` is the direct
parent).
+ ancestor_infos: list[TypeInfo] = []
+ ancestor_info: TypeInfo | None = type_info.parent_type_info
+ while ancestor_info is not None:
+ ancestor_infos.append(ancestor_info)
+ ancestor_info = ancestor_info.parent_type_info
+ ancestor_infos.reverse()
+ parent_type_key = ancestor_infos[-1].type_key if ancestor_infos else
None
# Detect __ffi_init__ from TypeMethod or TypeAttrColumn.
has_init = any(m.name == "__ffi_init__" for m in type_info.methods)
if not has_init:
has_init = _lookup_type_attr(type_info.type_index, "__ffi_init__")
is not None
- # Walk parent chain (parent-first) to collect all init-eligible fields.
+ # Collect init-eligible fields from the whole chain, inherited fields
first.
init_fields: list[InitFieldInfo] = []
if has_init:
- ti: TypeInfo | None = type_info
- chain: list[TypeInfo] = []
- while ti is not None:
- chain.append(ti)
- ti = ti.parent_type_info
- for ancestor_info in reversed(chain):
- for field in ancestor_info.fields:
+ for declaring_info in [*ancestor_infos, type_info]:
+ for field in declaring_info.fields:
if not field.c_init:
continue
init_fields.append(
InitFieldInfo(
name=field.name,
- schema=NamedTypeSchema(
- name=field.name,
-
schema=_parse_type_schema(field.metadata["type_schema"]),
- ),
+ schema=NamedTypeSchema.from_type_field(field),
kw_only=field.c_kw_only,
has_default=field.c_has_default,
)
)
return ObjectInfo(
- fields=[
- NamedTypeSchema(
- name=field.name,
- schema=_parse_type_schema(field.metadata["type_schema"]),
- )
- for field in type_info.fields
- ],
+ fields=[NamedTypeSchema.from_type_field(field) for field in
type_info.fields],
methods=[
FuncInfo(
schema=NamedTypeSchema(
@@ -198,4 +262,6 @@ class ObjectInfo:
parent_type_key=parent_type_key,
init_fields=init_fields,
has_init=has_init,
+ ancestors=[info.type_key for info in ancestor_infos],
+ total_size=type_info.total_size if type_info._has_type_metadata
else None,
)
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index f94d2418..4fefdc24 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -23,12 +23,12 @@ from pathlib import Path
import pytest
import tvm_ffi.stub.cli as stub_cli
from tvm_ffi import Object, method
-from tvm_ffi.core import TypeSchema
+from tvm_ffi.core import MISSING, TypeSchema,
_lookup_or_register_type_info_from_type_key
from tvm_ffi.dataclasses import py_class
from tvm_ffi.stub import consts as C
from tvm_ffi.stub.cli import _stage_2, _stage_3
-from tvm_ffi.stub.file_utils import CodeBlock, FileInfo
-from tvm_ffi.stub.generator import get_generator
+from tvm_ffi.stub.file_utils import CodeBlock, FileInfo, collect_files,
syntax_for
+from tvm_ffi.stub.generator import generator_names, get_generator
from tvm_ffi.stub.python_generator import consts as PC
from tvm_ffi.stub.python_generator.codegen import (
generate_python_all,
@@ -53,6 +53,12 @@ from tvm_ffi.stub.utils import (
ObjectInfo,
Options,
)
+from tvm_ffi.testing import (
+ TestObjectBase,
+ _TestCxxClassBase,
+ _TestCxxClassDerived,
+ _TestCxxClassDerivedDerived,
+)
_counter = itertools.count()
@@ -964,3 +970,124 @@ def test_stage_2_filters_prefix_and_marks_root(
sub_text = sub_api.read_text(encoding="utf-8")
assert 'LIB = _FFI_LOAD_LIB("demo-pkg", "demo_shared")' in root_text
assert "LIB =" not in sub_text
+
+
+def test_objectinfo_from_type_info_layout_facts() -> None:
+ """Reflected size/alignment/offset/default facts reach ``ObjectInfo``
unchanged."""
+ base = ObjectInfo.from_type_info(_TestCxxClassBase.__tvm_ffi_type_info__)
# ty: ignore[unresolved-attribute]
+ derived =
ObjectInfo.from_type_info(_TestCxxClassDerived.__tvm_ffi_type_info__) # ty:
ignore[unresolved-attribute]
+ dd =
ObjectInfo.from_type_info(_TestCxxClassDerivedDerived.__tvm_ffi_type_info__) #
ty: ignore[unresolved-attribute]
+
+ assert base.ancestors == ["ffi.Object"]
+ assert derived.ancestors == ["ffi.Object", "testing.TestCxxClassBase"]
+ assert dd.ancestors == [
+ "ffi.Object",
+ "testing.TestCxxClassBase",
+ "testing.TestCxxClassDerived",
+ ]
+ assert dd.parent_type_key == dd.ancestors[-1]
+
+ header = 24 # sizeof(TVMFFIObject): the fields of a root type start right
after it
+ v_i64, v_i32 = base.fields
+ assert (v_i64.size, v_i64.alignment, v_i64.offset) == (8, 8, header)
+ assert (v_i32.size, v_i32.alignment, v_i32.offset) == (4, 4, header + 8)
+ assert base.total_size == 40 # 36 bytes of data, padded to the 8-byte
alignment
+ v_f64, v_f32 = derived.fields
+ assert v_f64.offset == base.total_size # a derived type's fields follow
the parent's size
+ assert (v_f32.size, v_f32.alignment, v_f32.offset) == (4, 4, 48)
+ assert derived.total_size == 56
+ v_str, v_bool = dd.fields
+ assert (v_str.size, v_str.alignment, v_str.offset) == (16, 8,
derived.total_size)
+ assert (v_bool.size, v_bool.alignment, v_bool.offset) == (1, 1, 72)
+ assert dd.total_size == 80
+
+ assert v_f64.default is MISSING
+ assert not v_f64.default_is_factory
+ assert v_f32.default == 8.0
+ assert v_str.default == "default"
+ assert all(f.structural_eq is None for f in (*base.fields,
*derived.fields, *dd.fields))
+ assert not v_i64.frozen
+
+ # `def_ro` fields are frozen; `def_rw` fields with `default_value` keep
their default.
+ tob_type_info = TestObjectBase.__tvm_ffi_type_info__ # ty:
ignore[unresolved-attribute]
+ tob = {f.name: f for f in ObjectInfo.from_type_info(tob_type_info).fields}
+ assert tob["v_f64"].frozen
+ assert not tob["v_i64"].frozen
+ assert tob["v_i64"].default == 10
+
+ # The auto-init parameter list carries the same facts, parent fields first.
+ assert [f.schema.offset for f in dd.init_fields] == [24, 32, 40, 48, 56,
72]
+
+
+def test_objectinfo_total_size_requires_own_metadata() -> None:
+ """A type without its own metadata inherits a meaningless size, so it
reports ``None``."""
+ root =
ObjectInfo.from_type_info(_lookup_or_register_type_info_from_type_key("ffi.Object"))
+ assert root.total_size == 24
+ assert root.ancestors == []
+ assert root.parent_type_key is None
+ func =
ObjectInfo.from_type_info(_lookup_or_register_type_info_from_type_key("ffi.Function"))
+ assert func.total_size is None
+ assert func.ancestors == ["ffi.Object"]
+
+
+def test_named_type_schema_keeps_positional_signature() -> None:
+ """Synthetic schemas (function params, tests) carry no layout facts."""
+ schema = NamedTypeSchema("x", TypeSchema("int"))
+ assert (schema.size, schema.alignment, schema.offset) == (None, None, None)
+ assert schema.default is MISSING
+ assert not schema.default_is_factory
+ assert schema.structural_eq is None
+ assert not schema.frozen
+
+
+def test_rust_marker_syntax_parses_rs_file(tmp_path: Path) -> None:
+ """``//`` markers in a ``.rs`` file parse into the same block kinds as
``#`` markers."""
+ rs = tmp_path / "mod.rs"
+ rs.write_text(
+ "\n".join(
+ [
+ f"{C.RUST_SYNTAX.begin} object/demo.Foo",
+ "pub struct Foo;",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.ty_map} a -> b",
+ "",
+ ]
+ ),
+ encoding="utf-8",
+ )
+ assert syntax_for(rs) is C.RUST_SYNTAX
+ info = FileInfo.from_file(rs)
+ assert info is not None
+ assert info.syntax is C.RUST_SYNTAX
+ assert [block.kind for block in info.code_blocks] == ["object", "ty-map"]
+ assert info.code_blocks[0].param == "demo.Foo"
+ assert info.code_blocks[0].lines[1] == "pub struct Foo;"
+
+ # Python markers are plain text inside a Rust file.
+ py_markers = tmp_path / "other.rs"
+ py_markers.write_text(
+ f"{C.PYTHON_SYNTAX.begin} object/demo.Foo\n{C.PYTHON_SYNTAX.end}\n",
encoding="utf-8"
+ )
+ assert FileInfo.from_file(py_markers) is None
+
+
+def test_collect_files_filters_by_generator_exts(tmp_path: Path) -> None:
+ """Directory scans visit only the active generator's extensions; explicit
files always."""
+ py = tmp_path / "a.py"
+ py.write_text(f"{C.PYTHON_SYNTAX.begin}
import-section\n{C.PYTHON_SYNTAX.end}\n")
+ rs = tmp_path / "b.rs"
+ rs.write_text(f"{C.RUST_SYNTAX.begin}
import-section\n{C.RUST_SYNTAX.end}\n")
+
+ python = get_generator("python")
+ assert python.source_exts == frozenset({".py", ".pyi"})
+ assert [f.path for f in collect_files([tmp_path], python.source_exts)] ==
[py.resolve()]
+ assert [f.path for f in collect_files([rs], python.source_exts)] ==
[rs.resolve()]
+ both = collect_files([tmp_path], {".py", ".RS"})
+ assert {f.path for f in both} == {py.resolve(), rs.resolve()}
+
+
+def test_generator_registry_names() -> None:
+ """``--target`` choices follow the registered generators."""
+ assert generator_names() == ["python"]
+ with pytest.raises(ValueError, match="Known generators: python"):
+ get_generator("rust")
