(tvm-ffi) branch main updated: [CYTHON] Fix potential flaky issue in dispatcher (#63)

[tqchen]

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 cfff30b  [CYTHON] Fix potential flaky issue in dispatcher (#63)
cfff30b is described below

commit cfff30bd59e401e426ed6f3a3de5f7280ce5aed0
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Fri Sep 26 21:35:49 2025 -0400

    [CYTHON] Fix potential flaky issue in dispatcher (#63)
    
    The C++ dispatcher dispatches the argument passing by TYPE(obj) pointer
    which
    is non-owning. This means that there is the following edge case:
    - type A is registered through dispatcher
    - type A gets garbage collected (because it is a local type)
    - type B is created and uses the same memory address as type A
    
    Then when we pass in type B, it will mistakenly use the dispatch
    function for type A
    
    To prevent this, we keep alive the types that are registered through
    dispatcher
    by appending them to _DISPATCH_TYPE_KEEP_ALIVE
    
    NOTE that the total number of types that are registered through
    dispatcher is expected
    to be limited in practice so we can afford to keep them alive
---
 python/tvm_ffi/cython/function.pxi | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/python/tvm_ffi/cython/function.pxi 
b/python/tvm_ffi/cython/function.pxi
index ff79cff..095e3d6 100644
--- a/python/tvm_ffi/cython/function.pxi
+++ b/python/tvm_ffi/cython/function.pxi
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import ctypes
+import threading
 import os
 from numbers import Real, Integral
 
@@ -501,6 +502,10 @@ cdef int TVMFFIPyArgSetterDTypeFromNumpy_(
     out.v_dtype = NUMPY_DTYPE_TO_DTYPE[py_obj]
     return 0
 
+
+cdef _DISPATCH_TYPE_KEEP_ALIVE = set()
+cdef _DISPATCH_TYPE_KEEP_ALIVE_LOCK = threading.Lock()
+
 cdef int TVMFFIPyArgSetterFactory_(PyObject* value, TVMFFIPyArgSetter* out) 
except -1:
     """
     Factory function that creates an argument setter for a given Python 
argument type.
@@ -510,6 +515,24 @@ cdef int TVMFFIPyArgSetterFactory_(PyObject* value, 
TVMFFIPyArgSetter* out) exce
     # priortize native types over external types
     cdef object arg = <object>value
     cdef long long temp_ptr
+
+    # The C++ dispatcher dispatches the argument passing by TYPE(obj) pointer 
which
+    # is non-owning. This means that there is the following edge case:
+    # - type A is registered through dispatcher
+    # - type A gets garbage collected (because it is a local type)
+    # - type B is created and uses the same memory address as type A
+    #
+    # Then when we pass in type B, it will mistakenly use the dispatch 
function for type A
+    #
+    # To prevent this, we keep alive the types that are registered through 
dispatcher
+    # by adding them to _DISPATCH_TYPE_KEEP_ALIVE
+    #
+    # NOTE that the total number of types that are registered through 
dispatcher is expected
+    # to be limited in practice so we can afford to keep them alive
+    # Lock is used to ensure thread-safety for future thread-free python case
+    with _DISPATCH_TYPE_KEEP_ALIVE_LOCK:
+        _DISPATCH_TYPE_KEEP_ALIVE.add(type(arg))
+
     if arg is None:
         out.func = TVMFFIPyArgSetterNone_
         return 0