This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch python-move-if-unique in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
commit 195b6189557a6b380a2eb005b996eeb32fe243a7 Author: tqchen <[email protected]> AuthorDate: Fri Jul 17 11:44:28 2026 +0000 [FEAT][Python] Add move-if-unique object helper --- python/tvm_ffi/core.pyi | 1 + python/tvm_ffi/cython/base.pxi | 1 + python/tvm_ffi/cython/object.pxi | 21 +++++++++++++++++++++ python/tvm_ffi/cython/tvm_ffi_python_object.h | 15 +++++++++++++++ tests/python/test_function.py | 23 +++++++++++++++++++++++ 5 files changed, 61 insertions(+) diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi index 525d7a40..57492d0c 100644 --- a/python/tvm_ffi/core.pyi +++ b/python/tvm_ffi/core.pyi @@ -56,6 +56,7 @@ class CObject: def same_as(self, other: Any) -> bool: ... def is_(self, other: Any) -> bool: ... def _move(self) -> ObjectRValueRef: ... + def _move_if_unique(self) -> ObjectRValueRef | CObject: ... def __move_handle_from__(self, other: CObject) -> None: ... class CContainerBase(CObject): ... diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi index 56c4d3d1..6febf961 100644 --- a/python/tvm_ffi/cython/base.pxi +++ b/python/tvm_ffi/cython/base.pxi @@ -371,6 +371,7 @@ cdef extern from "tvm_ffi_python_object.h": void TVMFFIPyTpDealloc(void** ptr_to_chandle, PyObject* wrapper) noexcept void TVMFFIPyInstallTypeSlots(PyObject* type_obj) noexcept object TVMFFIPyMakeRetObject(void* chandle, PyObject* cls_type) + bint TVMFFIPyWrapperIsUniqueForMove(PyObject* wrapper) noexcept cdef extern from "tvm_ffi_python_helpers.h": diff --git a/python/tvm_ffi/cython/object.pxi b/python/tvm_ffi/cython/object.pxi index 30e3ea90..f792e319 100644 --- a/python/tvm_ffi/cython/object.pxi +++ b/python/tvm_ffi/cython/object.pxi @@ -335,6 +335,27 @@ class Object(CObject, metaclass=_ObjectSlotsMeta): """ return ObjectRValueRef(self) + def _move_if_unique(self): + """Move this object only when its Python wrapper is uniquely referenced. + + Returns an rvalue reference when this wrapper has exactly one Python + reference. Otherwise, returns the object itself so FFI calls receive + it as an lvalue and retain the original wrapper. + + Returns + ------- + ObjectRValueRef or Object + An rvalue reference for a unique wrapper, or ``self`` when shared. + + Notes + ----- + Free-threaded Python conservatively returns ``self`` because its + reference count does not establish exclusive access. + """ + if TVMFFIPyWrapperIsUniqueForMove(<PyObject*>self): + return ObjectRValueRef(self) + return self + def __move_handle_from__(self, other: CObject) -> None: """Steal the FFI handle from ``other``. diff --git a/python/tvm_ffi/cython/tvm_ffi_python_object.h b/python/tvm_ffi/cython/tvm_ffi_python_object.h index ed8389ed..23d5048d 100644 --- a/python/tvm_ffi/cython/tvm_ffi_python_object.h +++ b/python/tvm_ffi/cython/tvm_ffi_python_object.h @@ -48,6 +48,21 @@ #define Py_TPFLAGS_MANAGED_DICT 0 #endif +/*! \brief Whether a Python-visible ``_move_if_unique`` call has one external wrapper owner. + * + * The bound-method call and Cython argument wrapper each temporarily own a reference, so a + * classic-GIL build observes three references when the caller has exactly one. Free-threaded + * builds conservatively return false: ``Py_REFCNT`` does not establish exclusive access there, + * and an lvalue preserves correctness at the cost of skipping the move optimization. + */ +TVM_FFI_INLINE bool TVMFFIPyWrapperIsUniqueForMove(PyObject* wrapper) { +#ifdef Py_GIL_DISABLED + return false; +#else + return Py_REFCNT(wrapper) == 3; +#endif +} + #include <atomic> #include <cassert> #include <cstring> diff --git a/tests/python/test_function.py b/tests/python/test_function.py index 908e2e21..c905dce8 100644 --- a/tests/python/test_function.py +++ b/tests/python/test_function.py @@ -229,6 +229,29 @@ def test_rvalue_ref() -> None: check_callback_move() +def test_move_if_unique() -> None: + discard = tvm_ffi.convert(lambda _: None) + + shared = tvm_ffi.convert([1, 2]) + alias = shared + assert shared._move_if_unique() is shared + discard(shared._move_if_unique()) + assert shared.__ctypes_handle__().value is not None + assert alias is shared + + del alias + if hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled(): + candidate = shared._move_if_unique() + assert candidate is shared + discard(candidate) + assert shared.__ctypes_handle__().value is not None + else: + candidate = shared._move_if_unique() + assert isinstance(candidate, tvm_ffi.core.ObjectRValueRef) + discard(candidate) + assert shared.__ctypes_handle__().value is None + + def test_echo_with_opaque_object() -> None: class MyObject: def __init__(self, value: Any) -> None:
