Skip to content

[FFI] Optimize atomic decref in Object #18077

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 18, 2025
Merged

[FFI] Optimize atomic decref in Object #18077

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 0 additions & 28 deletions ffi/include/tvm/ffi/base_details.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,34 +139,6 @@ namespace tvm {
namespace ffi {
namespace details {

/********** Atomic Operations *********/

TVM_FFI_INLINE int32_t AtomicIncrementRelaxed(int32_t* ptr) {
#ifdef _MSC_VER
return _InterlockedIncrement(reinterpret_cast<volatile long*>(ptr)) - 1; // NOLINT(*)
#else
return __atomic_fetch_add(ptr, 1, __ATOMIC_RELAXED);
#endif
}

TVM_FFI_INLINE int32_t AtomicDecrementRelAcq(int32_t* ptr) {
#ifdef _MSC_VER
return _InterlockedDecrement(reinterpret_cast<volatile long*>(ptr)) + 1; // NOLINT(*)
#else
return __atomic_fetch_sub(ptr, 1, __ATOMIC_ACQ_REL);
#endif
}

TVM_FFI_INLINE int32_t AtomicLoadRelaxed(const int32_t* ptr) {
int32_t* raw_ptr = const_cast<int32_t*>(ptr);
#ifdef _MSC_VER
// simply load the variable ptr out
return (reinterpret_cast<const volatile long*>(raw_ptr))[0]; // NOLINT(*)
#else
return __atomic_load_n(raw_ptr, __ATOMIC_RELAXED);
#endif
}

// for each iterator
template <bool stop, std::size_t I, typename F>
struct for_each_dispatcher {
Expand Down
33 changes: 30 additions & 3 deletions ffi/include/tvm/ffi/object.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,14 @@ class Object {
* \return The usage count of the cell.
* \note We use stl style naming to be consistent with known API in shared_ptr.
*/
int32_t use_count() const { return details::AtomicLoadRelaxed(&(header_.ref_counter)); }
int32_t use_count() const {
// only need relaxed load of counters
#ifdef _MSC_VER
return (reinterpret_cast<const volatile long*>(&header_.ref_counter))[0]; // NOLINT(*)
#else
return __atomic_load_n(&(header_.ref_counter), __ATOMIC_RELAXED);
#endif
}

// Information about the object
static constexpr const char* _type_key = "object.Object";
Expand Down Expand Up @@ -220,15 +227,35 @@ class Object {

private:
/*! \brief increase reference count */
void IncRef() { details::AtomicIncrementRelaxed(&(header_.ref_counter)); }
void IncRef() {
#ifdef _MSC_VER
_InterlockedIncrement(reinterpret_cast<volatile long*>(&header_.ref_counter)); // NOLINT(*)
#else
__atomic_fetch_add(&(header_.ref_counter), 1, __ATOMIC_RELAXED);
#endif
}

/*! \brief decrease reference count and delete the object */
void DecRef() {
if (details::AtomicDecrementRelAcq(&(header_.ref_counter)) == 1) {
#ifdef _MSC_VER
if (_InterlockedDecrement( //
reinterpret_cast<volatile long*>(&header_.ref_counter)) == 0) { // NOLINT(*)
// full barrrier is implicit in InterlockedDecrement
if (header_.deleter != nullptr) {
header_.deleter(&(this->header_));
}
}
#else
// first do a release, note we only need to acquire for deleter
if (__atomic_fetch_sub(&(header_.ref_counter), 1, __ATOMIC_RELEASE) == 1) {
// only acquire when we need to call deleter
// in this case we need to ensure all previous writes are visible
__atomic_thread_fence(__ATOMIC_ACQUIRE);
if (header_.deleter != nullptr) {
header_.deleter(&(this->header_));
}
}
#endif
}

// friend classes
Expand Down
Loading