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 13436f0 [FFI][ABI] TVMFFIObject ABI to further align with torch
intrusive_ptr (#57)
13436f0 is described below
commit 13436f01111bc4218feb440a29a2e421bc148cc4
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 25 20:16:31 2025 -0400
[FFI][ABI] TVMFFIObject ABI to further align with torch intrusive_ptr (#57)
This PR updates the TVMFFIObject ABI to use
strong_ref_count/weak_ref_count
into u32 and which can be packed into the first u64 field in the object
header.
This change will open up further optimizations that combines weak/strong
into
a single u64 atomic that skips the need of extra atomic reading of weak
value
during deletion. The technique was recently in torch
We also reordered the counter to be first since most rc ops are frequent
and
it would be more cache friendly. We do give up the same type_index first
location
with TVMFFIAny, which is OK as they optimize for different goals.
---
docs/concepts/abi_overview.md | 9 +++++----
include/tvm/ffi/c_api.h | 22 +++++++++++++---------
include/tvm/ffi/object.h | 24 ++++++++++++------------
python/tvm_ffi/cython/base.pxi | 4 +++-
tests/cpp/test_c_ffi_abi.cc | 31 -------------------------------
5 files changed, 33 insertions(+), 57 deletions(-)
diff --git a/docs/concepts/abi_overview.md b/docs/concepts/abi_overview.md
index d134e6b..b93397a 100644
--- a/docs/concepts/abi_overview.md
+++ b/docs/concepts/abi_overview.md
@@ -191,9 +191,10 @@ we adopt a unified object storage format, defined as
follows:
```c++
typedef struct TVMFFIObject {
- int32_t type_index;
+ uint32_t strong_ref_count;
uint32_t weak_ref_count;
- uint64_t strong_ref_count;
+ int32_t type_index;
+ uint32_t __padding;
union {
void (*deleter)(struct TVMFFIObject* self, int flags);
int64_t __ensure_align;
@@ -203,9 +204,9 @@ typedef struct TVMFFIObject {
`TVMFFIObject` defines a common 24-byte intrusive header that all in-memory
objects share:
-- `type_index` helps us identify the type being stored, which is consistent
with `TVMFFIAny.type_index`.
-- `weak_ref_count` stores the weak atomic reference counter of the object.
- `strong_ref_count` stores the strong atomic reference counter of the object.
+- `weak_ref_count` stores the weak atomic reference counter of the object.
+- `type_index` helps us identify the type being stored, which is consistent
with `TVMFFIAny.type_index`.
- `deleter` should be called when either the strong or weak ref counter goes
to zero.
- The flags are set to indicate the event of either weak or strong going to
zero, or both.
- When `strong_ref_count` gets to zero, the deleter needs to call the
destructor of the object.
diff --git a/include/tvm/ffi/c_api.h b/include/tvm/ffi/c_api.h
index 62d9001..d2b9fab 100644
--- a/include/tvm/ffi/c_api.h
+++ b/include/tvm/ffi/c_api.h
@@ -217,22 +217,26 @@ typedef enum {
/*!
* \brief C-based type of all FFI object header that allocates on heap.
- * \note TVMFFIObject and TVMFFIAny share the common type_index header
*/
typedef struct {
+ // Ref counter goes first to align ABI with most intrusive ptr designs.
+ // It is also likely more efficient as rc operations can be quite common
+ // ABI note: Strong ref counter and weak ref counter can be packed into a
single 64-bit field
+ // Hopefully in future being able to use 64bit atomic that avoids extra
reading of
+ // weak counter during deletion.
+ /*! \brief Strong reference counter of the object. */
+ uint32_t strong_ref_count;
+ /*!
+ * \brief Weak reference counter of the object, for compatiblity with
weak_ptr design.
+ */
+ uint32_t weak_ref_count;
/*!
* \brief type index of the object.
* \note The type index of Object and Any are shared in FFI.
*/
int32_t type_index;
- /*!
- * \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;
+ /*! \brief Extra padding to ensure 8 bytes alignment. */
+ uint32_t __padding;
#if !defined(TVM_FFI_DOXYGEN_MODE)
union {
#endif
diff --git a/include/tvm/ffi/object.h b/include/tvm/ffi/object.h
index 64db8d7..e5f955c 100644
--- a/include/tvm/ffi/object.h
+++ b/include/tvm/ffi/object.h
@@ -250,7 +250,7 @@ class Object {
int32_t use_count() const {
// only need relaxed load of counters
#ifdef _MSC_VER
- return (reinterpret_cast<const volatile
__int64*>(&header_.strong_ref_count))[0]; // NOLINT(*)
+ return (reinterpret_cast<const volatile
long*>(&header_.strong_ref_count))[0]; // NOLINT(*)
#else
return __atomic_load_n(&(header_.strong_ref_count), __ATOMIC_RELAXED);
#endif
@@ -293,8 +293,8 @@ class Object {
/*! \brief increase strong reference count, the caller must already hold a
strong reference */
void IncRef() {
#ifdef _MSC_VER
- _InterlockedIncrement64(
- reinterpret_cast<volatile __int64*>(&header_.strong_ref_count)); //
NOLINT(*)
+ _InterlockedIncrement(
+ reinterpret_cast<volatile long*>(&header_.strong_ref_count)); //
NOLINT(*)
#else
__atomic_fetch_add(&(header_.strong_ref_count), 1, __ATOMIC_RELAXED);
#endif
@@ -306,12 +306,12 @@ class Object {
*/
bool TryPromoteWeakPtr() {
#ifdef _MSC_VER
- uint64_t old_count =
- (reinterpret_cast<const volatile
__int64*>(&header_.strong_ref_count))[0]; // NOLINT(*)
+ uint32_t old_count =
+ (reinterpret_cast<const volatile
long*>(&header_.strong_ref_count))[0]; // NOLINT(*)
while (old_count > 0) {
- uint64_t new_count = old_count + 1;
- uint64_t old_count_loaded = _InterlockedCompareExchange64(
- reinterpret_cast<volatile __int64*>(&header_.strong_ref_count),
new_count, old_count);
+ uint32_t new_count = old_count + 1;
+ uint32_t old_count_loaded = _InterlockedCompareExchange(
+ reinterpret_cast<volatile long*>(&header_.strong_ref_count),
new_count, old_count);
if (old_count == old_count_loaded) {
return true;
}
@@ -319,12 +319,12 @@ class Object {
}
return false;
#else
- uint64_t old_count = __atomic_load_n(&(header_.strong_ref_count),
__ATOMIC_RELAXED);
+ uint32_t old_count = __atomic_load_n(&(header_.strong_ref_count),
__ATOMIC_RELAXED);
while (old_count > 0) {
// must do CAS to ensure that we are the only one that increases the
reference count
// avoid condition when two threads tries to promote weak to strong at
same time
// or when strong deletion happens between the load and the CAS
- uint64_t new_count = old_count + 1;
+ uint32_t new_count = old_count + 1;
if (__atomic_compare_exchange_n(&(header_.strong_ref_count), &old_count,
new_count, true,
__ATOMIC_ACQ_REL, __ATOMIC_RELAXED)) {
return true;
@@ -347,8 +347,8 @@ class Object {
void DecRef() {
#ifdef _MSC_VER
// use simpler impl in windows to ensure correctness
- if (_InterlockedDecrement64(
//
- reinterpret_cast<volatile __int64*>(&header_.strong_ref_count)) ==
0) { // NOLINT(*)
+ if (_InterlockedDecrement(
//
+ reinterpret_cast<volatile long*>(&header_.strong_ref_count)) == 0)
{ // NOLINT(*)
// full barrrier is implicit in InterlockedDecrement
if (header_.deleter != nullptr) {
header_.deleter(&(this->header_),
kTVMFFIObjectDeleterFlagBitMaskStrong);
diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi
index 5c1ba1e..ff532ea 100644
--- a/python/tvm_ffi/cython/base.pxi
+++ b/python/tvm_ffi/cython/base.pxi
@@ -111,8 +111,10 @@ cdef extern from "tvm/ffi/c_api.h":
ctypedef void* TVMFFIObjectHandle
ctypedef struct TVMFFIObject:
+ uint32_t strong_ref_count
+ uint32_t weak_ref_count
int32_t type_index
- int32_t ref_counter
+ int32_t __padding
void (*deleter)(TVMFFIObject* self)
ctypedef struct TVMFFIAny:
diff --git a/tests/cpp/test_c_ffi_abi.cc b/tests/cpp/test_c_ffi_abi.cc
deleted file mode 100644
index e6c6116..0000000
--- a/tests/cpp/test_c_ffi_abi.cc
+++ /dev/null
@@ -1,31 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#include <gtest/gtest.h>
-#include <tvm/ffi/c_api.h>
-
-namespace {
-
-TEST(ABIHeaderAlignment, Default) {
- TVMFFIObject value;
- value.type_index = 10;
- EXPECT_EQ(reinterpret_cast<TVMFFIAny*>(&value)->type_index, 10);
- static_assert(sizeof(TVMFFIObject) == 24);
-}
-
-} // namespace
