[FFI][ABI] Introduce weak rc support (#18259)

This PR adds weak ref counter support to the FFI ABI.
Weak rc is useful when we want to break cyclic dependencies.

- When a strong rc goes to zero, we call the destructor of the object, but not freeing the memory
- When both strong and weak rc goes to zero, we call the memory free operation

The weak rc mechanism is useful when we want to break cyclic dependencies in object, where the
weak rc can keep memory alive but the destructor is called.
As of now, because we deliberately avoid cyles in codebase, we do not have strong use-case for weak rc.

However, given weak rc is common practice in shared_ptr, Rust RC, and also used in torch's c10::intrusive_ptr.
It is better to make sure the ABI is future compatible to such use-cases before we freeze.

This PR implements weak rc as a u32 counter and strong rc as a u64 counter, with the following
design consideration.

- Weak rc is very rarely used and u32 is sufficient.
- Keeping weak rc in u32 allows us to keep object header size to 24 bytes,
  saving extra 8 bytes(considering alignment)

We also need to update deleter to take flags that consider both weak and
strong deletion events. The implementation tries to optimize common case where
both strong and weak goes to 0 at the same time and call deleter once
with both flags set.

Author: tqchen

ffi/include/tvm/ffi/c_api.h

@@ -156,6 +156,36 @@ typedef enum {
 /*! \brief Handle to Object from C API's pov */
 typedef void* TVMFFIObjectHandle;
 
+/*!
+ * \brief bitmask of the object deleter flag.
+ */
+#ifdef __cplusplus
+enum TVMFFIObjectDeleterFlagBitMask : int32_t {
+#else
+typedef enum {
+#endif
+  /*!
+   * \brief deleter action when strong reference count becomes zero.
+   * Need to call destructor of the object but not free the memory block.
+   */
+  kTVMFFIObjectDeleterFlagBitMaskStrong = 1 << 0,
+  /*!
+   * \brief deleter action when weak reference count becomes zero.
+   * Need to free the memory block.
+   */
+  kTVMFFIObjectDeleterFlagBitMaskWeak = 1 << 1,
+  /*!
+   * \brief deleter action when both strong and weak reference counts become zero.
+   * \note This is the most common case.
+   */
+  kTVMFFIObjectDeleterFlagBitMaskBoth =
+      (kTVMFFIObjectDeleterFlagBitMaskStrong | kTVMFFIObjectDeleterFlagBitMaskWeak),
+#ifdef __cplusplus
+};
+#else
+} TVMFFIObjectDeleterFlagBitMask;
+#endif
+
 /*!
  * \brief C-based type of all FFI object header that allocates on heap.
  * \note TVMFFIObject and TVMFFIAny share the common type_index header
@@ -166,11 +196,22 @@ typedef struct TVMFFIObject {
    * \note The type index of Object and Any are shared in FFI.
    */
   int32_t type_index;
-  /*! \brief Reference counter of the object. */
-  int32_t ref_counter;
+  /*!
+   * \brief Weak reference counter of the object, for compatiblity with weak_ptr design.
+   * \note Use u32 to ensure that overall object stays within 24-byte boundary, usually
+   *       manipulation of weak counter is less common than strong counter.
+   */
+  uint32_t weak_ref_count;
+  /*! \brief Strong reference counter of the object. */
+  uint64_t strong_ref_count;
   union {
-    /*! \brief Deleter to be invoked when reference counter goes to zero. */
-    void (*deleter)(struct TVMFFIObject* self);
+    /*!
+     * \brief Deleter to be invoked when strong reference counter goes to zero.
+     * \param self The self object handle.
+     * \param flags The flags to indicate deletion behavior.
+     * \sa TVMFFIObjectDeleterFlagBitMask
+     */
+    void (*deleter)(struct TVMFFIObject* self, int flags);
     /*!
      * \brief auxilary field to TVMFFIObject is always 8 bytes aligned.
      * \note This helps us to ensure cross platform compatibility.
@@ -307,13 +348,19 @@ typedef struct {
 // Section: Basic object API
 //------------------------------------------------------------
 /*!
- * \brief Free an object handle by decreasing reference
+ * \brief Increas the strong reference count of an object handle
+ * \param obj The object handle.
+ * \note Internally we increase the reference counter of the object.
+ * \return 0 when success, nonzero when failure happens
+ */
+TVM_FFI_DLL int TVMFFIObjectIncRef(TVMFFIObjectHandle obj);
+
+/*!
+ * \brief Free an object handle by decreasing strong reference
  * \param obj The object handle.
- * \note Internally we decrease the reference counter of the object.
- *       The object will be freed when every reference to the object are removed.
  * \return 0 when success, nonzero when failure happens
  */
-TVM_FFI_DLL int TVMFFIObjectFree(TVMFFIObjectHandle obj);
+TVM_FFI_DLL int TVMFFIObjectDecRef(TVMFFIObjectHandle obj);
 
 /*!
  * \brief Convert type key to type index.
@@ -470,7 +517,7 @@ TVM_FFI_DLL int TVMFFIDataTypeFromString(const TVMFFIByteArray* str, DLDataType*
 * \param dtype The DLDataType to convert.
 * \param out The output string.
 * \return 0 when success, nonzero when failure happens
-* \note out is a String object that needs to be freed by the caller via TVMFFIObjectFree.
+* \note out is a String object that needs to be freed by the caller via TVMFFIObjectDecRef.
 The content of string can be accessed via TVMFFIObjectGetByteArrayPtr.
 
 * \note The input dtype is a pointer to the DLDataType to avoid ABI compatibility issues.

ffi/include/tvm/ffi/memory.h

@@ -33,7 +33,7 @@ namespace tvm {
 namespace ffi {
 
 /*! \brief Deleter function for obeject */
-typedef void (*FObjectDeleter)(TVMFFIObject* obj);
+typedef void (*FObjectDeleter)(TVMFFIObject* obj, int flags);
 
 /*!
  * \brief Allocate an object using default allocator.
@@ -75,7 +75,8 @@ class ObjAllocatorBase {
     static_assert(std::is_base_of<Object, T>::value, "make can only be used to create Object");
     T* ptr = Handler::New(static_cast<Derived*>(this), std::forward<Args>(args)...);
     TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr);
-    ffi_ptr->ref_counter = 1;
+    ffi_ptr->strong_ref_count = 1;
+    ffi_ptr->weak_ref_count = 1;
     ffi_ptr->type_index = T::RuntimeTypeIndex();
     ffi_ptr->deleter = Handler::Deleter();
     return details::ObjectUnsafe::ObjectPtrFromOwned<T>(ptr);
@@ -96,7 +97,8 @@ class ObjAllocatorBase {
     ArrayType* ptr =
         Handler::New(static_cast<Derived*>(this), num_elems, std::forward<Args>(args)...);
     TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr);
-    ffi_ptr->ref_counter = 1;
+    ffi_ptr->strong_ref_count = 1;
+    ffi_ptr->weak_ref_count = 1;
     ffi_ptr->type_index = ArrayType::RuntimeTypeIndex();
     ffi_ptr->deleter = Handler::Deleter();
     return details::ObjectUnsafe::ObjectPtrFromOwned<ArrayType>(ptr);
@@ -136,14 +138,18 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
     static FObjectDeleter Deleter() { return Deleter_; }
 
    private:
-    static void Deleter_(TVMFFIObject* objptr) {
+    static void Deleter_(TVMFFIObject* objptr, int flags) {
       T* tptr = details::ObjectUnsafe::RawObjectPtrFromUnowned<T>(objptr);
-      // It is important to do tptr->T::~T(),
-      // so that we explicitly call the specific destructor
-      // instead of tptr->~T(), which could mean the intention
-      // call a virtual destructor(which may not be available and is not required).
-      tptr->T::~T();
-      delete reinterpret_cast<StorageType*>(tptr);
+      if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) {
+        // It is important to do tptr->T::~T(),
+        // so that we explicitly call the specific destructor
+        // instead of tptr->~T(), which could mean the intention
+        // call a virtual destructor(which may not be available and is not required).
+        tptr->T::~T();
+      }
+      if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
+        delete reinterpret_cast<StorageType*>(tptr);
+      }
     }
   };
 
@@ -182,15 +188,19 @@ class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
     static FObjectDeleter Deleter() { return Deleter_; }
 
    private:
-    static void Deleter_(TVMFFIObject* objptr) {
+    static void Deleter_(TVMFFIObject* objptr, int flags) {
       ArrayType* tptr = details::ObjectUnsafe::RawObjectPtrFromUnowned<ArrayType>(objptr);
-      // It is important to do tptr->ArrayType::~ArrayType(),
-      // so that we explicitly call the specific destructor
-      // instead of tptr->~ArrayType(), which could mean the intention
-      // call a virtual destructor(which may not be available and is not required).
-      tptr->ArrayType::~ArrayType();
-      StorageType* p = reinterpret_cast<StorageType*>(tptr);
-      delete[] p;
+      if (flags & kTVMFFIObjectDeleterFlagBitMaskStrong) {
+        // It is important to do tptr->ArrayType::~ArrayType(),
+        // so that we explicitly call the specific destructor
+        // instead of tptr->~ArrayType(), which could mean the intention
+        // call a virtual destructor(which may not be available and is not required).
+        tptr->ArrayType::~ArrayType();
+      }
+      if (flags & kTVMFFIObjectDeleterFlagBitMaskWeak) {
+        StorageType* p = reinterpret_cast<StorageType*>(tptr);
+        delete[] p;
+      }
     }
   };
 };