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

Make `a.x is a.x`, `id(a.x)` stable across drop+refetch, and `f(x) is x`
for FFI returns, on both the GIL and free-threaded (3.14t) builds, by
binding one Python wrapper to one C++ FFI object ("chandle") for the
object's lifetime. Implements the PyObjectTying design.

## Allocation layout

A two-layer custom-allocator hook lives in core libtvm_ffi.so:
`TVMFFIObjectAllocHeader { delete_space }`, `TVMFFICustomAllocator
{ allocate, context }`, plus `TVMFFIGetCustomAllocator` /
`TVMFFISetCustomAllocator`. libtvm_ffi installs a builtin default at
registry init, so every Object always carries the 8-byte base header and
`TVMFFIGetCustomAllocator` never returns NULL. The Python Cython module
overrides the global default at module load with `TVMFFIPyAllocate`,
which prepends a 16-byte `PyCustomAllocHeader` encoding the wrapper
binding. `make_object` / `make_inplace_array_object` / `PyClassDeleter`
and the Rust `ObjectArc::new[_with_extra_items]` paths all funnel through
the registry, so Python-defined types and Rust-allocated objects share
the layout and lifetime semantics.

## Binding state machine

State is concentrated in `tvm_ffi_python_helpers.h`. Each header word
(`tagged_pyobj`) holds the wrapper back-pointer plus low tag bits, giving
a four-state machine:

  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. The `tagged_pyobj` word doubles
as a per-word spin-lock (a Locked tag bit, `__atomic_*` CAS) that
serializes every binding transition; the lock is held only across short,
park-free word/header edits, while alloc/revival run lock-released. The
wait/back-off (`TVMFFIPyLockYield`) detaches the thread state so a
concurrent stop-the-world GC is never starved.

The free-threaded revival hazard is Cython's generated `tp_dealloc`,
which bumps the wrapper refcount before running `__dealloc__`: 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. 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 GIL build is byte-identical: all
free-threaded machinery is behind `#ifdef Py_GIL_DISABLED`, and on the
GIL build the post-bump `__dealloc__` hook runs the transition directly
(safe under the GIL) while the slot installer is a no-op stub.

## Robustness

A double-free in the builtin allocator is fixed along the way: allocate
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. A 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.

## `_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.

