diff --git a/ffi/include/tvm/ffi/base_details.h b/ffi/include/tvm/ffi/base_details.h index afc1b30ad66d..ed147b131804 100644 --- a/ffi/include/tvm/ffi/base_details.h +++ b/ffi/include/tvm/ffi/base_details.h @@ -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(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(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(ptr); -#ifdef _MSC_VER - // simply load the variable ptr out - return (reinterpret_cast(raw_ptr))[0]; // NOLINT(*) -#else - return __atomic_load_n(raw_ptr, __ATOMIC_RELAXED); -#endif -} - // for each iterator template struct for_each_dispatcher { diff --git a/ffi/include/tvm/ffi/object.h b/ffi/include/tvm/ffi/object.h index 72e6f0a1f8ec..21dd8f3f7fa7 100644 --- a/ffi/include/tvm/ffi/object.h +++ b/ffi/include/tvm/ffi/object.h @@ -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(&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"; @@ -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(&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(&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