[FFI][CYTHON] Release GIL when calling into long running functions (#11461)

Unlike ctypes, Cython by default do not release GIL when
calling into C API functions. This causes problems when the
function is long running. As the particular calling thread will
block other python threads by holding the GIL.

This PR explicitly releases GIL when calling into possible
long running functions. It fixes the timeout issue in
PopenPool which previously relied on another python thread
for timeout.

Added a regression test-case by changing sleep to sleep
in FFI, which previously will indefinitely block the popen tests.

Author: tqchen

python/tvm/_ffi/_cython/base.pxi

@@ -94,17 +94,34 @@ ctypedef int (*TVMPackedCFunc)(
 
 ctypedef void (*TVMPackedCFuncFinalizer)(void* resource_handle)
 
+# NOTE: All of TVM's C API function can be called without gil.
+# for API functions that can be run long(e.g. FuncCall)
+# we need to explicitly release the GIL as follows.
+#
+# cdef myfunc():
+#     cdef int c_api_ret_code
+#     with nogil:
+#         c_api_ret_code = TVMAPIFunc(...)
+#     CHECK_CALL(c_apt_ret_code)
+#
+# Explicitly releasing the GIL enables other python threads
+# to continue running while we are in TVMAPIFunc.
+# Not releasing GIL explicitly is OK(and perhaps desirable)
+# for short-running functions, as frequent unlocking also takes time,
+# the python interpreter will release GIL in a set period.
+#
+# We mark the possibly long running function as nogil below.
 cdef extern from "tvm/runtime/c_runtime_api.h":
     void TVMAPISetLastError(const char* msg)
     const char *TVMGetLastError()
     int TVMFuncGetGlobal(const char* name,
-                         TVMPackedFuncHandle* out);
+                         TVMPackedFuncHandle* out)
     int TVMFuncCall(TVMPackedFuncHandle func,
                     TVMValue* arg_values,
                     int* type_codes,
                     int num_args,
                     TVMValue* ret_val,
-                    int* ret_type_code)
+                    int* ret_type_code) nogil
     int TVMFuncFree(TVMPackedFuncHandle func)
     int TVMCFuncSetReturn(TVMRetValueHandle ret,
                           TVMValue* value,
@@ -119,15 +136,15 @@ cdef extern from "tvm/runtime/c_runtime_api.h":
                       tvm_index_t ndim,
                       DLDataType dtype,
                       DLDevice dev,
-                      DLTensorHandle* out)
-    int TVMArrayFree(DLTensorHandle handle)
+                      DLTensorHandle* out) nogil
+    int TVMArrayFree(DLTensorHandle handle) nogil
     int TVMArrayCopyFromTo(DLTensorHandle src,
                            DLTensorHandle to,
-                           TVMStreamHandle stream)
+                           TVMStreamHandle stream) nogil
     int TVMArrayFromDLPack(DLManagedTensor* arr_from,
-                           DLTensorHandle* out)
+                           DLTensorHandle* out) nogil
     int TVMArrayToDLPack(DLTensorHandle arr_from,
-                         DLManagedTensor** out)
+                         DLManagedTensor** out) nogil
     void TVMDLManagedTensorCallDeleter(DLManagedTensor* dltensor)
     int TVMObjectFree(ObjectHandle obj)
     int TVMObjectGetTypeIndex(ObjectHandle obj, unsigned* out_index)
@@ -155,7 +172,8 @@ cdef inline c_str(pystr):
     return pystr.encode("utf-8")
 
 
-cdef inline int CALL(int ret) except -2:
+cdef inline int CHECK_CALL(int ret) except -2:
+    """Check the return code of the C API function call"""
     # -2 brings exception
     if ret == -2:
         return -2
@@ -193,6 +211,6 @@ cdef _init_env_api():
     #
     # When the functions are not registered, the signals will be handled
     # only when the FFI function returns.
-    CALL(TVMBackendRegisterEnvCAPI(c_str("PyErr_CheckSignals"), <void*>PyErr_CheckSignals))
+    CHECK_CALL(TVMBackendRegisterEnvCAPI(c_str("PyErr_CheckSignals"), <void*>PyErr_CheckSignals))
 
 _init_env_api()

python/tvm/_ffi/_cython/ndarray.pxi

