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 0cfb9a94 [FFI] Make Object::unique() require no weak references (#731)
0cfb9a94 is described below
commit 0cfb9a945b0a76be870bbda6d0bb41cf4cde11c7
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 2 15:00:35 2026 -0400
[FFI] Make Object::unique() require no weak references (#731)
## Summary
- Make `Object::unique()` require one strong reference and no external
weak references by comparing the packed count with
`kCombinedRefCountBothOne`.
- Delegate `ObjectPtr<T>::unique()` to the object-level implementation
while leaving strong-only `use_count()` semantics unchanged.
- Add focused coverage for uniqueness while a `WeakObjectPtr` is live.
## Rationale
`unique()` gates copy-on-write and in-place mutation. A live weak
pointer can still upgrade while the object has a strong owner, so
strong-only uniqueness is insufficient for safely choosing in-place
mutation. Comparing the combined counter also observes both halves in
one atomic load.
## Compatibility
Out-of-tree users that retain weak references may now copy instead of
mutating in place. Current production call sites without weak references
retain their existing behavior.
---
include/tvm/ffi/object.h | 17 ++++++++++++++---
tests/cpp/test_object.cc | 18 ++++++++++++++++++
2 files changed, 32 insertions(+), 3 deletions(-)
diff --git a/include/tvm/ffi/object.h b/include/tvm/ffi/object.h
index 7a4007f7..e66d0dee 100644
--- a/include/tvm/ffi/object.h
+++ b/include/tvm/ffi/object.h
@@ -179,9 +179,20 @@ class Object {
}
/*!
- * \return Whether the object.use_count() == 1.
+ * \return Whether the object has one strong reference and no external weak
references.
+ * \note Checking both weak and strong count is needed to ensure correctness
in decisions such as
+ * copy-on-write in multi-threaded setting.
*/
- bool unique() const { return use_count() == 1; }
+ bool unique() const {
+#ifdef _MSC_VER
+ return (reinterpret_cast<const volatile uint64_t*>(
+ &header_.combined_ref_count))[0] == // NOLINT(*)
+ kCombinedRefCountBothOne;
+#else
+ return __atomic_load_n(&(header_.combined_ref_count), __ATOMIC_RELAXED) ==
+ kCombinedRefCountBothOne;
+#endif
+ }
/*!
* \return The usage count of the cell.
@@ -496,7 +507,7 @@ class ObjectPtr {
/*! \return The use count of the ptr, for debug purposes */
int use_count() const { return data_ != nullptr ? data_->use_count() : 0; }
/*! \return whether the reference is unique */
- bool unique() const { return data_ != nullptr && data_->use_count() == 1; }
+ bool unique() const { return data_ != nullptr && data_->unique(); }
/*! \return Whether two ObjectPtr do not equal each other */
bool operator==(const ObjectPtr<T>& other) const { return data_ ==
other.data_; }
/*! \return Whether two ObjectPtr equals each other */
diff --git a/tests/cpp/test_object.cc b/tests/cpp/test_object.cc
index 0c4b066a..a34efd98 100644
--- a/tests/cpp/test_object.cc
+++ b/tests/cpp/test_object.cc
@@ -353,6 +353,24 @@ TEST(Object, WeakObjectPtr) {
EXPECT_TRUE(expired_lock == nullptr);
}
+TEST(Object, UniqueWithWeakObjectPtr) {
+ ObjectPtr<TIntObj> strong_ptr = make_object<TIntObj>(42);
+
+ EXPECT_TRUE(strong_ptr.unique());
+ EXPECT_TRUE(strong_ptr->unique());
+ EXPECT_EQ(strong_ptr.use_count(), 1);
+
+ WeakObjectPtr<TIntObj> weak_ptr(strong_ptr);
+ EXPECT_FALSE(strong_ptr.unique());
+ EXPECT_FALSE(strong_ptr->unique());
+ EXPECT_EQ(strong_ptr.use_count(), 1);
+
+ weak_ptr.reset();
+ EXPECT_TRUE(strong_ptr.unique());
+ EXPECT_TRUE(strong_ptr->unique());
+ EXPECT_EQ(strong_ptr.use_count(), 1);
+}
+
TEST(Object, WeakObjectPtrAssignment) {
// Test copy construction
ObjectPtr<TIntObj> new_strong = make_object<TIntObj>(100);
