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 035975a feat: Improve Critical Path Perf in Cython `make_ret_object`
(#47)
035975a is described below
commit 035975a7e6804d1d23b07942d7704c30b3fadda0
Author: Junru Shao <[email protected]>
AuthorDate: Tue Sep 23 06:54:56 2025 -0700
feat: Improve Critical Path Perf in Cython `make_ret_object` (#47)
This PR improves critical path perf on TVM FFI when returning an object
from FFI function calls.
It achieves this by establishing a direct mapping from `type_index` to
Python `type_cls`, thereby eliminating an unnecessary indirection during
the object instantiation process. The change aims to enhance the overall
efficiency of object handling at the FFI layer.
---
python/tvm_ffi/core.pyi | 1 +
python/tvm_ffi/cython/object.pxi | 35 +++++++++++++++++++++++------------
2 files changed, 24 insertions(+), 12 deletions(-)
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 7c623af..53c54dd 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -83,6 +83,7 @@ class PyNativeObject:
def _set_class_object(cls: type) -> None: ...
def _register_object_by_index(type_index: int, type_cls: type) -> TypeInfo: ...
def _object_type_key_to_index(type_key: str) -> int | None: ...
+def _set_type_cls(type_index: int, type_cls: type) -> None: ...
def _lookup_type_info_from_type_key(type_key: str) -> TypeInfo: ...
class Error(Object):
diff --git a/python/tvm_ffi/cython/object.pxi b/python/tvm_ffi/cython/object.pxi
index 4cc737a..2777d82 100644
--- a/python/tvm_ffi/cython/object.pxi
+++ b/python/tvm_ffi/cython/object.pxi
@@ -268,18 +268,16 @@ cdef inline object make_ret_object(TVMFFIAny result):
cdef object cls
tindex = result.type_index
- if tindex < len(TYPE_INDEX_TO_INFO):
- type_info = TYPE_INDEX_TO_INFO[tindex]
- if type_info is not None:
- cls = type_info.type_cls
- if cls is not None:
- if issubclass(cls, PyNativeObject):
- obj = Object.__new__(Object)
- (<Object>obj).chandle = result.v_obj
- return cls.__from_tvm_ffi_object__(cls, obj)
- obj = cls.__new__(cls)
+ if tindex < len(TYPE_INDEX_TO_CLS):
+ cls = TYPE_INDEX_TO_CLS[tindex]
+ if cls is not None:
+ if issubclass(cls, PyNativeObject):
+ obj = Object.__new__(Object)
(<Object>obj).chandle = result.v_obj
- return obj
+ return cls.__from_tvm_ffi_object__(cls, obj)
+ obj = cls.__new__(cls)
+ (<Object>obj).chandle = result.v_obj
+ return obj
# object is not found in registered entry
# in this case we need to report an warning
@@ -351,16 +349,28 @@ def _type_info_create_from_type_key(object type_cls, str
type_key):
def _register_object_by_index(int type_index, object type_cls):
- global TYPE_INDEX_TO_INFO, TYPE_KEY_TO_INFO
+ global TYPE_INDEX_TO_INFO, TYPE_KEY_TO_INFO, TYPE_INDEX_TO_CLS
cdef str type_key = _type_index_to_key(type_index)
cdef object info = _type_info_create_from_type_key(type_cls, type_key)
+ assert len(TYPE_INDEX_TO_INFO) == len(TYPE_INDEX_TO_CLS)
if (extra := type_index + 1 - len(TYPE_INDEX_TO_INFO)) > 0:
TYPE_INDEX_TO_INFO.extend([None] * extra)
+ TYPE_INDEX_TO_CLS.extend([None] * extra)
+ TYPE_INDEX_TO_CLS[type_index] = type_cls
TYPE_INDEX_TO_INFO[type_index] = info
TYPE_KEY_TO_INFO[type_key] = info
return info
+def _set_type_cls(int type_index, object type_cls):
+ global TYPE_INDEX_TO_INFO, TYPE_INDEX_TO_CLS
+ assert len(TYPE_INDEX_TO_INFO) == len(TYPE_INDEX_TO_CLS)
+ type_info = TYPE_INDEX_TO_INFO[type_index]
+ assert type_info.type_cls is None, f"Type already registered for
{type_info.type_key}"
+ type_info.type_cls = type_cls
+ TYPE_INDEX_TO_CLS[type_index] = type_cls
+
+
def _lookup_type_info_from_type_key(type_key: str) -> TypeInfo:
if info := TYPE_KEY_TO_INFO.get(type_key, None):
return info
@@ -369,6 +379,7 @@ def _lookup_type_info_from_type_key(type_key: str) ->
TypeInfo:
return info
+cdef list TYPE_INDEX_TO_CLS = []
cdef list TYPE_INDEX_TO_INFO = []
cdef dict TYPE_KEY_TO_INFO = {}