@@ -31,9 +31,12 @@ cdef void _c_dlpack_deleter(object pycaps):
 def _from_dlpack(object dltensor):
     cdef DLManagedTensor* ptr
     cdef DLTensorHandle chandle
+    cdef int c_api_ret_code
     if pycapsule.PyCapsule_IsValid(dltensor, _c_str_dltensor):
         ptr = <DLManagedTensor*>pycapsule.PyCapsule_GetPointer(dltensor, _c_str_dltensor)
-        CALL(TVMArrayFromDLPack(ptr, &chandle))
+        with nogil:
+            c_api_ret_code = TVMArrayFromDLPack(ptr, &chandle)
+        CHECK_CALL(c_api_ret_code)
         # set name and destructor to be empty
         pycapsule.PyCapsule_SetDestructor(dltensor, NULL)
         pycapsule.PyCapsule_SetName(dltensor, _c_str_used_dltensor)
@@ -82,12 +85,18 @@ cdef class NDArrayBase:
         self.c_is_view = is_view
 
     def __dealloc__(self):
+        cdef int c_api_ret_code
         if self.c_is_view == 0:
-            CALL(TVMArrayFree(self.chandle))
+            with nogil:
+                c_api_ret_code = TVMArrayFree(self.chandle)
+            CHECK_CALL(c_api_ret_code)
 
     def _copyto(self, target_nd):
         """Internal function that implements copy to target ndarray."""
-        CALL(TVMArrayCopyFromTo(self.chandle, (<NDArrayBase>target_nd).chandle, NULL))
+        cdef int c_api_ret_code
+        with nogil:
+            c_api_ret_code = TVMArrayCopyFromTo(self.chandle, (<NDArrayBase>target_nd).chandle, NULL)
+        CHECK_CALL(c_api_ret_code)
         return target_nd
 
     def to_dlpack(self):
@@ -98,9 +107,12 @@ cdef class NDArrayBase:
         dlpack : DLPack tensor view of the array data
         """
         cdef DLManagedTensor* dltensor
+        cdef int c_api_ret_code
         if self.c_is_view != 0:
             raise ValueError("to_dlpack do not work with memory views")
-        CALL(TVMArrayToDLPack(self.chandle, &dltensor))
+        with nogil:
+            c_api_ret_code = TVMArrayToDLPack(self.chandle, &dltensor)
+        CHECK_CALL(c_api_ret_code)
         return pycapsule.PyCapsule_New(dltensor, _c_str_dltensor, _c_dlpack_deleter)
 
 

python/tvm/_ffi/_cython/object.pxi

@@ -38,7 +38,7 @@ cdef inline object make_ret_object(void* chandle):
     cdef object handle
     object_type = OBJECT_TYPE
     handle = ctypes_handle(chandle)
-    CALL(TVMObjectGetTypeIndex(chandle, &tindex))
+    CHECK_CALL(TVMObjectGetTypeIndex(chandle, &tindex))
 
     if tindex < len(OBJECT_TYPE):
         cls = OBJECT_TYPE[tindex]
@@ -101,7 +101,7 @@ cdef class ObjectBase:
             self._set_handle(value)
 
     def __dealloc__(self):
-        CALL(TVMObjectFree(self.chandle))
+        CHECK_CALL(TVMObjectFree(self.chandle))
 
     def __init_handle_by_constructor__(self, fconstructor, *args):
         """Initialize the handle by calling constructor function.

python/tvm/_ffi/_cython/packed_func.pxi

@@ -46,7 +46,7 @@ cdef int tvm_callback(TVMValue* args,
             tcode == kTVMNDArrayHandle or
             tcode == kTVMObjectRefArg or
             tcode > kTVMExtBegin):
-            CALL(TVMCbArgToReturn(&value, &tcode))
+            CHECK_CALL(TVMCbArgToReturn(&value, &tcode))
 
         if tcode != kTVMDLTensorHandle:
             pyargs.append(make_ret(value, tcode))
