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:

Reply via email to