## Tests

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` is refactored for the new
aliasing-aware use_count expectations. Verified: FT crash oracle clean,
FT suite 2321 passed, GIL suite 2346 passed, tying tests 43/43 on both
builds; Rust suite passes.

Author: cyx-6

CMakeLists.txt

@@ -70,6 +70,7 @@ set(_tvm_ffi_objs_sources
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/dtype.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/container.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/init_once.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/custom_allocator.cc"
 )
 
 set(_tvm_ffi_extra_objs_sources
@@ -274,20 +275,11 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
     VERBATIM
   )
 
-  if (Python_VERSION VERSION_GREATER_EQUAL "3.12" AND NOT PYTHON_IS_FREE_THREADED)
-    # >= Python3.12, use Use_SABI version
-    python_add_library(tvm_ffi_cython MODULE "${_core_cpp}" USE_SABI 3.12)
-    target_link_libraries(tvm_ffi_cython PRIVATE Python::SABIModule)
-    set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
-    if (NOT WIN32)
-      target_link_libraries(tvm_ffi_cython PRIVATE Python::Module)
-      set_target_properties(tvm_ffi_cython PROPERTIES SUFFIX ".abi3.so")
-    endif ()
-  else ()
-    # before Python3.12, use WITH_SOABI version
-    python_add_library(tvm_ffi_cython MODULE "${_core_cpp}" WITH_SOABI)
-    set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
-  endif ()
+  # The PyObject-tying impl in tvm_ffi_python_helpers.h uses full Python C API (Py_IncRef,
+  # PyObject_GC_Del, atomic header reads), so we build against the per-version ABI rather than the
+  # limited (abi3) ABI.
+  python_add_library(tvm_ffi_cython MODULE "${_core_cpp}" WITH_SOABI)
+  set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
   target_include_directories(
     tvm_ffi_cython PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython
   )

include/tvm/ffi/c_api.h

@@ -582,6 +582,66 @@ TVM_FFI_DLL int TVMFFIObjectDecRef(TVMFFIObjectHandle obj);
 TVM_FFI_DLL int TVMFFIObjectCreateOpaque(void* handle, int32_t type_index,
                                          void (*deleter)(void* handle), TVMFFIObjectHandle* out);
 
+//-----------------------------------------------------------------------
+// Section: ObjectAllocHeader and CustomAllocator
+//-----------------------------------------------------------------------
+/*!
+ * \brief Mandatory header placed immediately before each TVMFFIObject body.
+ *
+ * This header may be used by TVMFFIObject::deleter to reclaim space when a
+ * custom allocator is present. It can also be set to NULL if
+ * TVMFFIObject::deleter directly calls system free. This section must be
+ * available for each Object so a frontend can rely on this field to confirm
+ * if the object came from a certain allocator.
+ */
+typedef struct {
+  /*!
+   * \brief Free the allocation.
+   * \param ptr The pointer to the space of the object.
+   * \note ``ptr`` points to the space of TVMFFIObject and does not include
+   *       the TVMFFIObjectAllocHeader.
+   */
+  void (*delete_space)(void* ptr);
+} TVMFFIObjectAllocHeader;
+
+/*!
+ * \brief Custom allocator entry registered with TVMFFISetCustomAllocator.
+ */
+typedef struct {
+  /*!
+   * \brief Allocate the space for an Object body.
+   * \param size The size requested for the object body.
+   * \param alignment The alignment requirement for the object body.
+   * \param type_index Type index of the object.
+   * \param context The ``context`` field of the registered allocator.
+   * \return Pointer to the space of the object, or NULL on failure (with
+   *         the error reported via ``TVMFFIErrorSetRaised``).
+   * \note The returned pointer must be preceded by a
+   *       ``TVMFFIObjectAllocHeader`` whose ``delete_space`` releases the
+   *       full underlying allocation when invoked.
+   */
+  void* (*allocate)(size_t size, size_t alignment, int32_t type_index, void* context);
+  /*! \brief Allocator context passed unmodified to ``allocate``. */
+  void* context;
+} TVMFFICustomAllocator;
+
+/*!
+ * \brief Get the process-wide custom allocator.
+ * \return The currently registered allocator (never NULL).
+ * \note ``TVMFFIGetCustomAllocator`` always returns a valid allocator and
+ *       can be overridden by ``TVMFFISetCustomAllocator``.
+ */
+TVM_FFI_DLL TVMFFICustomAllocator* TVMFFIGetCustomAllocator(void);
+
+/*!
+ * \brief Register the process-wide custom allocator.
+ * \param allocator Pointer to a TVMFFICustomAllocator, or NULL to restore
+ *                  the builtin default.
+ * \return 0 on success, nonzero on failure.
+ * \note ``allocator`` must be alive throughout the lifetime of the process.
+ */
+TVM_FFI_DLL int TVMFFISetCustomAllocator(TVMFFICustomAllocator* allocator);
+
 /*!
  * \brief Convert type key to type index.
  * \param type_key The key of the type.

include/tvm/ffi/memory.h

@@ -27,6 +27,7 @@
 
 #include <cstddef>
 #include <cstdlib>
+#include <new>
 #include <type_traits>
 #include <utility>
 
@@ -79,6 +80,39 @@ TVM_FFI_INLINE void* AlignedAlloc(size_t size) {
 #endif
 }
 
+/*!
+ * \brief Allocate aligned memory with a runtime-known alignment.
+ *
+ * Sibling of the templated ``AlignedAlloc<align>`` for callers that
+ * receive ``align`` as a parameter (e.g. custom-allocator
+ * implementations dispatching on ``TVMFFICustomAllocator::allocate``'s
+ * runtime ``alignment`` argument).
+ *
+ * \param size The size.
+ * \param align The alignment, must be a power of 2.
+ * \return The pointer to the allocated memory.
+ */
+TVM_FFI_INLINE void* AlignedAllocRuntime(size_t size, size_t align) {
+#ifdef _MSC_VER
+  if (void* ptr = _aligned_malloc(size, align)) {
+    return ptr;
+  }
+  throw std::bad_alloc();
+#else
+  if (align <= alignof(std::max_align_t)) {
+    if (void* ptr = std::malloc(size)) {
+      return ptr;
+    }
+    throw std::bad_alloc();
+  }
+  void* ptr;
+  if (posix_memalign(&ptr, align, size) != 0) {
+    throw std::bad_alloc();
+  }
+  return ptr;
+#endif
+}
+
 /*!
  * \brief Free aligned memory.
  * \param data The pointer to the memory to free.
@@ -168,7 +202,11 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
       // class with non-virtual destructor.
       // We are fine here as we captured the right deleter during construction.
       // This is also the right way to get storage type for an object pool.
-      void* data = AlignedAlloc<alignof(T)>(sizeof(T));
+      static_assert(alignof(T) <= alignof(::std::max_align_t),
+                    "Object types with alignment > max_align_t are not supported "
+                    "by the custom allocator hook");
+      TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
+      void* data = alloc->allocate(sizeof(T), alignof(T), T::RuntimeTypeIndex(), alloc->context);
       new (data) T(std::forward<Args>(args)...);
       return reinterpret_cast<T*>(data);
     }
@@ -187,7 +225,8 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
         tptr->T::~T();
       }
       if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
-        AlignedFree(static_cast<void*>(tptr));
+        ObjectUnsafe::GetObjectAllocHeaderFromPtr(static_cast<void*>(tptr))
+            ->delete_space(static_cast<void*>(tptr));
       }
     }
   };
@@ -215,12 +254,17 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
       static_assert(
           alignof(ArrayType) % alignof(ElemType) == 0 && sizeof(ArrayType) % alignof(ElemType) == 0,
           "element alignment constraint");
+      static_assert(alignof(ArrayType) <= alignof(::std::max_align_t),
+                    "Object types with alignment > max_align_t are not supported "
+                    "by the custom allocator hook");
       size_t size = sizeof(ArrayType) + sizeof(ElemType) * num_elems;
       // round up to the nearest multiple of align
       constexpr size_t align = alignof(ArrayType);
       // C++ standard always guarantees that alignof operator returns a power of 2
       size_t aligned_size = (size + (align - 1)) & ~(align - 1);
-      void* data = AlignedAlloc<align>(aligned_size);
+      TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
+      void* data =
+          alloc->allocate(aligned_size, align, ArrayType::RuntimeTypeIndex(), alloc->context);
       new (data) ArrayType(std::forward<Args>(args)...);
       return reinterpret_cast<ArrayType*>(data);
     }
@@ -239,7 +283,8 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
         tptr->ArrayType::~ArrayType();
       }
       if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
-        AlignedFree(static_cast<void*>(tptr));
+        ObjectUnsafe::GetObjectAllocHeaderFromPtr(static_cast<void*>(tptr))
+            ->delete_space(static_cast<void*>(tptr));
       }
     }
   };

include/tvm/ffi/object.h

@@ -1141,6 +1141,16 @@ struct ObjectUnsafe {
     return const_cast<TVMFFIObject*>(&(src->header_));
   }
 
+  /*!
+   * \brief Recover the TVMFFIObjectAllocHeader for a TVMFFIObject pointer.
+   * \param ptr The pointer to the space of the object.
+   * \return The header set by the allocator that produced ``ptr``.
+   */
+  TVM_FFI_INLINE static TVMFFIObjectAllocHeader* GetObjectAllocHeaderFromPtr(void* ptr) {
+    return reinterpret_cast<TVMFFIObjectAllocHeader*>(static_cast<char*>(ptr) -
+                                                      sizeof(TVMFFIObjectAllocHeader));
+  }
+
 // Suppress -Winvalid-offsetof: we intentionally use offsetof on non-standard-layout types
 // to avoid undefined behavior from null pointer arithmetic that sanitizers flag.
 #if defined(__clang__) || defined(__GNUC__)

pyproject.toml

@@ -232,14 +232,19 @@ docstring-code-line-length = 80
 [tool.cibuildwheel]
 build-verbosity = 1
 
-# only build up to cp312, cp312
-# will be abi3 and can be used in future versions
-# ship 314t threaded nogil version
-build = ["cp38-*", "cp39-*", "cp310-*", "cp311-*", "cp312-*", "cp314t-*"]
+# Per-Python-version wheels (no abi3 / limited API).
+build = [
+  "cp38-*",
+  "cp39-*",
+  "cp310-*",
+  "cp311-*",
+  "cp312-*",
+  "cp313-*",
+  "cp314t-*",
+]
 skip = ["*musllinux*"]
 # we only need to test on cp312
 test-skip = ["cp38-*", "cp39-*", "cp310-*", "cp311-*"]
-# focus on testing abi3 wheel
 build-frontend = "build[uv]"
 test-command = "pytest {package}/tests/python -vvs"
 test-groups = ["test"]

python/tvm_ffi/cython/base.pxi

@@ -361,6 +361,21 @@ def _env_get_current_stream(int device_type, int device_id):
 
 
 cdef extern from "tvm_ffi_python_helpers.h":
+    int TVMFFIPyRegisterDefaultAllocator() noexcept
+    void TVMFFIPyMarkPythonFinalizing() noexcept
+
+    void TVMFFIPyRebindPyObject(void* chandle, PyObject* expect, PyObject* neo) noexcept
+    void TVMFFIPyTpDealloc(void** ptr_to_chandle, PyObject* wrapper) noexcept
+    void TVMFFIPyInstallTypeSlots(PyObject* type_obj) noexcept
+    object TVMFFIPyMakeRetObject(void* chandle, PyObject* cls_type)
+
+    # Pre-bump tp_dealloc installer (free-threaded builds only; a no-op symbol on the GIL build).
+    # Called once per cdef CObject-family carrier right after the class is defined: it replaces
+    # Cython's generated tp_dealloc with a hand-built slot that runs the binding transition before
+    # any resurrection bump. One installer for every carrier (the slot is universal); a layout
+    # guard inside it fails loudly at import if a carrier's struct ever drifts.
+    void TVMFFIPyWrapDealloc(PyObject* type_obj) noexcept
+
     # no need to expose fields of the call context setter data structure
     ctypedef struct TVMFFIPyCallContext:
         int device_type
@@ -542,5 +557,11 @@ cdef _init_env_api():
 
 _init_env_api()
 
+
+CHECK_CALL(TVMFFIPyRegisterDefaultAllocator())
+
+import atexit as _tvm_ffi_atexit
+_tvm_ffi_atexit.register(TVMFFIPyMarkPythonFinalizing)
+
 # ensure testing is linked and we can run testcases
 TVMFFITestingDummyTarget()

python/tvm_ffi/cython/error.pxi

@@ -122,6 +122,10 @@ cdef class Error(CObject):
         return make_ret_object(any_val)
 
 
+# Install the free-threaded pre-bump tp_dealloc slot on this cdef carrier (no-op on the GIL).
+TVMFFIPyWrapDealloc(<PyObject*>Error)
+
+
 cdef inline Error move_from_last_error():
     # raise last error
     error = Error.__new__(Error)

python/tvm_ffi/cython/function.pxi

@@ -545,9 +545,17 @@ cdef int TVMFFIPyArgSetterObjectRValueRef_(
     PyObject* py_arg, TVMFFIAny* out
 ) except -1:
     """Setter for ObjectRValueRef"""
-    cdef object arg = <object>py_arg
+    cdef CObject src = (<object>py_arg).obj
+    # need to detach from chandle
+    # there are two possible outcomes after the call:
+    #   chandle gets moved, so it is set to NULL
+    #   callee did not move chandle, in such case src.chandle is valid
+    #     but chandle is no longer attached to PyObject
+    # we need to carefully handle chandle and PyObject recycling in both cases.
+    # These logics are implemented in TVMFFIPyTpDealloc (CObject.__dealloc__).
+    TVMFFIPyRebindPyObject(src.chandle, <PyObject*>src, NULL)
     out.type_index = kTVMFFIObjectRValueRef
-    out.v_ptr = &((<CObject>(arg.obj)).chandle)
+    out.v_ptr = &(src.chandle)
     return 0
 
 
@@ -1075,6 +1083,10 @@ cdef class Function(CObject):
         return func
 
 
+# Install the free-threaded pre-bump tp_dealloc slot on this cdef carrier (no-op on the GIL).
+TVMFFIPyWrapDealloc(<PyObject*>Function)
+
+
 def _register_global_func(name: str, pyfunc: Callable[..., Any] | Function, override: bool) -> Function:
     cdef TVMFFIObjectHandle chandle
     cdef int c_api_ret_code
@@ -1089,6 +1101,7 @@ def _register_global_func(name: str, pyfunc: Callable[..., Any] | Function, over
 
 
 def _get_global_func(name: str, allow_missing: bool):
+    # PyObject tying is not applicable here.
     cdef TVMFFIObjectHandle chandle
     cdef ByteArrayArg name_arg = ByteArrayArg(c_str(name))