[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 by attaching a 16-byte tagged-pointer header to every
Object allocation. Implements Tianqi's PyObjectTying doc (2026-05-01)
with adaptations forced by Py_LIMITED_API + Cython 3.x constraints.

Design:
  - Two-layer custom-allocator hook in core libtvm_ffi.so:
    `TVMFFICustomAllocHeader { delete_space }`,
    `TVMFFICustomAllocator { allocate }`, plus
    `TVMFFIGetCustomAllocator` / `TVMFFISetCustomAllocator`. libtvm_ffi
    installs a builtin default at registry construction so every
    Object always carries the 8-byte base header. The Python Cython
    module overrides the global default at module load with
    `TVMFFIPyAllocate`, which prepends a 16-byte `PyCustomAllocHeader`
    encoding the Python wrapper binding.
  - Tagged-pointer state machine. `PyCustomAllocHeader` is a single
    `uintptr_t tagged_wrapper` field plus the base header (16 B
    total). Low bit of `tagged_wrapper` is the "live" tag:
      (a) tagged_wrapper == 0                  : never wrapped
      (b) tagged_wrapper == (uintptr_t)w | 1   : live wrapper at w
      (c) tagged_wrapper == (uintptr_t)w       : preserved-mem at w
    CPython PyObjects are >= 8-byte aligned (PyObject_HEAD = 16 B), so
    bit 0 is always free for the tag. Eliminates the 8-byte reserved
    pad an earlier draft needed.
  - Universal cache-on for `make_ret`. Every FFI return (function
    results, callback args, FieldGetter) funnels through
    `make_ret_object` which returns the canonical wrapper for a
    chandle when one exists. Combined with state-(c) revive on the
    cached path, this delivers `a.x is a.x`, stable `id(a.x)` across
    drop+refetch, and `f(x) is x` when a function returns the same
    chandle it received. The doc proposed a `tp_alloc` + TLS hand-off
    for revive; not viable under Py_LIMITED_API
    (`PyType_SetSlot` not in 3.12 SABI; Cython 3.x forces
    `CYTHON_USE_TYPE_SPECS=0` for limited-API targets). The inline
    `make_ret_object` revive branch is the SABI-compatible equivalent
    -- bypasses `tp_alloc` by returning the cached PyObject directly,
    strictly cheaper than the doc's design (no alloc+dealloc cycle on
    revive).
  - State (c) via Cython `def __del__` on `CObject`, mapping to
    `tp_finalize` (PEP 442). When other C++ holders keep the chandle
    alive past the wrapper's Python refcount hitting 0, clear the
    live bit (keep the pointer), DecRef the chandle, `Py_INCREF(self)`
    to resurrect. Wrapper memory survives at the same address until
    revive or chandle death. Shim for
    `PyObject_CallFinalizerFromDealloc` (added to SABI in 3.13)
    using `PyType_GetSlot` + `Py_SET_REFCNT`/`Py_REFCNT` on 3.12;
    paired with `-DCYTHON_USE_TP_FINALIZE=1`.
  - Shutdown guard. `g_tvm_ffi_python_alive` flag is set to 0 by an
    atexit handler registered at Cython module init; read by
    `TVMFFIPyDeleteSpace` before `PyGILState_Ensure` to avoid GIL
    acquire after `Py_Finalize` has begun. State-(c) wrapper bytes
    on chandles still alive at that point are intentionally leaked
    -- process is exiting.
  - Frontend-allocation detection by `delete_space` pointer comparison
    (`TVMFFIPyIsCanonical`) avoids a flag bit on `TVMFFIObject`.
    Pre-Python-init chandles (statically-initialized global functions
    in libtvm_ffi.so) carry only the base header; the Python side
    detects this and skips the binding install.

