[FEAT][Python] Tie Python wrapper lifetime to underlying C++ FFI object (#593)
## Summary
Bind one Python wrapper to one C++ FFI object ("chandle") for the
object's lifetime, so identity is stable: `a.x is a.x`, `id(a.x)` stable
across drop+refetch, and `f(x) is x` for FFI returns. Works on **both
the GIL and free-threaded (3.14t) builds**. Implements the
*PyObjectTying* design.
Before:
```python
a = MyClass(Inner(...))
assert a.x is a.x # False — fresh wrapper per attribute access
assert id(a.x) == id(a.x) # flaky
```
After: both hold, and identity is preserved across a wrapper
death-and-revive cycle whenever the C++ object outlives the wrapper.
## Allocation layout
A two-layer custom-allocator hook lives in core libtvm_ffi:
- `TVMFFIObjectAllocHeader { delete_space }` — 8-byte base header
preceding every Object body.
- `TVMFFICustomAllocator { allocate, context }` — process-wide registry;
libtvm_ffi installs a builtin default at registry init so
`TVMFFIGetCustomAllocator` never returns NULL.
- `TVMFFIGetCustomAllocator` / `TVMFFISetCustomAllocator` — frontends
override the default at module load.
The Python Cython module overrides the global default with
`TVMFFIPyAllocate`, which prepends a 16-byte `PyCustomAllocHeader`
encoding the wrapper binding. The Rust crate
(`ObjectArc::new[_with_extra_items]`) and the Python-defined types in
`extra/dataclass.cc` route through the same registry, so layout and
lifetime semantics are uniform across frontends.
## Binding state machine
State is concentrated in
`python/tvm_ffi/cython/tvm_ffi_python_helpers.h`. Each header word
(`tagged_pyobj`) holds the wrapper back-pointer plus low tag bits,
giving a four-state machine:
| State | Meaning |
|-----------|---------|
| Detached | No wrapper bound |
| Active | Wrapper bound and owns a +1 on the chandle |
| Inactive | Wrapper dead, but its allocation is cached for in-place
revival |
| InTransit | A dealloc is mid-flight (handshake bit) |
Every FFI return funnels through `make_ret_object` (C++ entry
`TVMFFIPyMakeRetObject`), which returns the canonical wrapper for a
chandle when one exists, **reviving an Inactive cached allocation in
place** so a re-fetched wrapper keeps a stable `id()` at the same
address. The cache-vs-free handshake spans three slots — a pre-bump
`tp_dealloc` opens it, `tp_free` settles it, and the C++ weak deleter
(`TVMFFIPyDeleteSpace`) reclaims the block — coordinated so a chandle
that outlives its wrapper keeps the cached bytes, and a genuinely dead
chandle frees them exactly once.
Frontend-allocation is detected by `delete_space` pointer comparison
(`TVMFFIPyIsCanonical`), avoiding a flag bit on `TVMFFIObject`. Chandles
created before the Python allocator is registered
(statically-initialized global functions in libtvm_ffi.so) carry only
the base header and are skipped.
## Free-threaded build
The same tying runs on `Py_GIL_DISABLED`. All free-threaded machinery is
behind `#ifdef Py_GIL_DISABLED`; **the GIL build is byte-identical**
(verified by a function-body-map diff).
- **Per-word spin-lock.** The `tagged_pyobj` word doubles as a spin-lock
(a Locked tag bit, `__atomic_*` CAS) serializing every binding
transition. The lock is held only across short, park-free word/header
edits; alloc/revival run lock-released. The back-off
(`TVMFFIPyLockYield`) detaches the thread state so a concurrent
stop-the-world GC is never starved.
- **The revival UAF, and the fix.** Cython's generated `tp_dealloc`
bumps the wrapper refcount before running `__dealloc__`. On
free-threaded builds that bump makes `PyUnstable_TryIncRef` spuriously
succeed on a wrapper being torn down, so a concurrent `make_ret`
Active-hit could revive a corpse (borrowed-ref UAF). The fix
**replaces** Cython's `tp_dealloc` on each cdef CObject-family carrier
with one hand-built `TVMFFIPyTpDeallocSlot` that runs the binding
transition (bracketed by `PyErr_Get/SetRaisedException`) **before** any
bump, then — stripped of the now-dead bump — GC-untrack (guarded by
GC-ness), a generic `__dict__` clear (guarded by a real
`tp_dictoffset`), and `tp_free`. A plain carrier runs exactly
`transition; tp_free`; Function fires both guards; both are faithful to
Cython's originals minus the bump.
- **Total carrier coverage.** The six carriers (CObject, CContainerBase,
OpaquePyObject, Error, Tensor, Function) are a closed compile-time set,
each wrapped once at init via `TVMFFIPyWrapDealloc`; every heap subtype
is covered for free through `subtype_dealloc`'s base-walk to the nearest
carrier. An import-time **layout guard** `Py_FatalError`s if a non-GC
carrier ever gains a `__dict__`, turning silent owned-member drift into
a loud failure.
- **Active-hit revival** uses `PyUnstable_TryIncRef` / `EnableTryIncRef`
to close the borrowed-read UAF (the CPython "weakmap" pattern, whose
`tp_dealloc` support requirement this implements).
Verified safe through CPython 3.15/3.16: the
`PyUnstable_TryIncRef`/`PyMutex` contract is unchanged, and the 3.15
rule that managed dict/weakref implies `HAVE_GC` does not affect us — no
carrier uses a managed dict.
## Robustness
- **align>8 double-free fix (latent).** The builtin allocator offset the
body by `round_up(header, alignment)` but free subtracted a fixed
`sizeof(header)`, symmetric only for alignment ≤ 8. A 16-aligned
reflection dataclass (`alignof(max_align_t)=16`) was freed 8 bytes
early. Both sides now use a fixed `alignof(max_align_t)` body offset. It
was masked on the GIL build because the symmetric custom allocator
shadows the builtin one for all Python objects.
- **Shutdown guard.** `TVMFFIPyMarkPythonFinalizing` (wired to atexit)
flips an atomic flag read before any `PyGILState_Ensure`, to avoid
acquiring the GIL after Python finalization has begun. Wrapper bytes on
chandles still alive at exit are intentionally leaked — the process is
exiting.
## `_move()` semantics
Under universal cache-on, callback args alias the caller's wrapper and
FFI returns of the same chandle alias the caller's wrapper (one wrapper,
one chandle ref). `_move()` is kept as an API: the rvalue setter on
either side eager-detaches the canonical binding before the C++
`AnyViewToOwnedAny` transfer nulls the source chandle, so a downstream
cache lookup never sees a stale back-pointer.
## Test plan
- [x] New `tests/python/test_pyobject_tying.py` covers
Active/Inactive/InTransit transitions, cache-on aliasing, `_move()`
under cache-on, pickle stress, threading stress, GC integration,
multi-chandle isolation, the weakref limitation, free-threaded
concurrent carrier-type stress (Function/Error/multi-level-heap), and
OpaquePyObject roundtrip/leak.
- [x] `test_function.py::test_rvalue_ref` refactored for cache-on
aliasing semantics.
- [x] Free-threaded (3.14t): crash oracle clean, full suite **2333
passed**, tying tests **43/43**.
- [x] GIL (3.13): full suite **2358 passed**, tying tests **43/43**.
- [x] Rust suite passes.
- [ ] CI: lint, clang-tidy, doc, C++/Python/Rust on Linux x86_64 +
aarch64, macOS arm64, Windows AMD64.
## Out-of-scope follow-ups
- Name-keyed dict cache for `_get_global_func` to deliver id-stability
for static-init Functions (whose chandles predate the Python allocator
and so carry only the base header). Tracked as a TODO in
`function.pxi::_get_global_func`.
Co-authored-by: Yaoyao Ding <dingyaoyao.cs@gmail.com>📚 Documentation | 🚀 Quickstart
Apache TVM FFI is an open ABI and FFI for machine learning systems. It is a minimal, framework-agnostic, yet flexible open convention with the following systems in mind:
These enable broad interoperability across frameworks, libraries, DSLs, and agents; the ability to ship one wheel for multiple frameworks and Python versions (including free-threaded Python); and consistent infrastructure across environments.
Install TVM-FFI with pip, uv or from source:
pip install apache-tvm-ffi pip install torch-c-dlpack-ext # compatibility package for torch <= 2.9
C ABI stability is our top priority.
Status: RFC Main features are complete and ABI stable. We recognize potential needs for evolution to ensure it works best for the machine learning systems community, and would like to work together with the community for such evolution. We plan to stay in the RFC stage for three months from the v0.1.0 release.
Releases during the RFC stage will be 0.X.Y, where bumps in X indicate C ABI-breaking changes and Y indicates other changes. We anticipate the RFC stage will last for three months, then we will start following Semantic Versioning (major.minor.patch) going forward.
Our documentation site includes: