diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml index ab0486de6..7d3d6a2e0 100644 --- a/.github/workflows/ci_test.yml +++ b/.github/workflows/ci_test.yml @@ -209,6 +209,19 @@ jobs: run: | cargo test + # Build the Rust stubgen example against the installed package. Its post-build + # step refreshes the bindings in place, so a diff means the checked-in file was stale. + - name: Run rust stubgen example + if: ${{ matrix.os == 'ubuntu-latest' }} + working-directory: examples/rust_stubgen + env: + CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }} + run: | + cmake -B build -DCMAKE_BUILD_TYPE=Release + cmake --build build + git diff --exit-code -- rust/src/generated + cd rust && cargo run + # Build + test the orcjit wheel. On Linux the action downloads a prebuilt LLVM # tarball from a GitHub Release (built from source under the manylinux # gcc-toolset so the ABI stays under the floor); macOS/Windows install LLVM via diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index 98087e2e9..7b299d461 100644 --- a/docs/packaging/stubgen.rst +++ b/docs/packaging/stubgen.rst @@ -390,7 +390,8 @@ When you run the tool, it: ``prefix`` - Demand a Namespace Rust target only. Adds an ``object/`` block to the file for every object registered directly under the prefix that no processed file defines yet; ``skip`` - leaves one out. Code outside the blocks is preserved. + leaves one out, and a key the ``tvm_ffi`` crate binds itself (every ``ffi.*`` key) + is never added. Code outside the blocks is preserved. .. code-block:: rust diff --git a/examples/rust_stubgen/CMakeLists.txt b/examples/rust_stubgen/CMakeLists.txt index 970ad3ee1..ec22b32c5 100644 --- a/examples/rust_stubgen/CMakeLists.txt +++ b/examples/rust_stubgen/CMakeLists.txt @@ -36,7 +36,7 @@ tvm_ffi_configure_target( STUB_TARGET rust STUB_INIT - ON + OFF STUB_DIR "./rust/src/generated" ) diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md index f4d257e3d..35906b312 100644 --- a/examples/rust_stubgen/README.md +++ b/examples/rust_stubgen/README.md @@ -19,16 +19,17 @@ `tvm-ffi-stubgen --target rust` turns the reflection metadata of a C++ library into Rust bindings. This example registers one object, `rust_stubgen.IntPair` -(`src/int_pair.cc`), and lets CMake regenerate `rust/src/generated/` after -every build. +(`src/int_pair.cc`), and binds it in `rust/src/generated/rust_stubgen/mod.rs`: +a file that mixes generated blocks with hand-written code, and that CMake +refreshes in place after every build. Every object gets a `#[repr(C)]` wrapper, a reference type, `Deref`, and the upcasts along its ancestor chain. `IntPair` is plain data, so its reflected fields account for every byte and the binding is *complete*: the struct mirrors the fields at their real offsets and widths, a `const` assertion pins its size -and alignment to the reflected facts, and a generated `new` allocates the object -in Rust. `main.rs` builds one that way, reads `pair.a` directly, and hands it to -a C++ function that reads it back. +and alignment to the reflected facts, and a generated allocator builds the +object in Rust. `main.rs` builds one that way, reads `pair.a` directly, and +hands it to a C++ function that reads it back. An object whose layout cannot be reproduced (a polymorphic one, say, with a vtable in front of the object header) is bound *opaquely* instead: the struct @@ -42,7 +43,7 @@ derived type depth matches the registry. ## Build and run ```bash -# 1. Build the C++ library; the post-build step runs the stub generator. +# 1. Build the C++ library; the post-build step refreshes the bindings in place. cmake -B build -DCMAKE_BUILD_TYPE=Release cmake --build build @@ -50,27 +51,71 @@ cmake --build build cd rust && cargo run ``` +The program prints: + +```text +a=2 b=1 kind=PairKind(0) +sum=3 +``` + The Rust crate depends on the `tvm-ffi` crate of this repository and needs `tvm-ffi-config` on `PATH` (activate the virtual environment where the `apache-tvm-ffi` package is installed). +## Demand-driven layout + +The binding file opens with a `prefix` line: + +```rust +// tvm-ffi-stubgen(prefix): rust_stubgen +``` + +It makes the file own the `rust_stubgen` namespace: every object registered +directly under it gets an `object/` block in this file on the next +run, and a `skip` line (`// tvm-ffi-stubgen(skip): rust_stubgen.Internal`) +leaves one out. Add a second object to `int_pair.cc` and rebuild, and its block +appears after `IntPair`; delete a block and it comes back, so dropping a binding +means writing `skip`. Code outside the blocks is never touched. + +Nothing requires a `generated/` directory: the blocks can live in any `.rs` +file the command line points at. This example keeps them under +`rust/src/generated/` only to make the split visible. + +## Hand-written constructors + +`IntPair::new` is hand-written. The `custom-new` directive makes the generator +name its lossless allocator `from_complete_fields` instead of `new`, and the +hand-written constructor calls it after deriving `kind` from the operands: + +```rust +// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair + +impl IntPair { + pub fn new(a: i64, b: i64) -> Self { + let kind = if a <= b { PairKind::Ordered } else { PairKind::Unordered }; + Self::from_complete_fields(a, b, kind) + } +} +``` + +Without the directive the block would define `IntPair::new` too and the crate +would not compile. + ## Directives The generated file keeps one-line directives the generator reads on every run. -`rust/src/generated/rust_stubgen/mod.rs` declares the integer field `kind` as an -open newtype: +Besides `prefix` and `custom-new`, this example declares the integer field +`kind` as an open newtype: ```rust // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } ``` -Four more directives are available: `field` names the Rust type of a field +Three more are available: `field` names the Rust type of a field (`// tvm-ffi-stubgen(field): rust_stubgen.IntPair.a -> MyInt`), `nullable` wraps it in `Option` (`// tvm-ffi-stubgen(nullable): rust_stubgen.IntPair.a`), -`upcast` adds a conversion to a hand-written typed view -(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`), and -`custom-new` names the generated allocator `from_complete_fields` when `new` -is hand-written (`// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair`). +and `upcast` adds a conversion to a hand-written typed view +(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`). ## Partial generation @@ -82,3 +127,20 @@ struct is named `Obj` (`// tvm-ffi-stubgen(ty-map): rust_stubgen.IntPair -> crate::hand::IntPair`). Anything else is an error listing the missing keys, so a partial binding never references a module that does not exist. + +## Workflow + +1. Write the skeleton: `rust/src/generated/mod.rs` with `pub mod rust_stubgen;` + and `rust_stubgen/mod.rs` holding nothing but the `prefix` line. +2. Build. The post-build step adds the import section and one block per + object. Read the generated code, then add directives and hand-written code + outside the blocks. +3. Rebuild whenever the C++ side changes, or run the post-build command + directly; a `repo: local` pre-commit hook can do the same: + + ```bash + tvm-ffi-stubgen rust/src/generated --target rust --dlls build/librust_stubgen.so + ``` + +4. In CI, run the same command with `--check`. It writes nothing and exits with + status 1 when a block is out of date (2 when a file fails to process). diff --git a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs index 01afb51b3..cda7a4d6a 100644 --- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs +++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs @@ -20,6 +20,9 @@ //! FFI bindings for `rust_stubgen` (generated by tvm-ffi-stubgen). +// Every object registered under `rust_stubgen` gets its block in this file; `skip` leaves one out. +// tvm-ffi-stubgen(prefix): rust_stubgen + // tvm-ffi-stubgen(begin): import-section use std::ops::Deref; use tvm_ffi::Error; @@ -32,6 +35,8 @@ use tvm_ffi::VALUE_ERROR; // The `kind` field is an integer on the C++ side; this directive types it as an // open integer newtype in Rust. // tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) { Unordered=0, Ordered=1 } +// `IntPair::new` is hand-written below, so the generated allocator is `from_complete_fields`. +// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair // tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -98,9 +103,19 @@ impl IntPairObj { impl IntPair { /// Lossless complete-field allocation. - pub fn new(a: i64, b: i64, kind: PairKind) -> Self { + pub fn from_complete_fields(a: i64, b: i64, kind: PairKind) -> Self { let obj = IntPairObj::new(a, b, kind); Self { base: ObjectArc::new(obj) } } } // tvm-ffi-stubgen(end) + +// Hand-written code lives outside the blocks; the generator never touches it. +impl IntPair { + /// Allocate a pair whose `kind` follows from the operands: ordered when + /// `a <= b`, unordered otherwise. + pub fn new(a: i64, b: i64) -> Self { + let kind = if a <= b { PairKind::Ordered } else { PairKind::Unordered }; + Self::from_complete_fields(a, b, kind) + } +} diff --git a/examples/rust_stubgen/rust/src/main.rs b/examples/rust_stubgen/rust/src/main.rs index f49fadbef..e3ca7d898 100644 --- a/examples/rust_stubgen/rust/src/main.rs +++ b/examples/rust_stubgen/rust/src/main.rs @@ -41,10 +41,11 @@ fn main() -> Result<()> { let _lib = Module::load_from_file(lib_path())?; // The object has a reproducible layout: it is allocated in Rust and its - // fields are plain struct members, on both sides of the ABI. - let pair = IntPair::new(1, 2, PairKind::Ordered); + // fields are plain struct members, on both sides of the ABI. `IntPair::new` + // is hand-written next to the generated block and derives `kind` itself. + let pair = IntPair::new(2, 1); println!("a={} b={} kind={:?}", pair.a, pair.b, pair.kind); - assert_eq!(pair.kind, PairKind::Ordered); + assert_eq!(pair.kind, PairKind::Unordered); let sum: i64 = tvm_ffi::cached_global_func!("rust_stubgen.IntPairSum") .call_tuple((pair.clone(),))? diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 4f96e9aed..af8f8c247 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -98,7 +98,7 @@ def __main__() -> int: # Stage 2b. Add the object blocks a `tvm-ffi-stubgen(prefix)` file asks for. This runs # after `--init`, which rewrites files on disk and reloads them. - failed += _roll_out_prefixes(files) + failed += _roll_out_prefixes(files, generator) # Stage 3: Process # - `tvm-ffi-stubgen(begin): global/...` @@ -166,15 +166,14 @@ def _stage_1( ty_map[lhs.strip()] = rhs.strip() -def _roll_out_prefixes(files: list[FileInfo]) -> int: +def _roll_out_prefixes(files: list[FileInfo], generator: Generator) -> int: """Append an ``object/`` block for each registered object under a file's ``prefix``. - Keys with a block in any file of the run, or named by ``skip``, are left alone; an - ``import-section`` is added when the file has none. Returns the number of bad files. + Keys with a block in any file of the run, named by ``skip``, or bound by the target's + runtime are left alone; an ``import-section`` is added when the file has none. Returns + the number of bad files. """ - defined = { - code.param for file in files for code in file.code_blocks if code.kind == "object" - } | C.BUILTIN_TYPE_KEYS + defined = {code.param for file in files for code in file.code_blocks if code.kind == "object"} registry = collect_type_keys() owners: dict[str, Path] = {} failed = 0 @@ -200,7 +199,11 @@ def _roll_out_prefixes(files: list[FileInfo]) -> int: ) continue skipped = {c.param[1].strip() for c in directives if c.param[0] == "skip"} - keys = [key for key in registry[prefix] if key not in defined and key not in skipped] + keys = [ + key + for key in registry[prefix] + if key not in defined and key not in skipped and not generator.is_builtin(key) + ] blocks = file.code_blocks if not any(c.kind == "import-section" for c in blocks): at = blocks.index(head) + 1 @@ -251,7 +254,7 @@ def _find_or_insert_file(path: Path) -> FileInfo: } defined_objs: set[str] = { # ty: ignore[invalid-assignment] code.param for file in files for code in file.code_blocks if code.kind == "object" - } | C.BUILTIN_TYPE_KEYS + } skipped: set[str] = { code.param[1].strip() for file in files @@ -275,7 +278,9 @@ def _find_or_insert_file(path: Path) -> FileInfo: [] if prefix in defined_func_prefixes else global_funcs.get(prefix, []), key=lambda f: f.schema.name, ) - objs = sorted(set(obj_names) - defined_objs - skipped) + objs = sorted( + key for key in set(obj_names) - defined_objs - skipped if not generator.is_builtin(key) + ) object_infos = toposort_objects(objs) if not funcs and not object_infos: continue diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py index add7503a5..c22522a42 100644 --- a/python/tvm_ffi/stub/generator.py +++ b/python/tvm_ffi/stub/generator.py @@ -114,6 +114,14 @@ def extra_export_names(self, imports: Any) -> set[str]: """Return extra public-export names implied by the collected imports.""" ... + def is_builtin(self, type_key: str) -> bool: + """Whether the target's runtime binds ``type_key`` itself. + + Such a key gets no ``object/`` block: neither ``--init`` scaffolds one nor a + ``prefix`` directive rolls one out. + """ + ... + # --- per-block generation (mutates `code.lines`) ------------------------ def generate_global_funcs_block( diff --git a/python/tvm_ffi/stub/python_generator/generator.py b/python/tvm_ffi/stub/python_generator/generator.py index cd81821d7..5d784818c 100644 --- a/python/tvm_ffi/stub/python_generator/generator.py +++ b/python/tvm_ffi/stub/python_generator/generator.py @@ -75,6 +75,10 @@ def extra_export_names(self, imports: PythonImports) -> set[str]: """Return extra ``__all__`` names implied by the collected imports.""" return {"LIB"} if imports.has_lib_load else set() + def is_builtin(self, type_key: str) -> bool: + """Whether the ``tvm_ffi`` package binds ``type_key`` itself.""" + return type_key in C.BUILTIN_TYPE_KEYS + # --- per-block generation (mutates `code.lines`) ------------------------ def generate_global_funcs_block( diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py b/python/tvm_ffi/stub/rust_generator/codegen.py index c5da9b4f4..7385dab7f 100644 --- a/python/tvm_ffi/stub/rust_generator/codegen.py +++ b/python/tvm_ffi/stub/rust_generator/codegen.py @@ -52,7 +52,13 @@ from ..layout import Verdict, classify from ..lib_state import object_info_from_type_key from . import consts as C_RUST -from .utils import RustImports, builtin_mirror_name, render_rust_type, rust_ident +from .utils import ( + RustImports, + builtin_mirror_name, + is_crate_type_key, + render_rust_type, + rust_ident, +) if TYPE_CHECKING: from collections.abc import Container @@ -147,7 +153,7 @@ def _provider(self, type_key: str) -> str | None: ``"crate"`` for builtin ``ffi.*`` types, ``"mapped"`` for a ``ty-map`` to a hand-written binding, ``"generated"`` for an ``object/`` block in this run. """ - if type_key.partition(".")[0] in C_RUST.RUST_MOD_MAP: + if is_crate_type_key(type_key): return "crate" if type_key in self.ty_map: return "mapped" diff --git a/python/tvm_ffi/stub/rust_generator/generator.py b/python/tvm_ffi/stub/rust_generator/generator.py index 30b44dbdc..edfc7c078 100644 --- a/python/tvm_ffi/stub/rust_generator/generator.py +++ b/python/tvm_ffi/stub/rust_generator/generator.py @@ -32,7 +32,7 @@ generate_rust_import_section, generate_rust_object, ) -from .utils import RustImports, RustUse +from .utils import RustImports, RustUse, is_crate_type_key if TYPE_CHECKING: from collections.abc import Container @@ -75,6 +75,10 @@ def extra_export_names(self, imports: RustImports) -> set[str]: """No extra export names for Rust.""" return set() + def is_builtin(self, type_key: str) -> bool: + """Every ``ffi.*`` key is bound by the crate (the rule ``codegen`` resolves types with).""" + return is_crate_type_key(type_key) + # --- per-block generation ------------------------------------------------- def generate_global_funcs_block( diff --git a/python/tvm_ffi/stub/rust_generator/utils.py b/python/tvm_ffi/stub/rust_generator/utils.py index d38e35bcc..dcd68e3dd 100644 --- a/python/tvm_ffi/stub/rust_generator/utils.py +++ b/python/tvm_ffi/stub/rust_generator/utils.py @@ -58,6 +58,11 @@ def as_use_line(self) -> str: return f"use {self.path};" +def is_crate_type_key(type_key: str) -> bool: + """Whether the ``tvm_ffi`` crate binds ``type_key`` itself (every ``ffi.*`` key).""" + return type_key.partition(".")[0] in C.RUST_MOD_MAP + + def builtin_mirror_name(type_key: str) -> str: """Name of the header-only stand-in for a builtin type (``ffi.IntEnum -> FfiIntEnumObj``). diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 97a446d21..05449b8c3 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -1369,7 +1369,7 @@ def test_roll_out_matches_the_prefix_exactly( monkeypatch.setattr( stub_cli, "collect_type_keys", - lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"], "ffi": ["ffi.Object"]}, + lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"], "ffi": ["ffi.Array", "ffi.Object"]}, ) monkeypatch.setattr( stub_cli, @@ -1392,7 +1392,7 @@ def test_roll_out_matches_the_prefix_exactly( ) info = FileInfo.from_file(src) assert info is not None - assert stub_cli._roll_out_prefixes([info]) == 0 + assert stub_cli._roll_out_prefixes([info], RUST) == 0 assert [line for block in info.code_blocks for line in block.lines] == [ f"{C.RUST_SYNTAX.directive('prefix')} a.", f"{begin} import-section", @@ -1405,11 +1405,12 @@ def test_roll_out_matches_the_prefix_exactly( end, "mod tail {}", ] - # Builtin type keys live in the crate and are never rolled out. + # Every `ffi.*` key lives in the crate and is never rolled out, `ffi.Array` included: a + # generated non-generic `Array` would shadow the crate's `Array`. src.write_text(f"{C.RUST_SYNTAX.directive('prefix')} ffi\n", encoding="utf-8") info = FileInfo.from_file(src) assert info is not None - assert stub_cli._roll_out_prefixes([info]) == 0 + assert stub_cli._roll_out_prefixes([info], RUST) == 0 assert not any(block.kind == "object" for block in info.code_blocks)