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

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
@@ -279,20 +280,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_object.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
   )
@@ -332,6 +324,9 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
   install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/tvm_ffi_python_helpers.h
           DESTINATION include/
   )
+  install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/tvm_ffi_python_object.h
+          DESTINATION include/
+  )
 endif ()
 
 # ######### Install the related for normal cmake library ##########

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>
 
@@ -49,33 +50,30 @@ namespace details {
 /*!
  * \brief Allocate aligned memory.
  * \param size The size.
- * \tparam align The alignment, must be a power of 2.
+ * \param align The alignment, must be a power of 2.
  * \return The pointer to the allocated memory.
  */
-template <size_t align>
-TVM_FFI_INLINE void* AlignedAlloc(size_t size) {
-  static_assert(align != 0 && (align & (align - 1)) == 0, "align must be a power of 2");
+TVM_FFI_INLINE void* AlignedAlloc(size_t size, size_t align) {
 #ifdef _MSC_VER
   // MSVC have to use _aligned_malloc
   if (void* ptr = _aligned_malloc(size, align)) {
     return ptr;
   }
   throw std::bad_alloc();
 #else
-  if constexpr (align <= alignof(std::max_align_t)) {
+  if (align <= alignof(std::max_align_t)) {
     // malloc guarantees alignment of std::max_align_t
     if (void* ptr = std::malloc(size)) {
       return ptr;
     }
     throw std::bad_alloc();
-  } else {
-    void* ptr;
-    // for other alignments, use posix_memalign
-    if (posix_memalign(&ptr, align, size) != 0) {
-      throw std::bad_alloc();
-    }
-    return ptr;
   }
+  void* ptr;
+  // for other alignments, use posix_memalign
+  if (posix_memalign(&ptr, align, size) != 0) {
+    throw std::bad_alloc();
+  }
+  return ptr;
 #endif
 }
 
@@ -150,7 +148,12 @@ class ObjAllocatorBase {
 // Simple allocator that uses new/delete.
 class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
  private:
-  /*! \brief Guard a simple-allocator allocation until ownership is transferred to an object. */
+  /*! \brief Guard a custom-allocator allocation until ownership is transferred to an object.
+   *         If placement construction throws, free the raw block through its allocator's
+   *         ``delete_space`` (the same path the live deleter uses). ``data`` is a body pointer
+   *         offset past the ``TVMFFIObjectAllocHeader``, so a plain ``AlignedFree`` would be wrong;
+   *         the header (with ``delete_space`` wired) is set up by ``allocate`` before construction,
+   *         so this is valid even though the body was never constructed. */
   class AllocGuard {
    public:
     explicit AllocGuard(void* data) noexcept : data_(data) {}
@@ -159,7 +162,7 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
 
     ~AllocGuard() noexcept {
       if (data_ != nullptr) {
-        AlignedFree(data_);
+        ObjectUnsafe::GetObjectAllocHeaderFromPtr(data_)->delete_space(data_);
       }
     }
 
@@ -189,7 +192,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);
       AllocGuard alloc_guard(data);
       new (data) T(std::forward<Args>(args)...);
       alloc_guard.Release();
@@ -210,7 +217,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));
       }
     }
   };
@@ -238,12 +246,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);
       AllocGuard alloc_guard(data);
       new (data) ArrayType(std::forward<Args>(args)...);
       alloc_guard.Release();
@@ -264,7 +277,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

@@ -1151,6 +1151,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

@@ -59,7 +59,7 @@ dev = [
   "clang-tidy",
   "ipdb",
   "ipython",
-  "cython>=3.0",
+  "cython>=3.2.8",
   "cmake",
   "scikit-build-core",
   "tomli",
@@ -100,7 +100,7 @@ tvm-ffi-config = "tvm_ffi.config:__main__"
 tvm-ffi-stubgen = "tvm_ffi.stub.cli:__main__"
 
 [build-system]
-requires = ["scikit-build-core>=0.10.0", "cython>=3.0", "setuptools-scm"]
+requires = ["scikit-build-core>=0.10.0", "cython>=3.2.8", "setuptools-scm"]
 build-backend = "scikit_build_core.build"
 
 [tool.scikit-build]
@@ -232,14 +232,11 @@ 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 = ["cp39-*", "cp310-*", "cp311-*", "cp312-*", "cp314t-*"]
+# Per-Python-version wheels (no abi3 / limited API).
+build = ["cp39-*", "cp310-*", "cp311-*", "cp312-*", "cp313-*", "cp314t-*"]
 skip = ["*musllinux*"]
 # we only need to test on cp312
 test-skip = ["cp39-*", "cp310-*", "cp311-*"]
-# focus on testing abi3 wheel
 build-frontend = "build[uv]"
 test-command = "pytest {package}/tests/python -vvs"
 test-groups = ["test"]
@@ -288,6 +285,9 @@ allowed-unresolved-imports = [
 
 [tool.uv]
 exclude-newer = "14 days"
+# TODO(2026-07-14): drop exclude-newer-package once the rolling 14-day cutoff reaches
+# Cython 3.2.8 (uploaded 2026-06-30); until then, override so cython>=3.2.8 resolves.
+exclude-newer-package = { cython = "2026-07-01" }
 
 [tool.uv.dependency-groups]
 docs = { requires-python = ">=3.13" }

python/tvm_ffi/cython/base.pxi

@@ -360,6 +360,18 @@ def _env_get_current_stream(int device_type, int device_id):
     return <uint64_t>current_stream
 
 
+# PyObject-tying state machine (binds one Python wrapper to one C++ chandle). The C++ impl lives in
+# tvm_ffi_python_object.h; see its top-of-file banner for the design.
+cdef extern from "tvm_ffi_python_object.h":
+    int TVMFFIPyRegisterDefaultAllocator() noexcept
+    void TVMFFIPyMarkPythonFinalizing() noexcept
+
+    void TVMFFIPyCompareAndRebindPyObject(void* chandle, PyObject* expect, PyObject* new_object) noexcept
+    void TVMFFIPyTpDealloc(void** ptr_to_chandle, PyObject* wrapper) noexcept
+    void TVMFFIPyInstallTypeSlots(PyObject* type_obj) noexcept
+    object TVMFFIPyMakeRetObject(void* chandle, PyObject* cls_type)
+
+
 cdef extern from "tvm_ffi_python_helpers.h":
     # no need to expose fields of the call context setter data structure
     ctypedef struct TVMFFIPyCallContext:
@@ -542,5 +554,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/function.pxi

@@ -545,9 +545,13 @@ 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
+    # Eager-detach the canonical binding before passing the chandle by rvalue ref: the callee
+    # either moves it (src.chandle -> NULL) or leaves it, and either way ``src`` must no longer be
+    # the canonical wrapper. Recycling in both cases is handled in TVMFFIPyTpDealloc on src death.
+    TVMFFIPyCompareAndRebindPyObject(src.chandle, <PyObject*>src, NULL)
     out.type_index = kTVMFFIObjectRValueRef
-    out.v_ptr = &((<CObject>(arg.obj)).chandle)
+    out.v_ptr = &(src.chandle)
     return 0
 
 
@@ -1089,6 +1093,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))