[FFI][REFACTOR] Refactor python ffi call mechanism for perf (#18302)

This PR refactors python ffi call mechanism. Previously the argument setting can
become an as things can be sensitive to the if checking order. This PR refactors
the calling to leverage a C++ based dispatcher where each dispatch functor can be
registered from Cython.

Author: tqchen

ffi/CMakeLists.txt

@@ -215,6 +215,7 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
     Python_add_library(tvm_ffi_cython MODULE "${core_cpp}" WITH_SOABI)
     set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
   endif()
+  target_include_directories(tvm_ffi_cython PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython)
   target_compile_features(tvm_ffi_cython PRIVATE cxx_std_17)
   target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_header)
   target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_shared)

ffi/include/tvm/ffi/container/tensor.h

@@ -30,6 +30,8 @@
 #include <tvm/ffi/error.h>
 #include <tvm/ffi/type_traits.h>
 
+#include <atomic>
+#include <memory>
 #include <utility>
 
 namespace tvm {
@@ -123,18 +125,26 @@ class TensorObj : public Object, public DLTensor {
   static constexpr const uint32_t _type_index = TypeIndex::kTVMFFITensor;
   TVM_FFI_DECLARE_OBJECT_INFO_STATIC(StaticTypeKey::kTVMFFITensor, TensorObj, Object);
   /// \endcond
-
+  ~TensorObj() {
+    // deleting the cached dl managed tensor versioned
+    // need to acquire the value in case it is released by another thread
+    DLManagedTensorVersioned* cached =
+        cached_dl_managed_tensor_versioned_.load(std::memory_order_acquire);
+    if (cached != nullptr) {
+      delete cached;
+    }
+  }
   /*!
    * \brief Move a Tensor to a DLPack managed tensor.
    * \return The converted DLPack managed tensor.
    */
   DLManagedTensor* ToDLPack() const {
+    TensorObj* self = const_cast<TensorObj*>(this);
     DLManagedTensor* ret = new DLManagedTensor();
-    TensorObj* from = const_cast<TensorObj*>(this);
-    ret->dl_tensor = *static_cast<DLTensor*>(from);
-    ret->manager_ctx = from;
+    ret->dl_tensor = *static_cast<DLTensor*>(self);
+    ret->manager_ctx = self;
     ret->deleter = DLManagedTensorDeleter;
-    details::ObjectUnsafe::IncRefObjectHandle(from);
+    details::ObjectUnsafe::IncRefObjectHandle(self);
     return ret;
   }
 
@@ -143,23 +153,49 @@ class TensorObj : public Object, public DLTensor {
    * \return The converted DLPack managed tensor.
    */
   DLManagedTensorVersioned* ToDLPackVersioned() const {
-    DLManagedTensorVersioned* ret = new DLManagedTensorVersioned();
     TensorObj* from = const_cast<TensorObj*>(this);
-    ret->version.major = DLPACK_MAJOR_VERSION;
-    ret->version.minor = DLPACK_MINOR_VERSION;
-    ret->dl_tensor = *static_cast<DLTensor*>(from);
-    ret->manager_ctx = from;
-    ret->deleter = DLManagedTensorVersionedDeleter;
-    ret->flags = 0;
+    // if cache is set, directly return it
+    // we need to use acquire to ensure that write to DLManagedTensorVersioned
+    // from another thread is visible to this thread.
+    DLManagedTensorVersioned* cached =
+        cached_dl_managed_tensor_versioned_.load(std::memory_order_acquire);
+    // if cache is not set, create a new one
+    if (cached == nullptr) {
+      DLManagedTensorVersioned* ret = new DLManagedTensorVersioned();
+      ret->version.major = DLPACK_MAJOR_VERSION;
+      ret->version.minor = DLPACK_MINOR_VERSION;
+      ret->dl_tensor = *static_cast<DLTensor*>(from);
+      ret->manager_ctx = from;
+      ret->deleter = EmbeddedDLManagedTensorVersionedDeleter;
+      ret->flags = 0;
+      DLManagedTensorVersioned* expected = nullptr;
+      // success set must release the new value to all other threads
+      // failure set must acquire, since the expected value is now coming
+      // from another thread that released this value
+      if (std::atomic_compare_exchange_strong_explicit(&cached_dl_managed_tensor_versioned_,
+                                                       &expected, ret, std::memory_order_release,
+                                                       std::memory_order_acquire)) {
+        // set is succes
+        cached = ret;
+      } else {
+        // delete the ret value as another thread raced to set this one first
+        delete ret;
+        cached = expected;
+      }
+      // at this point, cached is the value that officially set to the field
+    }
+    // inc the ref count of the from object
     details::ObjectUnsafe::IncRefObjectHandle(from);
-    return ret;
+    return cached;
   }
 
  protected:
   /*! \brief Internal data to back returning shape. */
   Optional<Shape> shape_data_;
   /*! \brief Internal data to back returning strides. */
   Optional<Shape> strides_data_;
+  /*! \brief cached data to back returning DLManagedTensorVersioned. */
+  mutable std::atomic<DLManagedTensorVersioned*> cached_dl_managed_tensor_versioned_ = nullptr;
 
   /*!
    * \brief Deleter for DLManagedTensor.
@@ -175,10 +211,9 @@ class TensorObj : public Object, public DLTensor {
    * \brief Deleter for DLManagedTensorVersioned.
    * \param tensor The DLManagedTensorVersioned to be deleted.
    */
-  static void DLManagedTensorVersionedDeleter(DLManagedTensorVersioned* tensor) {
+  static void EmbeddedDLManagedTensorVersionedDeleter(DLManagedTensorVersioned* tensor) {
     TensorObj* obj = static_cast<TensorObj*>(tensor->manager_ctx);
     details::ObjectUnsafe::DecRefObjectHandle(obj);
-    delete tensor;
   }
 
   friend class Tensor;

ffi/python/tvm_ffi/cython/base.pxi

@@ -72,7 +72,7 @@ cdef extern from "dlpack/dlpack.h":
 
     ctypedef struct DLManagedTensorVersioned:
         DLPackVersion version
-        DLManagedTensor dl_tensor
+        DLTensor dl_tensor
         void* manager_ctx
         void (*deleter)(DLManagedTensorVersioned* self)
         uint64_t flags
@@ -195,6 +195,7 @@ cdef extern from "tvm/ffi/c_api.h":
         const TVMFFITypeMetadata* metadata
 
     int TVMFFIObjectDecRef(TVMFFIObjectHandle obj) nogil
+    int TVMFFIObjectIncRef(TVMFFIObjectHandle obj) nogil
     int TVMFFIObjectCreateOpaque(void* handle, int32_t type_index,
                                  void (*deleter)(void*), TVMFFIObjectHandle* out) nogil
     int TVMFFIObjectGetTypeIndex(TVMFFIObjectHandle obj) nogil
@@ -243,6 +244,58 @@ cdef extern from "tvm/ffi/extra/c_env_api.h":
                                   TVMFFIStreamHandle* opt_out_original_stream) nogil
 
 
+cdef extern from "tvm_ffi_python_helpers.h":
+    # no need to expose fields of the call context
+    ctypedef struct TVMFFIPyCallContext:
+        int device_type
+        int device_id
+        TVMFFIStreamHandle stream
+
+     # setter data structure
+    ctypedef int (*DLPackPyObjectCExporter)(
+        void* py_obj, DLManagedTensorVersioned** out, TVMFFIStreamHandle* env_stream
+    ) except -1
+
+    ctypedef struct TVMFFIPyArgSetter:
+        int (*func)(TVMFFIPyArgSetter* handle, TVMFFIPyCallContext* ctx,  PyObject* py_arg, TVMFFIAny* out) except -1
+        DLPackPyObjectCExporter dlpack_c_exporter
+
+    ctypedef int (*TVMFFIPyArgSetterFactory)(PyObject* value, TVMFFIPyArgSetter* out) except -1
+    # The main call function
+    int TVMFFIPyFuncCall(
+        TVMFFIPyArgSetterFactory setter_factory,
+        void* chandle,
+        PyObject* py_arg_tuple,
+        TVMFFIAny* result,
+        int* c_api_ret_code
+    ) except -1
+
+    int TVMFFIPyCallFieldSetter(
+        TVMFFIPyArgSetterFactory setter_factory,
+        TVMFFIFieldSetter field_setter,
+        void* field_ptr,
+        PyObject* py_arg,
+        int* c_api_ret_code
+    ) except -1
+
+    int TVMFFIPyPyObjectToFFIAny(
+        TVMFFIPyArgSetterFactory setter_factory,
+        PyObject* py_arg,
+        TVMFFIAny* out,
+        int* c_api_ret_code
+    ) except -1
+
+    size_t TVMFFIPyGetDispatchMapSize() noexcept
+
+    void TVMFFIPyPushTempFFIObject(TVMFFIPyCallContext* ctx, TVMFFIObjectHandle arg) noexcept
+    void TVMFFIPyPushTempPyObject(TVMFFIPyCallContext* ctx, PyObject* arg) noexcept
+    # the predefined setters for common POD types
+    int TVMFFIPyArgSetterFloat_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1
+    int TVMFFIPyArgSetterInt_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1
+    int TVMFFIPyArgSetterBool_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1
+    int TVMFFIPyArgSetterNone_(TVMFFIPyArgSetter*, TVMFFIPyCallContext*, PyObject* arg, TVMFFIAny* out) except -1
+
+
 cdef class ByteArrayArg:
     cdef TVMFFIByteArray cdata
     cdef object py_data