[FEAT][Python] Tie Python wrapper lifetime to underlying C++ FFI object

[cyx-6]

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:

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_FatalErrors 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

• 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.
• test_function.py::test_rvalue_ref refactored for cache-on aliasing semantics.
• Free-threaded (3.14t): crash oracle clean, full suite 2333 passed, tying tests 43/43.
• GIL (3.13): full suite 2358 passed, tying tests 43/43.
• 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.

[gemini-code-assist]

Code Review

This pull request introduces 'PyObjectTying,' a mechanism to link the lifetime of Python wrappers to their underlying C++ FFI objects. By prepending a PyCustomizeAllocHeader to object allocations, the implementation enables stable object identity and address preservation across wrapper finalization cycles. The changes include modifications to the C++ SimpleObjAllocator, Cython bindings to utilize tp_finalize, and updates to the reflection system. Feedback from the review highlights a critical thread-safety issue in the callback used to free cached wrapper memory, which lacks GIL protection, and a potential memory leak when installing handles if a previous phantom wrapper is not correctly released.

[junrushao]

Let's double check and test it out, making sure this __del__ method will be properly invoked. I had some really negative experience with overloading __del__ because it could possibly prevent Python GC from recycling memory in certain cases. Not 100% sure if it applies to Cython's cdef class's __del__ method though.

[cyx-6]

a tp_finalize-free version is updated. we are now using tp_alloc + tp_free as alternative.

[tqchen]

comment, this is a new class, so it must be in detached state, attach it

[tqchen]

NOTE: we don't try to attach py object for tensor types as it may appear in callback and there maybe translation depending on the context

[tqchen]

When a chandle is destroyed, we can only be in Detached or Inactive state, if it is in Inactive state, the PyObject is reclaimed

[tqchen]

need clarification here, since Detached menas chandle == NULL

[tqchen]

comment Detached

[tqchen]

no need to DecRef here? just return and pass to dealloc

[junrushao]

Btw - does this approach also work for String and Bytes? I’m asking because I wasn’t sure if small string optimization is compatible with this design

[cyx-6]

@junrushao Not working for String and Bytes. For small ones, they are not heap-allocated. For large ones, they are not sub-classes of CObject. Both are out of the scope of this PR.

[tqchen]

try inline

[tqchen]

use std::swap

[tqchen]

PyObject*

[tqchen]

TLS

[tqchen]

static PyObject* custom_tp_alloc(PyTypeObject *type, Py_ssize_t nitems) {
    PyCustomizeAllocHeader *header = get_current_header();

    // Hot Path: Reuse
    if (header && header->cached_mem) {
        PyObject *obj = header->cached_mem;
        
        // Inline-friendly initialization
        PyObject_VarInit((PyVarObject *)obj, type, nitems);
        
        // Only memset the delta between header and body
        // This is where your real time is spent
        size_t head_sz = (nitems > 0 || type->tp_itemsize > 0) ? sizeof(PyVarObject) : sizeof(PyObject);
        memset((char*)obj + head_sz, 0, type->tp_basicsize - head_sz);

        if (PyType_IS_GC(type)) {
            PyObject_GC_Track(obj);
        }
        return obj;
    }

    // Cold Path: Slow allocation
    return PyType_GenericAlloc(type, nitems);
}

[tqchen]

add a runtime offset check test case

[tqchen]

void** TVMFFICyObjectGetCHandlePtr(PyObject* ptr)


cdef void** TVMFFICyObjectGetCHandlePtr(PyObject* ptr):
   return &((<CObject>ptr).chandle)

[tqchen]

move TVMFFIPyTpDealloc after TVMFFIPyTpAlloc, TVM_FFI_PY_INACTIVE_SENTINEL/TVMFFIPyTpDealloc/ TVMFFIPyTpFree/ sit together, document the overall logic in a section, TVM_FFI_PY_INACTIVE_SENTINEL=> TVM_FFI_PY_CHANDLE_INACTIVE_SENTINEL

[tqchen]

Document, after dealloc, there are two possibilities:

• chandle is set to CHANDLE_INACTIVE_SENTINEL, tpfree will not be triggered
• other cases, normal free

[tqchen]

cross check

[tqchen]

document

\return whether returned PyObject state is active

[junrushao]

Any updates on this issue?