The GitHub Actions job "mainline-only" on tvm-ffi.git/main has failed.
Run started by GitHub user tqchen (triggered by tqchen).

Head commit for run:
dac85897eb501440df67fc8d8800cd5d95d36f33 / Yaxing Cai <[email protected]>
[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 <[email protected]>

Report URL: https://github.com/apache/tvm-ffi/actions/runs/29474569608

With regards,
GitHub Actions via GitBox


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to