Trade-off on `_move()` semantics. Universal cache-on means callback
args alias the caller's wrapper (one wrapper, one chandle ref) and
function returns of the same chandle alias the caller's wrapper.
`_move()` semantics are kept as an API: the rvalue setter on either
caller or callback side eager-detaches the canonical-wrapper binding
before the C++ AnyViewToOwnedAny transfer nulls the source chandle.
`test_function.py::test_rvalue_ref` refactored for the new
aliasing-aware use_count expectations.

Refcount fix. The `make_ret_object` state-(b) cache hit path had a
redundant `Py_INCREF(cached)` leaking +1 wrapper ref per cache hit
(verified via `sys.getrefcount`). The `<object>(<PyObject*>...)` cast
already INCREFs for the owned reference; `return cached` transfers it
to the caller. Removed the explicit INCREF.

`PyClassDeleter` in `extra/dataclass.cc` and the `__ffi_new__` /
`__ffi_shallow_copy__` paths are routed through the same allocator
registry so Python-defined types share the layout and lifetime
semantics.

`TVMFFIPyArgSetterObjectRValueRef_` clears the source's binding
eagerly before the C++ side nulls its `chandle`; otherwise a
downstream cache lookup on the (still-alive) chandle would briefly
see a stale back-pointer.

Tests: full Python suite passes (2332 passed, 19 skipped, 2 xfailed).
New `tests/python/test_pyobject_tying.py` (32 tests) covers state
b/c, cache-on aliasing, `_move()` under cache-on, pickle stress,
threading stress, GC integration with state (c), multi-chandle
isolation, and the weakref limitation. `test_function.py::test_rvalue_ref`
refactored for cache-on aliasing semantics.

TODO carried in `function.pxi::_get_global_func`: name-keyed dict
cache for static-init Function id-stability. Most registry-resident
Functions are allocated at C++ static init (before the Python
allocator is registered) so their chandle has only the base header --
make_ret_object's cache can't reach them. Deferred to keep this
change small.

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
@@ -289,6 +290,17 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
   target_include_directories(
     tvm_ffi_cython PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython
   )
+  # Force-enable Cython's tp_finalize wiring under USE_SABI. Cython 3.x
+  # disables CYTHON_USE_TP_FINALIZE for limited-API targets because
+  # PyObject_CallFinalizerFromDealloc was added to the SABI only in 3.13.
+  # We need it for the PyObject-tying state-(b)->(c) transition (see
+  # CObject.__del__ in python/tvm_ffi/cython/object.pxi). Safety relies on
+  # the local PyObject_CallFinalizerFromDealloc shim in
+  # python/tvm_ffi/cython/tvm_ffi_python_helpers.h; do NOT remove the
+  # shim while this define is set. The override piggybacks on Cython 3.x's
+  # exact spelling of the call, so a Cython upgrade may require revisiting
+  # both.
+  target_compile_definitions(tvm_ffi_cython PRIVATE CYTHON_USE_TP_FINALIZE=1)
   target_compile_features(tvm_ffi_cython PRIVATE cxx_std_17)
   target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_header)
   target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_shared)

include/tvm/ffi/c_api.h

