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 4edf4f3 feat(cython): Expose TypeAttr via `_lookup_type_attr` (#247)
4edf4f3 is described below
commit 4edf4f30e50c52e082ad178eb91016378335040d
Author: Junru Shao <[email protected]>
AuthorDate: Sat Nov 8 17:31:55 2025 -0800
feat(cython): Expose TypeAttr via `_lookup_type_attr` (#247)
This PR introduces an API that exposes type attributes on Python side.
For now, there's no type attribute being defined, but I personally do
anticipate usecases (e.g. `__repr__`, `__metadata__`, etc.) being used
as part of stubgen & dataclass push.
---
include/tvm/ffi/c_api.h | 6 +++---
python/tvm_ffi/core.pyi | 1 +
python/tvm_ffi/cython/base.pxi | 5 +++++
python/tvm_ffi/cython/object.pxi | 9 +++++++++
src/ffi/object.cc | 8 ++++----
tests/python/test_metadata.py | 6 +++++-
6 files changed, 27 insertions(+), 8 deletions(-)
diff --git a/include/tvm/ffi/c_api.h b/include/tvm/ffi/c_api.h
index ebe1a99..e773208 100644
--- a/include/tvm/ffi/c_api.h
+++ b/include/tvm/ffi/c_api.h
@@ -156,9 +156,6 @@ typedef enum {
kTVMFFITensor = 70,
/*! \brief Array object. */
kTVMFFIArray = 71,
- //----------------------------------------------------------------
- // more complex objects
- //----------------------------------------------------------------
/*! \brief Map object. */
kTVMFFIMap = 72,
/*! \brief Runtime dynamic loaded module object. */
@@ -175,6 +172,9 @@ typedef enum {
* \sa TVMFFIObjectCreateOpaque
*/
kTVMFFIOpaquePyObject = 74,
+ //----------------------------------------------------------------
+ // more complex objects
+ //----------------------------------------------------------------
kTVMFFIStaticObjectEnd,
// [Section] Dynamic Boxed: [kTVMFFIDynObjectBegin, +oo)
/*! \brief Start of type indices that are allocated at runtime. */
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index d5d43b7..8141748 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -67,6 +67,7 @@ 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_info: TypeInfo, type_cls: type) -> None: ...
def _lookup_or_register_type_info_from_type_key(type_key: str) -> TypeInfo: ...
+def _lookup_type_attr(type_index: int, attr_key: str) -> Any: ...
class Error(Object):
def __init__(self, kind: str, message: str, backtrace: str) -> None: ...
diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi
index 597956b..a24a5d8 100644
--- a/python/tvm_ffi/cython/base.pxi
+++ b/python/tvm_ffi/cython/base.pxi
@@ -244,6 +244,10 @@ cdef extern from "tvm/ffi/c_api.h":
const TVMFFIMethodInfo* methods
const TVMFFITypeMetadata* metadata
+ ctypedef struct TVMFFITypeAttrColumn:
+ const TVMFFIAny* data
+ size_t size
+
int TVMFFIObjectDecRef(TVMFFIObjectHandle obj) nogil
int TVMFFIObjectIncRef(TVMFFIObjectHandle obj) nogil
int TVMFFIObjectCreateOpaque(void* handle, int32_t type_index,
@@ -285,6 +289,7 @@ cdef extern from "tvm/ffi/c_api.h":
TVMFFIShapeCell* TVMFFIShapeGetCellPtr(TVMFFIObjectHandle obj) nogil
DLTensor* TVMFFITensorGetDLTensorPtr(TVMFFIObjectHandle obj) nogil
DLDevice TVMFFIDLDeviceFromIntPair(int32_t device_type, int32_t device_id)
nogil
+ const TVMFFITypeAttrColumn* TVMFFIGetTypeAttrColumn(const TVMFFIByteArray*
attr_name) nogil
cdef extern from "tvm/ffi/extra/c_env_api.h":
ctypedef void* TVMFFIStreamHandle
diff --git a/python/tvm_ffi/cython/object.pxi b/python/tvm_ffi/cython/object.pxi
index 0c5bdac..49c7595 100644
--- a/python/tvm_ffi/cython/object.pxi
+++ b/python/tvm_ffi/cython/object.pxi
@@ -514,6 +514,15 @@ def _lookup_or_register_type_info_from_type_key(type_key:
str) -> TypeInfo:
return info
+def _lookup_type_attr(type_index: int32_t, attr_key: str) -> Any:
+ cdef ByteArrayArg attr_key_bytes = ByteArrayArg(c_str(attr_key))
+ cdef const TVMFFITypeAttrColumn* column =
TVMFFIGetTypeAttrColumn(&attr_key_bytes.cdata)
+ cdef TVMFFIAny data
+ if column == NULL or column.size <= type_index:
+ return None
+ return make_ret(column.data[type_index])
+
+
cdef list TYPE_INDEX_TO_CLS = []
cdef list TYPE_INDEX_TO_INFO = []
cdef dict TYPE_KEY_TO_INFO = {}
diff --git a/src/ffi/object.cc b/src/ffi/object.cc
index 810d3b1..1671ba8 100644
--- a/src/ffi/object.cc
+++ b/src/ffi/object.cc
@@ -315,7 +315,7 @@ class TypeTable {
}
private:
- TypeTable() {
+ TypeTable() : type_attr_columns_{}, type_attr_name_to_column_index_{} {
type_table_.reserve(TypeIndex::kTVMFFIDynObjectBegin);
for (int32_t i = 0; i < TypeIndex::kTVMFFIDynObjectBegin; ++i) {
type_table_.emplace_back(nullptr);
@@ -332,13 +332,13 @@ class TypeTable {
// reserve the static types
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFINone,
TypeIndex::kTVMFFINone);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIInt, TypeIndex::kTVMFFIInt);
- ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIFloat,
TypeIndex::kTVMFFIFloat);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIBool,
TypeIndex::kTVMFFIBool);
- ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIRawStr,
TypeIndex::kTVMFFIRawStr);
+ ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIFloat,
TypeIndex::kTVMFFIFloat);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIOpaquePtr,
TypeIndex::kTVMFFIOpaquePtr);
- ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDLTensorPtr,
TypeIndex::kTVMFFIDLTensorPtr);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDataType,
TypeIndex::kTVMFFIDataType);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDevice,
TypeIndex::kTVMFFIDevice);
+ ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIDLTensorPtr,
TypeIndex::kTVMFFIDLTensorPtr);
+ ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIRawStr,
TypeIndex::kTVMFFIRawStr);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIByteArrayPtr,
TypeIndex::kTVMFFIByteArrayPtr);
ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIObjectRValueRef,
TypeIndex::kTVMFFIObjectRValueRef);
diff --git a/tests/python/test_metadata.py b/tests/python/test_metadata.py
index 787ed91..116afd3 100644
--- a/tests/python/test_metadata.py
+++ b/tests/python/test_metadata.py
@@ -18,7 +18,7 @@ from typing import Any
import pytest
from tvm_ffi import get_global_func_metadata
-from tvm_ffi.core import TypeInfo, TypeSchema
+from tvm_ffi.core import TypeInfo, TypeSchema, _lookup_type_attr
from tvm_ffi.testing import _SchemaAllTypes
@@ -195,3 +195,7 @@ def test_mem_fn_as_global_func() -> None:
metadata: dict[str, Any] =
get_global_func_metadata("testing.TestIntPairSum")
type_schema: TypeSchema = TypeSchema.from_json_str(metadata["type_schema"])
assert str(type_schema) == "Callable[[testing.TestIntPair], int]"
+
+
+def test_class_metadata_none() -> None:
+ assert _lookup_type_attr(64, "__metadata__") is None