@@ -64,7 +64,7 @@ cdef int tvm_callback(TVMValue* args,
             raise ValueError("PackedFunction can only support one return value")
         temp_args = []
         make_arg(rv, &value, &tcode, temp_args)
-        CALL(TVMCFuncSetReturn(ret, &value, &tcode, 1))
+        CHECK_CALL(TVMCFuncSetReturn(ret, &value, &tcode, 1))
     return 0
 
 
@@ -90,10 +90,10 @@ def convert_to_tvm_func(object pyfunc):
     """
     cdef TVMPackedFuncHandle chandle
     Py_INCREF(pyfunc)
-    CALL(TVMFuncCreateFromCFunc(tvm_callback,
-                                <void*>(pyfunc),
-                                tvm_callback_finalize,
-                                &chandle))
+    CHECK_CALL(TVMFuncCreateFromCFunc(tvm_callback,
+                                      <void*>(pyfunc),
+                                      tvm_callback_finalize,
+                                      &chandle))
     return make_packed_func(chandle, False)
 
 
@@ -243,15 +243,20 @@ cdef inline int FuncCall3(void* chandle,
     temp_args = []
     for i in range(nargs):
         make_arg(args[i], &values[i], &tcodes[i], temp_args)
-    CALL(TVMFuncCall(chandle, &values[0], &tcodes[0],
-                     nargs, ret_val, ret_tcode))
+
+    with nogil:
+        c_api_ret_code = TVMFuncCall(chandle, &values[0], &tcodes[0],
+                                     nargs, ret_val, ret_tcode)
+
+    CHECK_CALL(c_api_ret_code)
     return 0
 
 cdef inline int FuncCall(void* chandle,
                          tuple args,
                          TVMValue* ret_val,
                          int* ret_tcode) except -1:
     cdef int nargs
+    cdef int c_api_ret_code
     nargs = len(args)
     if nargs <= 3:
         FuncCall3(chandle, args, nargs, ret_val, ret_tcode)
@@ -264,8 +269,11 @@ cdef inline int FuncCall(void* chandle,
     temp_args = []
     for i in range(nargs):
         make_arg(args[i], &values[i], &tcodes[i], temp_args)
-    CALL(TVMFuncCall(chandle, &values[0], &tcodes[0],
-                     nargs, ret_val, ret_tcode))
+
+    with nogil:
+        c_api_ret_code = TVMFuncCall(chandle, &values[0], &tcodes[0],
+                                     nargs, ret_val, ret_tcode)
+    CHECK_CALL(c_api_ret_code)
     return 0
 
 
@@ -314,7 +322,7 @@ cdef class PackedFuncBase:
 
     def __dealloc__(self):
         if self.is_global == 0:
-            CALL(TVMFuncFree(self.chandle))
+            CHECK_CALL(TVMFuncFree(self.chandle))
 
     def __call__(self, *args):
         cdef TVMValue ret_val
@@ -326,7 +334,7 @@ cdef class PackedFuncBase:
 
 def _get_global_func(name, allow_missing):
     cdef TVMPackedFuncHandle chandle
-    CALL(TVMFuncGetGlobal(c_str(name), &chandle))
+    CHECK_CALL(TVMFuncGetGlobal(c_str(name), &chandle))
     if chandle != NULL:
         return make_packed_func(chandle, True)
 

python/tvm/testing/popen_pool.py

@@ -16,8 +16,8 @@
 # under the License.
 # pylint: disable=invalid-name, missing-function-docstring
 """Common functions for popen_pool test cases"""
-import time
 import tvm
+from . import _ffi_api
 
 TEST_GLOBAL_STATE_1 = 0
 TEST_GLOBAL_STATE_2 = 0
@@ -72,4 +72,4 @@ def slow_summation(n):
 
 
 def timeout_job(n):
-    time.sleep(n * 1.5)
+    _ffi_api.sleep_in_ffi(n * 1.5)

src/support/ffi_testing.cc

@@ -28,6 +28,7 @@
 #include <tvm/te/tensor.h>
 #include <tvm/tir/expr.h>
 
+#include <chrono>
 #include <thread>
 
 namespace tvm {
@@ -159,4 +160,9 @@ runtime::Module NewFrontendTestModule() {
 
 TVM_REGISTER_GLOBAL("testing.FrontendTestModule").set_body_typed(NewFrontendTestModule);
 
+TVM_REGISTER_GLOBAL("testing.sleep_in_ffi").set_body_typed([](double timeout) {
+  std::chrono::duration<int64_t, std::nano> duration(static_cast<int64_t>(timeout * 1e9));
+  std::this_thread::sleep_for(duration);
+});
+
 }  // namespace tvm