junrushao commented on code in PR #593:
URL: https://github.com/apache/tvm-ffi/pull/593#discussion_r3575581097
##########
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);
Review Comment:
```suggestion
/*!
* \brief Reclaim storage allocated for a TVMFFIObject.
* \param ptr Pointer to the object body returned by the allocator.
*
* When used, ``TVMFFIObject::deleter`` invokes this callback after the
* object's weak lifetime ends. The callback is responsible for eventually
* reclaiming the complete allocator-owned allocation, including any
private
* prefix, and cleaning up associated frontend state.
*
* \note ``ptr`` points to the object body, not to the preceding
* ``TVMFFIObjectAllocHeader``. The object's C++ lifetime may already
* have ended, so the callback must not access its fields.
* \note This callback may be NULL if the object deleter reclaims the
* allocation directly.
*/
void (*delete_space)(void* ptr);
```
##########
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);
+
+/*!
Review Comment:
Please document if this API is thread-safe
##########
python/tvm_ffi/cython/base.pxi:
##########
@@ -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)
Review Comment:
nit: Move initialization time logics into `core.pyx`
##########
pyproject.toml:
##########
@@ -236,14 +236,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-*"]
Review Comment:
Why are we building cp313 though? I thought we will be using ABI3 so it can
be skipped?
##########
python/tvm_ffi/cython/pycallback.pxi:
##########
@@ -74,7 +74,13 @@ cdef int TVMFFIPyCallbackArgSetterObject_(
const TVMFFIAny* arg,
PyObject** out
) except -1:
- """Callback arg setter for generic static object types (type_index >=
kTVMFFIStaticObjectBegin)."""
+ """Callback arg setter for generic static object types (type_index >=
kTVMFFIStaticObjectBegin).
+
+ Funnels through ``make_ret_object`` so the callback receives the
+ canonical wrapper for the chandle. When the caller already has a
+ wrapper for this chandle, the callback's arg is the same Python
+ object — universal cache-on aliasing.
Review Comment:
```suggestion
"""Convert a generic FFI object argument for a Python callback.
The input contains a borrowed chandle, so this function increments the
chandle's reference count before calling ``make_ret_object(arg[0])``.
``make_ret_object`` consumes that reference and returns the Python value
passed to the callback.
For ``CObject``-based types that support wrapper tying,
``make_ret_object`` returns the existing canonical wrapper, revives its
previous allocation when eligible, or creates and binds a new wrapper.
If a live canonical wrapper already exists, the callback receives that
exact Python object. Chandles without Python allocator metadata are
wrapped normally but cannot reuse a canonical wrapper.
``PyNativeObject`` subclasses such as ``String``, ``Bytes``, and
``Shape`` are instead converted to their Python-native representations
and do not participate in wrapper tying. ``Tensor`` and
``OpaquePyObject`` are handled by dedicated callback argument setters
and therefore do not use this function.
When ``api`` is provided for a ``CContainerBase``, this function attaches
the exchange API to support lazy DLPack conversion.
"""
```
##########
src/ffi/extra/dataclass.cc:
##########
@@ -2041,16 +2044,20 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index,
int32_t total_size) {
RegisterFFIInit(type_index);
// Step 2. Register `__ffi_new__`
Function new_fn = Function::FromTyped([type_index, total_size]() ->
ObjectRef {
- void* obj_ptr = std::calloc(1, static_cast<size_t>(total_size));
- if (!obj_ptr) {
- TVM_FFI_THROW(RuntimeError) << "Failed to allocate " << total_size << "
bytes for type "
- << TypeIndexToTypeKey(type_index);
- }
+ // Route through the custom-allocator registry so the prepended
+ // TVMFFIObjectAllocHeader is in place for PyClassDeleter's Weak
+ // branch. Then memset to zero-init the T payload (mirrors the
+ // original calloc-based path).
+ size_t alloc_size = static_cast<size_t>(total_size);
+ TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
+ void* obj_ptr =
+ alloc->allocate(alloc_size, alignof(::std::max_align_t), type_index,
alloc->context);
+ std::memset(obj_ptr, 0, alloc_size);
Review Comment:
It means the prepended `TVMFFIObjectAllocHeader` is properly set by
`alloc->allocate(...)` right? If so, add this info to docstring of `allocate`
API
##########
CMakeLists.txt:
##########
@@ -294,20 +295,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")
Review Comment:
you may want to resolve the merge conflict here
##########
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;
Review Comment:
```suggestion
/*!
* \brief Allocator interface for TVM FFI object storage.
*
* A frontend may register this allocator with ``TVMFFISetCustomAllocator``
* to control allocations of subsequently created ``TVMFFIObject`` instances.
* This allows the frontend to prepend private metadata and coordinate object
* storage reclamation with language-specific state.
*
* Every returned object body must be immediately preceded by an initialized
* ``TVMFFIObjectAllocHeader``. An allocator may place additional private
* metadata before that header.
*
* \note The registered allocator and its context must remain valid for the
* lifetime of the process.
* \note The allocation callback may be invoked concurrently.
*/
typedef struct {
/*!
* \brief Allocate uninitialized storage for a TVM FFI object body.
* \param size Size of the object body in bytes, excluding the preceding
* ``TVMFFIObjectAllocHeader`` and any allocator-private
metadata.
* \param alignment Required alignment of the returned object body.
* \param type_index Runtime type index of the object being allocated.
* \param context The allocator-defined context stored in this entry.
* \return An aligned pointer to the object body, or NULL on failure.
*
* The returned pointer must be immediately preceded by an initialized
* ``TVMFFIObjectAllocHeader`` whose ``delete_space`` callback is
responsible
* for eventually reclaiming the complete allocator-owned allocation.
*
* On failure, the callback must report the error through
* ``TVMFFIErrorSetRaised`` before returning NULL.
*/
void* (*allocate)(size_t size, size_t alignment, int32_t type_index, void*
context);
/*!
* \brief Opaque allocator-defined context passed unchanged to
``allocate``.
*
* The TVM FFI runtime neither interprets nor owns this pointer. It may be
* NULL; otherwise, the referenced state must remain valid for the lifetime
* of the registered allocator.
*/
void* context;
} TVMFFICustomAllocator;
```
##########
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.
+ */
Review Comment:
```suggestion
/*!
* \brief Mandatory header placed immediately before each TVMFFIObject body.
*
* A ``TVMFFIObject::deleter`` may invoke this header's ``delete_space``
* callback when the object's storage lifetime ends. Frontend custom
allocators
* can use this mechanism to coordinate cleanup of language-specific state,
* such as a cached PyObject wrapper, with reclamation of the TVM FFI
allocation.
*
* The callback may be NULL if the object deleter releases the storage
directly.
*/
```
##########
include/tvm/ffi/object.h:
##########
@@ -1152,6 +1152,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));
Review Comment:
Does it always guarantee that there's no compiler-inserted padding before an
object, e.g., vptr?
##########
pyproject.toml:
##########
@@ -292,6 +289,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" }
Review Comment:
remove this line because 14 days have passed since 07/01
##########
python/tvm_ffi/cython/function.pxi:
##########
@@ -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.
Review Comment:
nit: remove this line?
##########
tests/python/test_function.py:
##########
@@ -177,40 +177,46 @@ def echo(x: Any) -> Any:
assert tvm_ffi.get_global_func("mytest.echo", allow_missing=True) is None
[email protected](
+ hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled(),
+ reason="PyObject-tying (wrapper aliasing) is disabled on free-threaded
Python",
Review Comment:
Is this still disabled? Can we re-enable this unittest?
##########
tests/python/test_function.py:
##########
@@ -177,40 +177,46 @@ def echo(x: Any) -> Any:
assert tvm_ffi.get_global_func("mytest.echo", allow_missing=True) is None
[email protected](
+ hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled(),
+ reason="PyObject-tying (wrapper aliasing) is disabled on free-threaded
Python",
+)
def test_rvalue_ref() -> None:
+ # Under universal cache-on the callback's arg aliases the caller's
+ # wrapper, so use_count inside is 1 (one wrapper, one chandle ref).
+ # ``_move()`` on either side detaches that wrapper's binding before
+ # the C++ AnyViewToOwnedAny transfer nulls the source chandle.
use_count = tvm_ffi.get_global_func("testing.object_use_count")
def callback(x: Any, expected_count: int) -> Any:
- # The use count of TVM FFI objects is decremented as part of
- # `ObjectRef.__del__`, which runs when the Python object is
- # destructed. However, Python object destruction is not
- # deterministic, and even CPython's reference-counting is
- # considered an implementation detail. Therefore, to ensure
- # correct results from this test, `gc.collect()` must be
- # explicitly called.
+ # ``gc.collect()`` ensures Python destructors have run so
+ # use_count reflects only live wrappers.
gc.collect()
assert expected_count == use_count(x)
return x._move()
f = tvm_ffi.convert(callback)
- def check0() -> None:
+ def check_caller_move() -> None:
+ # Caller passes ``x._move()``: callback receives a fresh canonical
+ # wrapper for the moved-in chandle.
x = tvm_ffi.convert([1, 2])
assert use_count(x) == 1
- f(x, 2)
f(x._move(), 1)
assert x.__ctypes_handle__().value is None
- def check1() -> None:
+ def check_callback_move() -> None:
+ # Callback returns ``x._move()``: caller sees a fresh canonical
+ # wrapper, distinct from the now-empty ``x``.
x = tvm_ffi.convert([1, 2])
assert use_count(x) == 1
- y = f(x, 2)
- f(x._move(), 2)
+ y = f(x, 1)
+ assert y is not x
assert x.__ctypes_handle__().value is None
assert y.__ctypes_handle__().value is not None
- check0()
- check1()
+ check_caller_move()
Review Comment:
not suggesting any change - did we accidentally swap check0 and check1 :p
##########
python/tvm_ffi/cython/tensor.pxi:
##########
@@ -497,7 +497,11 @@ cdef inline object make_tensor_from_chandle(
# call the deleter to free the memory since we will continue
to use the chandle
dlpack.deleter(dlpack)
pass
- # default return the tensor
+ # default return the tensor.
Review Comment:
Trying to understand this - it means id-consistency is not guaranteed for
tensors?
##########
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);
Review Comment:
BTW, do we guarantee zero-initialized memory?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