@@ -580,6 +580,83 @@ 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: Custom allocator hook for Object allocations
+//
+// Frontends (Python, Rust, etc.) can register a custom allocator that the
+// core ObjAllocator consults inside make_object<T> / make_inplace_array_object.
+// When a custom allocator is registered, every allocation of a ref-counted
+// Object is given the allocator's chance to prepend bookkeeping bytes ahead
+// of the object. The allocator returns the address of the embedded T
+// (= Object = TVMFFIObject); immediately before that address it must place a
+// TVMFFICustomAllocHeader whose ``delete_space`` field is invoked from the
+// C++ Weak deleter to release the entire block.
+//
+// The core libtvm_ffi.so does not depend on Python (or any other frontend).
+// libtvm_ffi installs a builtin default allocator at registry init time, so
+// every Object allocation always carries at least the base
+// TVMFFICustomAllocHeader; frontends that need richer per-allocation
+// bookkeeping (e.g. Python's ``py_object`` back-pointer) override the global
+// default with their own allocator.
+//-----------------------------------------------------------------------
+/*!
+ * \brief Header that sits immediately before the embedded TVMFFIObject. The
+ *        C++ deleter recovers this header by walking back a fixed offset
+ *        (``sizeof(TVMFFICustomAllocHeader)``) from the TVMFFIObject pointer
+ *        and invokes ``delete_space``.
+ *
+ * Frontends are free to define a derived layout that places extra fields
+ * before this base; the derived ``delete_space`` callback is responsible for
+ * recovering its enclosing struct from the ``tptr`` argument. Frontends can
+ * detect their own allocations by comparing this header's ``delete_space``
+ * pointer against their own callback (no flag bit needed).
+ */
+typedef struct {
+  /*!
+   * \brief Free the allocation that contains ``tptr``.
+   * \param tptr The pointer originally returned by
+   *             ``TVMFFICustomAllocator::allocate`` (i.e. the address of T,
+   *             not the start of the underlying malloc block).
+   */
+  void (*delete_space)(void* tptr);
+} TVMFFICustomAllocHeader;
+
+/*!
+ * \brief Custom allocator entry registered with TVMFFISetCustomAllocator.
+ */
+typedef struct {
+  /*!
+   * \brief Allocate ``size`` bytes for an Object of ``type_index`` with the
+   *        requested ``alignment``. Implementations must place a
+   *        TVMFFICustomAllocHeader immediately before the returned pointer
+   *        (at the fixed base offset documented on TVMFFICustomAllocHeader)
+   *        with a non-NULL ``delete_space``.
+   * \return Pointer to T's location (treated as TVMFFIObject* by the caller).
+   */
+  void* (*allocate)(size_t size, size_t alignment, int32_t type_index);
+} TVMFFICustomAllocator;
+
+/*!
+ * \brief Get the process-wide custom allocator.
+ *
+ * libtvm_ffi installs a builtin default at registry init, so the result is
+ * never NULL in practice. Frontends can override the global allocator via
+ * TVMFFISetCustomAllocator.
+ *
+ * \return The currently registered allocator.
+ */
+TVM_FFI_DLL TVMFFICustomAllocator* TVMFFIGetCustomAllocator(void);
+
+/*!
+ * \brief Register the process-wide custom allocator. Pass ``allocator==NULL``
+ *        to restore the builtin default.
+ *
+ * \param allocator Pointer to a TVMFFICustomAllocator with process-stable
+ *                  lifetime, or NULL to reset to the builtin default.
+ * \return 0 on success, nonzero on failure.
+ */
+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>
 
@@ -92,6 +93,37 @@ TVM_FFI_INLINE void AlignedFree(void* data) {
 #endif
 }
 
