This is an automated email from the ASF dual-hosted git repository.
junrushao 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 b64b46f [BUILD] Enable free threaded python build (#98)
b64b46f is described below
commit b64b46f32e845b650850d73a5828a2d3f07d3406
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 10 13:52:13 2025 -0400
[BUILD] Enable free threaded python build (#98)
This PR makes the cmake and tests to be compatible with free threaded
python.
Also fixes an issue where OpaqueObject is not recognized as Object type,
which is covered by the ffi ref count test.
---
.github/workflows/ci_test.yml | 4 ++--
CMakeLists.txt | 16 ++++++++++---
pyproject.toml | 6 +++--
python/tvm_ffi/cython/base.pxi | 2 ++
python/tvm_ffi/cython/core.pyx | 2 ++
python/tvm_ffi/cython/function.pxi | 19 +++++-----------
python/tvm_ffi/cython/tvm_ffi_python_helpers.h | 31 ++++++++++++++++++++++++++
src/ffi/object.cc | 8 ++++++-
tests/python/test_object.py | 13 ++++++-----
9 files changed, 74 insertions(+), 27 deletions(-)
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
index 2ea884c..e5361f3 100644
--- a/.github/workflows/ci_test.yml
+++ b/.github/workflows/ci_test.yml
@@ -79,10 +79,10 @@ jobs:
fail-fast: true
matrix:
include:
- - {os: ubuntu-latest, arch: x86_64, python_version: '3.13'}
+ - {os: ubuntu-latest, arch: x86_64, python_version: '3.14t'}
- {os: ubuntu-24.04-arm, arch: aarch64, python_version: '3.9'}
- {os: windows-latest, arch: AMD64, python_version: '3.12'}
- - {os: macos-14, arch: arm64, python_version: '3.10'}
+ - {os: macos-14, arch: arm64, python_version: '3.13'}
steps:
- uses: actions/checkout@v4
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 618588d..be2a84f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -192,8 +192,17 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/object.pxi
${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/string.pxi
)
- # set working directory to source so we can see the exact file name in
backtrace relatived to the
- # project source root
+ # Run a Python script to check for free-threaded build
+ execute_process(
+ COMMAND ${Python_EXECUTABLE} -c
+ "import sysconfig;
print(sysconfig.get_config_var('Py_GIL_DISABLED') == 1)"
+ OUTPUT_VARIABLE PYTHON_IS_FREE_THREADED
+ OUTPUT_STRIP_TRAILING_WHITESPACE
+ )
+ if (PYTHON_IS_FREE_THREADED)
+ message(STATUS "Free-threaded Python detected.")
+ endif ()
+
add_custom_command(
OUTPUT ${_core_cpp}
COMMAND ${Python_EXECUTABLE} -m cython --cplus ${_core_pyx} -o ${_core_cpp}
@@ -202,7 +211,8 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
DEPENDS ${_cython_sources}
VERBATIM
)
- if (Python_VERSION VERSION_GREATER_EQUAL "3.12")
+
+ 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)
set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
diff --git a/pyproject.toml b/pyproject.toml
index acf62e2..25a8137 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,8 @@ GitHub = "https://github.com/apache/tvm-ffi";
# setup tools is needed by torch jit for best perf
torch = ["torch", "setuptools", "ninja"]
cpp = ["ninja"]
-test = ["pytest", "numpy", "torch", "ninja"]
+# note pytorch does not yet ship with 3.14t
+test = ["pytest", "numpy", "ninja", "torch; python_version < '3.14'"]
[dependency-groups]
docs = [
@@ -197,7 +198,8 @@ build-verbosity = 1
# only build up to cp312, cp312
# will be abi3 and can be used in future versions
-build = ["cp39-*", "cp310-*", "cp311-*", "cp312-*"]
+# ship 314t threaded nogil version
+build = ["cp39-*", "cp310-*", "cp311-*", "cp312-*", "cp314t-*"]
skip = ["*musllinux*"]
# we only need to test on cp312
test-skip = ["cp39-*", "cp310-*", "cp311-*"]
diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi
index a8b4212..221bb04 100644
--- a/python/tvm_ffi/cython/base.pxi
+++ b/python/tvm_ffi/cython/base.pxi
@@ -339,6 +339,8 @@ cdef extern from "tvm_ffi_python_helpers.h":
int TVMFFIPyArgSetterInt_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*,
PyObject* arg, TVMFFIAny* out) except -1
int TVMFFIPyArgSetterBool_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*,
PyObject* arg, TVMFFIAny* out) except -1
int TVMFFIPyArgSetterNone_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*,
PyObject* arg, TVMFFIAny* out) except -1
+ # deleter for python objects
+ void TVMFFIPyObjectDeleter(void* py_obj) noexcept nogil
cdef class ByteArrayArg:
diff --git a/python/tvm_ffi/cython/core.pyx b/python/tvm_ffi/cython/core.pyx
index ca3a0ce..b11a27c 100644
--- a/python/tvm_ffi/cython/core.pyx
+++ b/python/tvm_ffi/cython/core.pyx
@@ -1,3 +1,5 @@
+# cython: freethreading_compatible = True
+# cython: language_level=3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
diff --git a/python/tvm_ffi/cython/function.pxi
b/python/tvm_ffi/cython/function.pxi
index 2fa75fb..4c7461c 100644
--- a/python/tvm_ffi/cython/function.pxi
+++ b/python/tvm_ffi/cython/function.pxi
@@ -709,12 +709,6 @@ def _get_global_func(name, allow_missing):
raise ValueError("Cannot find global function %s" % name)
-# handle callbacks
-cdef void tvm_ffi_pyobject_deleter(void* fhandle) noexcept with gil:
- local_pyobject = <object>(fhandle)
- Py_DECREF(local_pyobject)
-
-
cdef int tvm_ffi_callback(void* context,
const TVMFFIAny* packed_args,
int32_t num_args,
@@ -722,8 +716,9 @@ cdef int tvm_ffi_callback(void* context,
cdef list pyargs
cdef TVMFFIAny temp_result
cdef int c_api_ret_code
- local_pyfunc = <object>(context)
+ cdef object local_pyfunc = <object>(context)
pyargs = []
+
for i in range(num_args):
CHECK_CALL(TVMFFIAnyViewToOwnedAny(&packed_args[i], &temp_result))
pyargs.append(make_ret(temp_result))
@@ -736,11 +731,7 @@ cdef int tvm_ffi_callback(void* context,
result,
&c_api_ret_code
)
- if c_api_ret_code == 0:
- return 0
- elif c_api_ret_code == -2:
- raise_existing_error()
- return -1
+ return c_api_ret_code
except Exception as err:
set_last_ffi_error(err)
return -1
@@ -754,7 +745,7 @@ cdef inline int _convert_to_ffi_func_handle(
CHECK_CALL(TVMFFIFunctionCreate(
<void*>(pyfunc),
tvm_ffi_callback,
- tvm_ffi_pyobject_deleter,
+ TVMFFIPyObjectDeleter,
out_handle))
return 0
@@ -776,7 +767,7 @@ cdef inline int _convert_to_opaque_object_handle(
CHECK_CALL(TVMFFIObjectCreateOpaque(
<void*>(pyobject),
kTVMFFIOpaquePyObject,
- tvm_ffi_pyobject_deleter,
+ TVMFFIPyObjectDeleter,
out_handle))
return 0
diff --git a/python/tvm_ffi/cython/tvm_ffi_python_helpers.h
b/python/tvm_ffi/cython/tvm_ffi_python_helpers.h
index 3e204bc..e507f03 100644
--- a/python/tvm_ffi/cython/tvm_ffi_python_helpers.h
+++ b/python/tvm_ffi/cython/tvm_ffi_python_helpers.h
@@ -591,4 +591,35 @@ TVM_FFI_INLINE void
TVMFFIPyPushTempPyObject(TVMFFIPyCallContext* ctx, PyObject*
Py_IncRef(arg);
ctx->temp_py_objects[ctx->num_temp_py_objects++] = arg;
}
+
+//------------------------------------------------------------------------------------
+// Helpers for free-threaded python
+//------------------------------------------------------------------------------------
+#if defined(Py_GIL_DISABLED)
+// NOGIL case
+class TVMFFIPyWithGILIfNotFreeThreaded {
+ public:
+ TVMFFIPyWithGILIfNotFreeThreaded() = default;
+};
+#else
+// GIL case, need to ensure/release the GIL
+class TVMFFIPyWithGILIfNotFreeThreaded {
+ public:
+ TVMFFIPyWithGILIfNotFreeThreaded() noexcept { gstate_ = PyGILState_Ensure();
}
+ ~TVMFFIPyWithGILIfNotFreeThreaded() { PyGILState_Release(gstate_); }
+
+ private:
+ PyGILState_STATE gstate_;
+};
+#endif
+
+/*!
+ * \brief Deleter for Python objects
+ * \param py_obj The Python object to delete
+ */
+extern "C" void TVMFFIPyObjectDeleter(void* py_obj) noexcept {
+ TVMFFIPyWithGILIfNotFreeThreaded gil_state;
+ Py_DecRef(static_cast<PyObject*>(py_obj));
+}
+
#endif // TVM_FFI_PYTHON_HELPERS_H_
diff --git a/src/ffi/object.cc b/src/ffi/object.cc
index 3e7e294..b3927fe 100644
--- a/src/ffi/object.cc
+++ b/src/ffi/object.cc
@@ -339,7 +339,13 @@ class TypeTable {
TypeIndex::kTVMFFIObjectRValueRef);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFISmallStr,
TypeIndex::kTVMFFISmallStr);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFISmallBytes,
TypeIndex::kTVMFFISmallBytes);
- ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIOpaquePyObject,
TypeIndex::kTVMFFIOpaquePyObject);
+ // register opaque py whose type depth is 1
+ this->GetOrAllocTypeIndex(StaticTypeKey::kTVMFFIOpaquePyObject,
+ TypeIndex::kTVMFFIOpaquePyObject,
+ /*type_depth=*/1,
+ /*num_child_slots=*/0,
+ /*child_slots_can_overflow=*/false,
+ /*parent_type_index=*/TypeIndex::kTVMFFIObject);
// no need to reserve for object types as they will be registered
}
diff --git a/tests/python/test_object.py b/tests/python/test_object.py
index f679469..12dcf3f 100644
--- a/tests/python/test_object.py
+++ b/tests/python/test_object.py
@@ -107,17 +107,20 @@ class MyObject:
def test_opaque_object() -> None:
obj0 = MyObject("hello")
- assert sys.getrefcount(obj0) == 2
+ base_count = sys.getrefcount(obj0)
+ ref_count = tvm_ffi.get_global_func("testing.object_use_count")
+ assert sys.getrefcount(obj0) == base_count
obj0_converted = tvm_ffi.convert(obj0)
- assert sys.getrefcount(obj0) == 3
+ assert ref_count(obj0_converted) == 1
+ assert sys.getrefcount(obj0) == base_count + 1
assert isinstance(obj0_converted, tvm_ffi.core.OpaquePyObject)
obj0_cpy = obj0_converted.pyobject()
assert obj0_cpy is obj0
- assert sys.getrefcount(obj0) == 4
+ assert sys.getrefcount(obj0) == base_count + 2
obj0_converted = None
- assert sys.getrefcount(obj0) == 3
+ assert sys.getrefcount(obj0) == base_count + 1
obj0_cpy = None
- assert sys.getrefcount(obj0) == 2
+ assert sys.getrefcount(obj0) == base_count
def test_opaque_type_error() -> None:
