This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 60e0614d [CORE] Staging Object Tying Feature (#696)
60e0614d is described below
commit 60e0614d6ca0ae3c15df60596131f9f8a3a2e6c2
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Aug 1 10:08:38 2026 +0800
[CORE] Staging Object Tying Feature (#696)
## Summary
Object tying remains compiled but is deliberately staged dormant:
canonical-wrapper detection returns false, while C++, Python-defined
dataclasses, and Rust objects continue to use matching legacy allocation
and deletion paths.
This staging has three goals:
- let downstream projects upgrade without making existing generated
caches depend on the new allocator contract;
- keep the compatibility change compact by changing only the
activation/allocation boundary instead of removing the object-tying
implementation;
- preserve the custom allocator API and a clean path for a later
coordinated object-tying activation.
Legacy wrapper identity and reference-count expectations are restored,
while activation-only tying tests remain as skipped coverage for the
later rollout.
## Testing
- C++: 465 enabled tests passed (2 disabled).
- Python: 2,381 passed, 78 skipped, 2 expected failures.
- Rust: 138 unit/integration tests and 3 doctests passed; formatting is
clean.
- Pre-commit: all changed-file checks passed, including ruff, ty, and
clang-format.
- Order-balanced microbenchmarks place the dormant path about 1.3-7.9%
behind active tying and 4.7-12.5% ahead of a full revert.
---
include/tvm/ffi/memory.h | 28 ++-------
python/tvm_ffi/cython/tvm_ffi_python_object.h | 24 ++++----
rust/tvm-ffi/src/object.rs | 86 +++++++++++++++++----------
src/ffi/extra/dataclass.cc | 31 ++++------
tests/python/test_dataclass_py_class.py | 11 ++--
tests/python/test_function.py | 30 +++++-----
tests/python/test_pyobject_tying.py | 4 ++
7 files changed, 103 insertions(+), 111 deletions(-)
diff --git a/include/tvm/ffi/memory.h b/include/tvm/ffi/memory.h
index c4e07841..a06a76b0 100644
--- a/include/tvm/ffi/memory.h
+++ b/include/tvm/ffi/memory.h
@@ -148,12 +148,7 @@ class ObjAllocatorBase {
// Simple allocator that uses new/delete.
class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
private:
- /*! \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. */
+ /*! \brief Guard a simple-allocator allocation until ownership is
transferred to an object. */
class AllocGuard {
public:
explicit AllocGuard(void* data) noexcept : data_(data) {}
@@ -162,7 +157,7 @@ class SimpleObjAllocator : public
ObjAllocatorBase<SimpleObjAllocator> {
~AllocGuard() noexcept {
if (data_ != nullptr) {
- ObjectUnsafe::GetObjectAllocHeaderFromPtr(data_)->delete_space(data_);
+ AlignedFree(data_);
}
}
@@ -192,11 +187,7 @@ 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.
- 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);
+ void* data = AlignedAlloc(sizeof(T), alignof(T));
AllocGuard alloc_guard(data);
new (data) T(std::forward<Args>(args)...);
alloc_guard.Release();
@@ -217,8 +208,7 @@ class SimpleObjAllocator : public
ObjAllocatorBase<SimpleObjAllocator> {
tptr->T::~T();
}
if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
- ObjectUnsafe::GetObjectAllocHeaderFromPtr(static_cast<void*>(tptr))
- ->delete_space(static_cast<void*>(tptr));
+ AlignedFree(static_cast<void*>(tptr));
}
}
};
@@ -246,17 +236,12 @@ 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);
- TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
- void* data =
- alloc->allocate(aligned_size, align, ArrayType::RuntimeTypeIndex(),
alloc->context);
+ void* data = AlignedAlloc(aligned_size, align);
AllocGuard alloc_guard(data);
new (data) ArrayType(std::forward<Args>(args)...);
alloc_guard.Release();
@@ -277,8 +262,7 @@ class SimpleObjAllocator : public
ObjAllocatorBase<SimpleObjAllocator> {
tptr->ArrayType::~ArrayType();
}
if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
- ObjectUnsafe::GetObjectAllocHeaderFromPtr(static_cast<void*>(tptr))
- ->delete_space(static_cast<void*>(tptr));
+ AlignedFree(static_cast<void*>(tptr));
}
}
};
diff --git a/python/tvm_ffi/cython/tvm_ffi_python_object.h
b/python/tvm_ffi/cython/tvm_ffi_python_object.h
index ed8389ed..1672c784 100644
--- a/python/tvm_ffi/cython/tvm_ffi_python_object.h
+++ b/python/tvm_ffi/cython/tvm_ffi_python_object.h
@@ -660,24 +660,20 @@ inline void TVMFFIPyMarkPythonFinalizing() noexcept {
}
/*!
- * \brief True iff ``chandle`` was allocated through the Python custom
- * allocator (full ``PyCustomAllocHeader`` ahead of it). False for
- * allocations that came through libtvm_ffi's builtin default
- * (only the base ``TVMFFIObjectAllocHeader``).
+ * \brief Whether ``chandle`` participates in Python object tying.
*
- * Detection is by comparing ``base.delete_space`` against
- * ``TVMFFIPyDeleteSpace``: each frontend recognizes its own deleter
- * pointer, so multiple frontends can coexist without a flag bit on
- * ``TVMFFIObject``.
+ * Object tying is dormant during the compatibility rollout: public object
+ * constructors still use the legacy no-prefix allocation layout, so probing
+ * memory before an arbitrary object body would be invalid. Keep the complete
+ * tying implementation compiled, but route every object through the ordinary
+ * fresh-wrapper path until allocation is activated in a later release.
*
- * \param chandle The FFI object handle to test (NULL yields false).
- * \return True iff ``chandle`` carries a ``PyCustomAllocHeader`` (our
deleter).
+ * \param chandle The FFI object handle to test.
+ * \return False while object tying remains dormant.
*/
TVM_FFI_INLINE bool TVMFFIPyIsCanonical(void* chandle) {
- if (chandle == nullptr) return false;
- TVMFFIObjectAllocHeader* base = reinterpret_cast<TVMFFIObjectAllocHeader*>(
- static_cast<char*>(chandle) - sizeof(TVMFFIObjectAllocHeader));
- return base->delete_space == &TVMFFIPyDeleteSpace;
+ (void)chandle;
+ return false;
}
//---------------------------------------------------------------
diff --git a/rust/tvm-ffi/src/object.rs b/rust/tvm-ffi/src/object.rs
index 712aa35e..7adc439b 100644
--- a/rust/tvm-ffi/src/object.rs
+++ b/rust/tvm-ffi/src/object.rs
@@ -23,10 +23,7 @@ use crate::derive::ObjectRef;
use crate::type_traits::AnyCompatible;
pub use tvm_ffi_sys::TVMFFITypeIndex as TypeIndex;
/// Object related ABI handling
-use tvm_ffi_sys::{
- TVMFFIAny, TVMFFIGetCustomAllocator, TVMFFIGetTypeInfo, TVMFFIObject,
- COMBINED_REF_COUNT_BOTH_ONE,
-};
+use tvm_ffi_sys::{TVMFFIAny, TVMFFIGetTypeInfo, TVMFFIObject,
COMBINED_REF_COUNT_BOTH_ONE};
/// Object type is by default the TVMFFIObject
#[repr(C)]
@@ -239,7 +236,7 @@ pub mod unsafe_ {
use std::ffi::c_void;
use std::sync::atomic::{fence, Ordering};
- use tvm_ffi_sys::{TVMFFIObject, TVMFFIObjectAllocHeader};
+ use tvm_ffi_sys::TVMFFIObject;
use tvm_ffi_sys::TVMFFIObjectDeleterFlagBitMask::{
kTVMFFIObjectDeleterFlagBitMaskBoth,
kTVMFFIObjectDeleterFlagBitMaskStrong,
kTVMFFIObjectDeleterFlagBitMaskWeak,
@@ -316,8 +313,7 @@ pub mod unsafe_ {
(obj.combined_ref_count.load(Ordering::Relaxed) >> 32) as usize
}
- /// Generic object deleter for objects allocated through the registered
- /// `TVMFFICustomAllocator`.
+ /// Generic object deleter for objects allocated through Rust's global
allocator.
pub(crate) unsafe extern "C" fn object_deleter_for_new<T>(ptr: *mut
c_void, flags: i32)
where
T: super::ObjectCore,
@@ -327,13 +323,42 @@ pub mod unsafe_ {
std::ptr::drop_in_place(obj);
}
if flags & kTVMFFIObjectDeleterFlagBitMaskWeak as i32 != 0 {
- let header_ptr = (ptr as *mut u8)
- .sub(std::mem::size_of::<TVMFFIObjectAllocHeader>())
- as *mut TVMFFIObjectAllocHeader;
- let delete_space = (*header_ptr)
- .delete_space
- .expect("TVMFFIObjectAllocHeader::delete_space must be set by
the allocator");
- delete_space(ptr);
+ std::alloc::dealloc(ptr as *mut u8,
std::alloc::Layout::new::<T>());
+ }
+ }
+
+ pub(crate) unsafe extern "C" fn object_deleter_for_new_with_extra_items<T,
U>(
+ ptr: *mut c_void,
+ flags: i32,
+ ) where
+ T: super::ObjectCoreWithExtraItems<ExtraItem = U>,
+ {
+ let obj = ptr as *mut T;
+ if flags == kTVMFFIObjectDeleterFlagBitMaskBoth as i32 {
+ let extra_items_count = T::extra_items_count(&(*obj));
+ std::ptr::drop_in_place(obj);
+ let layout = std::alloc::Layout::from_size_align(
+ std::mem::size_of::<T>() + extra_items_count *
std::mem::size_of::<U>(),
+ std::mem::align_of::<T>(),
+ )
+ .unwrap();
+ std::alloc::dealloc(ptr as *mut u8, layout);
+ } else {
+ assert_eq!(std::mem::size_of::<T>() % std::mem::size_of::<u64>(),
0);
+ if flags & kTVMFFIObjectDeleterFlagBitMaskStrong as i32 != 0 {
+ let extra_items_count = T::extra_items_count(&(*obj));
+ std::ptr::drop_in_place(obj);
+ std::ptr::write(obj as *mut u64, extra_items_count as u64);
+ }
+ if flags & kTVMFFIObjectDeleterFlagBitMaskWeak as i32 != 0 {
+ let extra_items_count = std::ptr::read(obj as *mut u64) as
usize;
+ let layout = std::alloc::Layout::from_size_align(
+ std::mem::size_of::<T>() + extra_items_count *
std::mem::size_of::<U>(),
+ std::mem::align_of::<T>(),
+ )
+ .unwrap();
+ std::alloc::dealloc(ptr as *mut u8, layout);
+ }
}
}
}
@@ -367,23 +392,14 @@ unsafe impl ObjectCore for Object {
// ObjectArc
//---------------------
-/// Allocate via the process-wide `TVMFFICustomAllocator`.
-unsafe fn allocate_via_registry(size: usize, alignment: usize, type_index:
i32) -> *mut u8 {
- let alloc = TVMFFIGetCustomAllocator();
- let allocate_fn = (*alloc)
- .allocate
- .expect("TVMFFICustomAllocator::allocate must be set");
- allocate_fn(size, alignment, type_index, (*alloc).context) as *mut u8
-}
-
impl<T: ObjectCore> ObjectArc<T> {
pub fn new(data: T) -> Self {
unsafe {
- let raw_data_ptr = allocate_via_registry(
- std::mem::size_of::<T>(),
- std::mem::align_of::<T>(),
- T::type_index(),
- );
+ let layout = std::alloc::Layout::new::<T>();
+ let raw_data_ptr = std::alloc::alloc(layout);
+ if raw_data_ptr.is_null() {
+ std::alloc::handle_alloc_error(layout);
+ }
let ptr = raw_data_ptr as *mut T;
std::ptr::write(ptr, data);
// now override the header directly
@@ -413,9 +429,15 @@ impl<T: ObjectCore> ObjectArc<T> {
assert_eq!(std::mem::align_of::<T>() % std::mem::align_of::<U>(),
0);
assert_eq!(std::mem::size_of::<T>() % std::mem::align_of::<U>(),
0);
let extra_items_count = T::extra_items_count(&data);
- let total_size = std::mem::size_of::<T>() + extra_items_count *
std::mem::size_of::<U>();
- let raw_data_ptr =
- allocate_via_registry(total_size, std::mem::align_of::<T>(),
T::type_index());
+ let layout = std::alloc::Layout::from_size_align(
+ std::mem::size_of::<T>() + extra_items_count *
std::mem::size_of::<U>(),
+ std::mem::align_of::<T>(),
+ )
+ .unwrap();
+ let raw_data_ptr = std::alloc::alloc(layout);
+ if raw_data_ptr.is_null() {
+ std::alloc::handle_alloc_error(layout);
+ }
let ptr = raw_data_ptr as *mut T;
std::ptr::write(ptr, data);
// now override the header directly
@@ -425,7 +447,7 @@ impl<T: ObjectCore> ObjectArc<T> {
combined_ref_count:
AtomicU64::new(COMBINED_REF_COUNT_BOTH_ONE),
type_index: T::type_index(),
__padding: 0,
- deleter: Some(unsafe_::object_deleter_for_new::<T>),
+ deleter:
Some(unsafe_::object_deleter_for_new_with_extra_items::<T, U>),
},
);
// move into the object arc ptr
diff --git a/src/ffi/extra/dataclass.cc b/src/ffi/extra/dataclass.cc
index 4b6452d7..73a05a48 100644
--- a/src/ffi/extra/dataclass.cc
+++ b/src/ffi/extra/dataclass.cc
@@ -1742,9 +1742,7 @@ PyClassFieldStorageKind GetPyClassFieldStorageKind(const
TVMFFIFieldInfo* finfo)
*
* For the "strong" phase, iterates all reflected fields and destructs
* Any/ObjectRef values in-place (to release references). For the "weak"
- * phase, forwards to the prepended ``TVMFFIObjectAllocHeader``'s
- * ``delete_space`` (every Object carries one — libtvm_ffi's builtin
- * default allocator ensures it).
+ * phase, frees the underlying calloc'd memory.
*/
void PyClassDeleter(void* self_void, int flags) {
TVMFFIObject* self = static_cast<TVMFFIObject*>(self_void);
@@ -1768,7 +1766,7 @@ void PyClassDeleter(void* self_void, int flags) {
});
}
if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
-
details::ObjectUnsafe::GetObjectAllocHeaderFromPtr(self_void)->delete_space(self_void);
+ std::free(self_void);
}
}
@@ -2118,16 +2116,11 @@ 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, type_info]()
-> ObjectRef {
- // Route through the custom-allocator registry so the prepended
- // TVMFFIObjectAllocHeader is in place for PyClassDeleter's Weak
- // branch. Then memset to zero-init the payload (mirrors the original
- // calloc-based path); ConstructPyClassFields below placement-news the
- // non-trivial fields while the memset covers the POD ones.
- 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);
+ 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);
+ }
TVMFFIObject* ffi_obj = reinterpret_cast<TVMFFIObject*>(obj_ptr);
ffi_obj->type_index = type_index;
ffi_obj->combined_ref_count = details::kCombinedRefCountBothOne;
@@ -2142,12 +2135,10 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index,
int32_t total_size) {
// Step 3. Register `__ffi_shallow_copy__`
Function copy_fn =
Function::FromTyped([type_index, total_size, type_info](const Object*
src) -> ObjectRef {
- // Allocator + memset mirror RegisterFFINew above.
- 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);
+ void* obj_ptr = std::calloc(1, static_cast<size_t>(total_size));
+ if (!obj_ptr) {
+ TVM_FFI_THROW(RuntimeError) << "Failed to allocate for shallow copy";
+ }
TVMFFIObject* ffi_obj = reinterpret_cast<TVMFFIObject*>(obj_ptr);
ffi_obj->type_index = type_index;
ffi_obj->combined_ref_count = details::kCombinedRefCountBothOne;
diff --git a/tests/python/test_dataclass_py_class.py
b/tests/python/test_dataclass_py_class.py
index c001bbcb..b4b8986b 100644
--- a/tests/python/test_dataclass_py_class.py
+++ b/tests/python/test_dataclass_py_class.py
@@ -3689,23 +3689,20 @@ class TestNativeParentInheritance:
stored = holder.value
assert stored is not None
assert stored.same_as(target)
- # Under PyObject-tying the field accessor returns the canonical
wrapper,
- # so ``stored`` aliases ``target`` (no fresh wrapper, no extra C++
ref) and
- # the use count stays 2 rather than climbing to 3.
- assert stored is target
- assert use_count(target) == 2
+ assert stored is not target
+ assert use_count(target) == 3
del stored
gc.collect()
assert use_count(target) == 2
with pytest.raises(TypeError, match=r"testing\.TestObjectBase"):
holder.value = None # ty: ignore[invalid-assignment]
- assert holder.value is target
+ assert holder.value.same_as(target)
assert use_count(target) == 2
assert holder.optional_value is None
holder.optional_value = target
- assert holder.optional_value is target
+ assert holder.optional_value.same_as(target)
assert use_count(target) == 3
holder.optional_value = None
assert holder.optional_value is None
diff --git a/tests/python/test_function.py b/tests/python/test_function.py
index 908e2e21..2dd3ad86 100644
--- a/tests/python/test_function.py
+++ b/tests/python/test_function.py
@@ -192,41 +192,39 @@ def test_global_func() -> None:
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:
- # ``gc.collect()`` ensures Python destructors have run so
- # use_count reflects only live wrappers.
+ # 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()
assert expected_count == use_count(x)
return x._move()
f = tvm_ffi.convert(callback)
- def check_caller_move() -> None:
- # Caller passes ``x._move()``: callback receives a fresh canonical
- # wrapper for the moved-in chandle.
+ def check0() -> None:
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 check_callback_move() -> None:
- # Callback returns ``x._move()``: caller sees a fresh canonical
- # wrapper, distinct from the now-empty ``x``.
+ def check1() -> None:
x = tvm_ffi.convert([1, 2])
assert use_count(x) == 1
- y = f(x, 1)
- assert y is not x
+ y = f(x, 2)
+ f(x._move(), 2)
assert x.__ctypes_handle__().value is None
assert y.__ctypes_handle__().value is not None
- check_caller_move()
- check_callback_move()
+ check0()
+ check1()
def test_echo_with_opaque_object() -> None:
diff --git a/tests/python/test_pyobject_tying.py
b/tests/python/test_pyobject_tying.py
index 1bd7abb2..ae3ec1e6 100644
--- a/tests/python/test_pyobject_tying.py
+++ b/tests/python/test_pyobject_tying.py
@@ -48,6 +48,10 @@ import tvm_ffi
import tvm_ffi.testing
from tvm_ffi import dataclasses as dc
+pytestmark = pytest.mark.skip(
+ reason="PyObject tying is compiled but dormant until allocator activation"
+)
+
# ---------------------------------------------------------------------------
# Test type registration
# ---------------------------------------------------------------------------