+/*!
+ * \brief Fixed offset between the base TVMFFICustomAllocHeader and the
+ *        embedded T (= Object = TVMFFIObject).
+ *
+ * The C++ deleter recovers the base header from a chandle by walking back
+ * exactly this many bytes: ``tptr - kCustomAllocBaseOffset``. Frontends
+ * with a derived layout add their extra fields ahead of the base.
+ *
+ * We use ``sizeof(TVMFFICustomAllocHeader)`` (one function pointer = 8B)
+ * directly. Alignment of T is enforced separately by the frontend's
+ * choice of total padding before T (must be a multiple of alignof(T));
+ * see PyCustomAllocHeader for an example.
+ */
+constexpr size_t kCustomAllocBaseOffset = sizeof(TVMFFICustomAllocHeader);
+static_assert(alignof(::std::max_align_t) % alignof(TVMFFICustomAllocHeader) == 0,
+              "base header alignment must divide max_align_t so it remains "
+              "aligned when placed just before a max-aligned T");
+
+/*!
+ * \brief Recover the base TVMFFICustomAllocHeader from a TVMFFIObject pointer.
+ *
+ * Every Object allocation is routed through the custom-allocator registry
+ * (libtvm_ffi installs a builtin default at registry init), so this is
+ * always valid for chandles produced by ``make_object`` /
+ * ``make_inplace_array_object``.
+ */
+TVM_FFI_INLINE TVMFFICustomAllocHeader* GetCustomAllocHeader(void* tptr) {
+  return reinterpret_cast<TVMFFICustomAllocHeader*>(static_cast<char*>(tptr) -
+                                                    kCustomAllocBaseOffset);
+}
+
 /*!
  * \brief Base class of object allocators that implements make.
  *  Use curiously recurring template pattern.
@@ -111,7 +143,14 @@ class ObjAllocatorBase {
   ObjectPtr<T> make_object(Args&&... args) {
     using Handler = typename Derived::template Handler<T>;
     static_assert(std::is_base_of_v<Object, T>, "make can only be used to create Object");
-    T* ptr = Handler::New(static_cast<Derived*>(this), std::forward<Args>(args)...);
+    // Always go through the custom-allocator hook. libtvm_ffi installs a
+    // builtin default allocator at registry init time, so the result is
+    // never null and every Object always has a TVMFFICustomAllocHeader
+    // sitting immediately before it. Frontends with richer per-allocation
+    // bookkeeping (e.g. Python's PyCustomAllocHeader) override the global
+    // allocator with their own.
+    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
+    T* ptr = Handler::New(static_cast<Derived*>(this), alloc, std::forward<Args>(args)...);
     TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr);
     ffi_ptr->combined_ref_count = kCombinedRefCountBothOne;
     ffi_ptr->type_index = T::RuntimeTypeIndex();
@@ -132,8 +171,9 @@ class ObjAllocatorBase {
     using Handler = typename Derived::template ArrayHandler<ArrayType, ElemType>;
     static_assert(std::is_base_of_v<Object, ArrayType>,
                   "make_inplace_array can only be used to create Object");
-    ArrayType* ptr =
-        Handler::New(static_cast<Derived*>(this), num_elems, std::forward<Args>(args)...);
+    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
+    ArrayType* ptr = Handler::New(static_cast<Derived*>(this), alloc, num_elems,
+                                  std::forward<Args>(args)...);
     TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr);
     ffi_ptr->combined_ref_count = kCombinedRefCountBothOne;
     ffi_ptr->type_index = ArrayType::RuntimeTypeIndex();
@@ -154,21 +194,21 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
   class Handler {
    public:
     template <typename... Args>
-    static T* New(SimpleObjAllocator*, Args&&... args) {
-      // NOTE: the first argument is not needed for SimpleObjAllocator
-      // It is reserved for special allocators that needs to recycle
-      // the object to itself (e.g. in the case of object pool).
+    static T* New(SimpleObjAllocator*, TVMFFICustomAllocator* alloc, Args&&... args) {
+      // NOTE: the first argument is not needed for SimpleObjAllocator.
+      // It is reserved for special allocators that need to recycle the
+      // object to itself (e.g. an object pool).
       //
-      // In the case of an object pool, an allocator needs to create
-      // a special chunk memory that hides reference to the allocator
-      // and call allocator's release function in the deleter.
-
-      // NOTE2: Use inplace new to allocate
-      // This is used to get rid of warning when deleting a virtual
-      // 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));
+      // Use inplace new to construct T. This avoids the warning when
+      // deleting a virtual class with non-virtual destructor; we captured
+      // the right deleter during construction.
+      static_assert(alignof(T) <= alignof(::std::max_align_t),
+                    "Object types with alignment > max_align_t are not supported "
+                    "by the custom allocator hook");
+      // ``alloc`` is non-null because libtvm_ffi installs a builtin default
+      // allocator at registry init. The matching ``delete_space`` is stored
+      // in the prepended TVMFFICustomAllocHeader and invoked from Deleter_.
+      void* data = alloc->allocate(sizeof(T), alignof(T), T::RuntimeTypeIndex());
       new (data) T(std::forward<Args>(args)...);
       return reinterpret_cast<T*>(data);
     }
@@ -180,14 +220,13 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
       T* tptr =
           details::ObjectUnsafe::RawObjectPtrFromUnowned<T>(static_cast<TVMFFIObject*>(objptr));
       if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) {
-        // It is important to do tptr->T::~T(),
-        // so that we explicitly call the specific destructor
-        // instead of tptr->~T(), which could mean the intention
-        // call a virtual destructor(which may not be available and is not required).
+        // tptr->T::~T() (not tptr->~T()) so we always call the specific
+        // destructor and never pick up a virtual destructor that may not
+        // exist for this Object type.
         tptr->T::~T();
       }
       if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
-        AlignedFree(static_cast<void*>(tptr));
+        GetCustomAllocHeader(static_cast<void*>(tptr))->delete_space(static_cast<void*>(tptr));
       }
     }
   };
@@ -197,30 +236,20 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
   class ArrayHandler {
    public:
     template <typename... Args>
-    static ArrayType* New(SimpleObjAllocator*, size_t num_elems, Args&&... args) {
-      // NOTE: the first argument is not needed for ArrayObjAllocator
-      // It is reserved for special allocators that needs to recycle
-      // the object to itself (e.g. in the case of object pool).
-      //
-      // In the case of an object pool, an allocator needs to create
-      // a special chunk memory that hides reference to the allocator
-      // and call allocator's release function in the deleter.
-      // NOTE2: Use inplace new to allocate
-      // This is used to get rid of warning when deleting a virtual
-      // 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.
-
-      // for now only support elements that aligns with array header.
+    static ArrayType* New(SimpleObjAllocator*, TVMFFICustomAllocator* alloc, size_t num_elems,
+                          Args&&... args) {
+      // For now only support elements that align with the array header.
       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);
+      void* data = alloc->allocate(aligned_size, align, ArrayType::RuntimeTypeIndex());
       new (data) ArrayType(std::forward<Args>(args)...);
       return reinterpret_cast<ArrayType*>(data);
     }
@@ -232,14 +261,10 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
       ArrayType* tptr = details::ObjectUnsafe::RawObjectPtrFromUnowned<ArrayType>(
           static_cast<TVMFFIObject*>(objptr));
       if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) {
-        // It is important to do tptr->ArrayType::~ArrayType(),
-        // so that we explicitly call the specific destructor
-        // instead of tptr->~ArrayType(), which could mean the intention
-        // call a virtual destructor(which may not be available and is not required).
         tptr->ArrayType::~ArrayType();
       }
       if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
-        AlignedFree(static_cast<void*>(tptr));
+        GetCustomAllocHeader(static_cast<void*>(tptr))->delete_space(static_cast<void*>(tptr));
       }
     }
   };

python/tvm_ffi/cython/base.pxi

@@ -15,12 +15,17 @@
 # specific language governing permissions and limitations
 # under the License.
 import ctypes
-from libc.stdint cimport int32_t, int64_t, uint64_t, uint32_t, uint8_t, int16_t
+from libc.stdint cimport int32_t, int64_t, uint64_t, uint32_t, uint8_t, int16_t, uintptr_t
 from libc.string cimport memcpy
 from libcpp.vector cimport vector
 from cpython.bytes cimport PyBytes_AsStringAndSize, PyBytes_FromStringAndSize, PyBytes_AsString
 from cpython cimport Py_INCREF, Py_DECREF, Py_REFCNT
+from cpython.object cimport PyTypeObject
 from cpython cimport PyErr_CheckSignals, PyGILState_Ensure, PyGILState_Release, PyObject
+
+cdef extern from "Python.h":
+    void Py_IncRef(PyObject*) nogil
+    void Py_DecRef(PyObject*) nogil
 from cpython cimport pycapsule, PyCapsule_Destructor
 from cpython cimport PyErr_SetNone
 
@@ -162,6 +167,15 @@ cdef extern from "tvm/ffi/c_api.h":
         uint32_t __padding
         void (*deleter)(void* self, int flags)
 
+    ctypedef struct TVMFFICustomAllocHeader:
+        void (*delete_space)(void* tptr)
+
+    ctypedef struct TVMFFICustomAllocator:
+        void* (*allocate)(size_t size, size_t alignment, int32_t type_index)
+
+    TVMFFICustomAllocator* TVMFFIGetCustomAllocator() nogil
+    int TVMFFISetCustomAllocator(TVMFFICustomAllocator* allocator) nogil
+
     ctypedef struct TVMFFIAny:
         int32_t type_index
         int32_t zero_padding
@@ -361,6 +375,31 @@ def _env_get_current_stream(int device_type, int device_id):
 
 
 cdef extern from "tvm_ffi_python_helpers.h":
+    # 16-byte header prepended to every Object allocated through the
+    # registered Python custom allocator. The single ``tagged_wrapper``
+    # field encodes the binding state (a/b/c): low bit set means a live
+    # canonical wrapper (state b); low bit clear with non-zero high bits
+    # means preserved wrapper memory (state c); zero means never wrapped
+    # (state a). See tvm_ffi_python_helpers.h for the full lifecycle.
+    ctypedef struct PyCustomAllocHeader:
+        uintptr_t tagged_wrapper
+        TVMFFICustomAllocHeader base
+
+    const uintptr_t kPyHeaderLiveTag
+
+    PyCustomAllocHeader* TVMFFIPyHeader(void* tptr) noexcept
+    PyObject* TVMFFIPyHeaderWrapper(PyCustomAllocHeader* h) noexcept
+    bint TVMFFIPyHeaderIsLive(PyCustomAllocHeader* h) noexcept
+    int TVMFFIPyRegisterDefaultAllocator() noexcept
+    bint TVMFFIPyIsCanonical(void* chandle) noexcept
+    # Set to 0 by an atexit hook (registered in object.pxi). Read by
+    # TVMFFIPyDeleteSpace before acquiring the GIL — once Python is
+    # finalizing, state-(c) wrapper bytes are intentionally leaked
+    # rather than reclaimed. State (c)'s "next access reads same
+    # address" contract is moot after Py_Finalize starts (no Python
+    # code runs again) and the GIL acquire is unsafe.
+    void TVMFFIPyMarkPythonFinalizing() noexcept
+
     # no need to expose fields of the call context setter data structure
     ctypedef struct TVMFFIPyCallContext:
         int device_type
@@ -542,5 +581,21 @@ cdef _init_env_api():
 
 _init_env_api()
 
+
+# Register the Python custom allocator with core libtvm_ffi as the global
+# default. From this point on, every make_object<T> in the C++ side prepends
+# a PyCustomAllocHeader so we can tie a Python wrapper's lifetime to the
+# underlying C++ object. See tvm_ffi_python_helpers.h for the layout.
+CHECK_CALL(TVMFFIPyRegisterDefaultAllocator())
+
+# atexit-driven shutdown guard: flips the C-side flag that
+# TVMFFIPyDeleteSpace consults before acquiring the GIL. After this fires,
+# state-(c) wrapper bytes on chandles whose deleter hasn't yet run are
+# intentionally leaked (process is exiting; the OS reclaims). Without
+# this, a state-(c) deleter firing during/after Py_Finalize would call
+# PyGILState_Ensure on a teardown interpreter — segfault.
+import atexit as _tvm_ffi_atexit
+_tvm_ffi_atexit.register(TVMFFIPyMarkPythonFinalizing)
+
 # ensure testing is linked and we can run testcases
 TVMFFITestingDummyTarget()