This is an automated email from the ASF dual-hosted git repository.
cyx-6 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 f894595b [FEAT] Add core opaque pointer handler map (#708)
f894595b is described below
commit f894595bc5db1183a2bbd5eb9c2f2471d8dba1fb
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Aug 21 00:36:55 2026 -0400
[FEAT] Add core opaque pointer handler map (#708)
Third-party extension types cannot always define or accept
monkey-patching of `__tvm_ffi_opaque_ptr__`. Add a direct exact-class
`tvm_ffi.core._OPAQUE_PTR_HANDLERS` dictionary in the Cython core so
those objects can participate in opaque-pointer argument conversion
without changing their owning package.
Handlers are populated during application setup and treated as read-only
once FFI calls begin. Native conversions and the class-defined
opaque-pointer protocol retain precedence.
---
docs/guides/python_lang_guide.md | 21 +++++++++++++++++
python/tvm_ffi/core.pyi | 2 ++
python/tvm_ffi/cython/function.pxi | 29 ++++++++++++++++++++++++
python/tvm_ffi/cython/pyclass_type_converter.pxi | 6 ++++-
tests/python/test_function.py | 22 ++++++++++++++++++
5 files changed, 79 insertions(+), 1 deletion(-)
diff --git a/docs/guides/python_lang_guide.md b/docs/guides/python_lang_guide.md
index 83b0dd2e..1f98deb2 100644
--- a/docs/guides/python_lang_guide.md
+++ b/docs/guides/python_lang_guide.md
@@ -163,6 +163,27 @@ compile_kernel =
tvm_ffi.get_global_func("my_compiler.compile_kernel")
compile_kernel(workspace)
```
+When a third-party class cannot define or be monkey-patched with that
protocol, add an exact-class
+handler during application setup:
+
+```python
+import tvm_ffi
+
+
+class ForeignWorkspace:
+ __slots__ = ("address",)
+
+ def __init__(self, address: int):
+ self.address = address
+
+
+tvm_ffi.core._OPAQUE_PTR_HANDLERS[ForeignWorkspace] = lambda workspace:
workspace.address
+```
+
+Configure `tvm_ffi.core._OPAQUE_PTR_HANDLERS` before instances of those
classes are passed to FFI,
+then treat the mapping as read-only while FFI calls are running. Dispatch
matches exact classes
+rather than subclasses, and a class's own `__tvm_ffi_opaque_ptr__` takes
precedence.
+
## Container Types
TVM FFI provides five container types that split into **immutable**
(copy-on-write) and
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 29659bd6..668be688 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -206,6 +206,8 @@ class DLTensorTestWrapper:
def _dltensor_test_wrapper_c_dlpack_from_pyobject_as_intptr() -> int: ...
+_OPAQUE_PTR_HANDLERS: dict[type, Callable[[Any], int]]
+
class Function(Object):
@property
def release_gil(self) -> bool: ...
diff --git a/python/tvm_ffi/cython/function.pxi
b/python/tvm_ffi/cython/function.pxi
index c0486ccf..d2a73626 100644
--- a/python/tvm_ffi/cython/function.pxi
+++ b/python/tvm_ffi/cython/function.pxi
@@ -550,6 +550,32 @@ cdef int TVMFFIPyArgSetterFFIOpaquePtrCompatible_(
return 0
+_OPAQUE_PTR_HANDLERS = {}
+"""Map exact Python classes to their setup-time opaque-pointer handlers.
+
+Populate this dictionary before instances of registered classes are passed to
+TVM FFI, then treat it as read-only while FFI calls are running.
+"""
+
+
+cdef object _lookup_opaque_ptr_handler(object type_cls):
+ """Return the setup-time exact-class opaque pointer handler, if present."""
+ return _OPAQUE_PTR_HANDLERS.get(type_cls)
+
+
+cdef int TVMFFIPyArgSetterOpaquePtrHandler_(
+ TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
+ PyObject* py_arg, TVMFFIAny* out
+) except -1:
+ """Use the exact-class handler installed during application setup."""
+ cdef object arg = <object>py_arg
+ cdef object func = _OPAQUE_PTR_HANDLERS[type(arg)]
+ cdef long long long_ptr = <long long>func(arg)
+ out.type_index = kTVMFFIOpaquePtr
+ out.v_ptr = <void*>long_ptr
+ return 0
+
+
cdef int TVMFFIPyArgSetterObjectRValueRef_(
TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,
PyObject* py_arg, TVMFFIAny* out
@@ -876,6 +902,9 @@ cdef public int TVMFFICyArgSetterFactory(PyObject* value,
TVMFFIPyArgSetter* out
if hasattr(arg_class, "__tvm_ffi_opaque_ptr__"):
out.func = TVMFFIPyArgSetterFFIOpaquePtrCompatible_
return 0
+ if _lookup_opaque_ptr_handler(arg_class) is not None:
+ out.func = TVMFFIPyArgSetterOpaquePtrHandler_
+ return 0
if callable(arg):
out.func = TVMFFIPyArgSetterCallable_
return 0
diff --git a/python/tvm_ffi/cython/pyclass_type_converter.pxi
b/python/tvm_ffi/cython/pyclass_type_converter.pxi
index 43d362e2..8587e401 100644
--- a/python/tvm_ffi/cython/pyclass_type_converter.pxi
+++ b/python/tvm_ffi/cython/pyclass_type_converter.pxi
@@ -267,7 +267,11 @@ cdef CAny _tc_convert_opaque_ptr(_TypeConverter _conv,
object value, bint* chang
if value is None:
changed[0] = True
return CAny(ctypes.c_void_p(None))
- if hasattr(vtype, "__tvm_ffi_opaque_ptr__") or hasattr(vtype,
"__cuda_stream__"):
+ if (
+ hasattr(vtype, "__tvm_ffi_opaque_ptr__")
+ or _lookup_opaque_ptr_handler(vtype) is not None
+ or hasattr(vtype, "__cuda_stream__")
+ ):
changed[0] = True
return CAnyChecked(value, "ctypes.c_void_p", value)
raise _ConvertError(f"expected ctypes.c_void_p, got
{_tc_describe_value_type(value)}")
diff --git a/tests/python/test_function.py b/tests/python/test_function.py
index 2dd3ad86..c7ba2e9a 100644
--- a/tests/python/test_function.py
+++ b/tests/python/test_function.py
@@ -341,6 +341,28 @@ def test_function_with_opaque_ptr_protocol() -> None:
assert y.value == 10
+def test_function_with_opaque_ptr_handler() -> None:
+ """The setup map converts an exact third-party class to an opaque
pointer."""
+
+ class ForeignEvent:
+ __slots__ = ("address",)
+
+ def __init__(self, address: int) -> None:
+ self.address = address
+
+ handlers = tvm_ffi.core._OPAQUE_PTR_HANDLERS
+ handlers[ForeignEvent] = lambda value: value.address
+ try:
+ fecho = tvm_ffi.get_global_func("testing.echo")
+ event = ForeignEvent(0xCAFE)
+ result = fecho(event)
+ assert isinstance(result, ctypes.c_void_p)
+ assert result.value == 0xCAFE
+
tvm_ffi.core.TypeSchema.from_annotation(ctypes.c_void_p).check_value(event)
+ finally:
+ del handlers[ForeignEvent]
+
+
def test_function_with_dlpack_data_type_protocol() -> None:
class DLPackDataTypeProtocol:
def __init__(self, dlpack_data_type: tuple[int, int, int]) -> None: