123 KiB
123 KiB
<html lang="en">
<head>
</head>
</html>
LCOV - code coverage report | ||||||||||||||||||||||
![]() | ||||||||||||||||||||||
|
||||||||||||||||||||||
![]() |
Line data Source code 1 : #pragma once 2 : 3 : #include <c10/util/Exception.h> 4 : #include <c10/util/MaybeOwned.h> 5 : #include <atomic> 6 : #include <climits> 7 : #include <memory> 8 : 9 : namespace pybind11 { 10 : template <typename, typename...> 11 : class class_; 12 : } 13 : 14 : namespace c10 { 15 : class intrusive_ptr_target; 16 : namespace raw { 17 : namespace weak_intrusive_ptr { 18 : inline void incref(intrusive_ptr_target* self); 19 : } 20 : namespace intrusive_ptr { 21 : inline void incref(intrusive_ptr_target* self); 22 : } 23 : 24 : // constructor tag used by intrusive_ptr constructors 25 : struct DontIncreaseRefcount {}; 26 : } // namespace raw 27 : /** 28 : * intrusive_ptr<T> is an alternative to shared_ptr<T> that has better 29 : * performance because it does the refcounting intrusively 30 : * (i.e. in a member of the object itself). 31 : * Your class T needs to inherit from intrusive_ptr_target to allow it to be 32 : * used in an intrusive_ptr<T>. Your class's constructor should not allow 33 : *`this` to escape to other threads or create an intrusive_ptr from `this`. 34 : */ 35 : 36 : // Note [Stack allocated intrusive_ptr_target safety] 37 : // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 38 : // A well known problem with std::enable_shared_from_this is that it 39 : // allows you to create a std::shared_ptr from a stack allocated object, 40 : // which is totally bogus because the object will die once you return 41 : // from the stack. In intrusive_ptr, we can detect that this has occurred, 42 : // because we set the refcount/weakcount of objects which inherit from 43 : // intrusive_ptr_target to zero, *unless* we can prove that the object 44 : // was dynamically allocated (e.g., via make_intrusive). 45 : // 46 : // Thus, whenever you transmute a T* into a intrusive_ptr<T>, we check 47 : // and make sure that the refcount isn't zero (or, a more subtle 48 : // test for weak_intrusive_ptr<T>, for which the refcount may validly 49 : // be zero, but the weak refcount better not be zero), because that 50 : // tells us if the object was allocated by us. If it wasn't, no 51 : // intrusive_ptr for you! 52 : 53 : // NOLINTNEXTLINE(cppcoreguidelines-virtual-class-destructor) 54 : class C10_API intrusive_ptr_target { 55 : // Note [Weak references for intrusive refcounting] 56 : // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 57 : // Here's the scheme: 58 : // 59 : // - refcount == number of strong references to the object 60 : // weakcount == number of weak references to the object, 61 : // plus one more if refcount > 0 62 : // An invariant: refcount > 0 => weakcount > 0 63 : // 64 : // - c10::StorageImpl stays live as long as there are any strong 65 : // or weak pointers to it (weakcount > 0, since strong 66 : // references count as a +1 to weakcount) 67 : // 68 : // - finalizers are called and data_ptr is deallocated when refcount == 0 69 : // 70 : // - Once refcount == 0, it can never again be > 0 (the transition 71 : // from > 0 to == 0 is monotonic) 72 : // 73 : // - When you access c10::StorageImpl via a weak pointer, you must 74 : // atomically increment the use count, if it is greater than 0. 75 : // If it is not, you must report that the storage is dead. 76 : // 77 : mutable std::atomic<size_t> refcount_; 78 : mutable std::atomic<size_t> weakcount_; 79 : 80 : template <typename T, typename NullType> 81 : friend class intrusive_ptr; 82 : friend inline void raw::intrusive_ptr::incref(intrusive_ptr_target* self); 83 : 84 : template <typename T, typename NullType> 85 : friend class weak_intrusive_ptr; 86 : friend inline void raw::weak_intrusive_ptr::incref( 87 : intrusive_ptr_target* self); 88 : 89 : template <typename T> 90 : friend struct ExclusivelyOwnedTensorTraits; 91 : 92 : protected: 93 : // protected destructor. We never want to destruct intrusive_ptr_target* 94 : // directly. 95 520215158 : virtual ~intrusive_ptr_target() { 96 : // Disable -Wterminate and -Wexceptions so we're allowed to use assertions 97 : // (i.e. throw exceptions) in a destructor. 98 : // We also have to disable -Wunknown-warning-option and -Wpragmas, because 99 : // some other compilers don't know about -Wterminate or -Wexceptions and 100 : // will show a warning about unknown warning options otherwise. 101 : #if defined(_MSC_VER) && !defined(__clang__) 102 : #pragma warning(push) 103 : #pragma warning( \ 104 : disable : 4297) // function assumed not to throw an exception but does 105 : #else 106 : #pragma GCC diagnostic push 107 : #pragma GCC diagnostic ignored "-Wpragmas" 108 : #pragma GCC diagnostic ignored "-Wunknown-warning-option" 109 : #pragma GCC diagnostic ignored "-Wterminate" 110 : #pragma GCC diagnostic ignored "-Wexceptions" 111 : #endif 112 1040430316 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 113 : // Second condition is there to accommodate 114 : // unsafe_adapt_non_heap_allocated: since we are doing our own 115 : // deallocation in that case, it is correct for each 116 : // expected_decref to have happened (some user code tried to 117 : // decref and thus free the object, but it didn't happen right 118 : // away) or not (no user code tried to free the object, and 119 : // now it's getting destroyed through whatever mechanism the 120 : // caller of unsafe_adapt_non_heap_allocated wanted to 121 : // use). We choose our reference count such that the count 122 : // will not dip below INT_MAX regardless. 123 : refcount_.load() == 0 || refcount_.load() >= INT_MAX, 124 : "Tried to destruct an intrusive_ptr_target that still has intrusive_ptr to it; refcount was ", 125 : refcount_.load()); 126 1040430316 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 127 : // See ~intrusive_ptr for optimization that will frequently result in 1 128 : // at destruction time. 129 : weakcount_.load() == 1 || weakcount_.load() == 0 || 130 : weakcount_.load() == INT_MAX - 1 || weakcount_.load() == INT_MAX, 131 : "Tried to destruct an intrusive_ptr_target that still has weak_intrusive_ptr to it"); 132 : #if defined(_MSC_VER) && !defined(__clang__) 133 : #pragma warning(pop) 134 : #else 135 : #pragma GCC diagnostic pop 136 : #endif 137 520215158 : } 138 : 139 378768 : constexpr intrusive_ptr_target() noexcept : refcount_(0), weakcount_(0) {} 140 : 141 : // intrusive_ptr_target supports copy and move: but refcount and weakcount 142 : // don't participate (since they are intrinsic properties of the memory 143 : // location) 144 : intrusive_ptr_target(intrusive_ptr_target&& /*other*/) noexcept 145 : : intrusive_ptr_target() {} 146 : 147 : intrusive_ptr_target& operator=(intrusive_ptr_target&& /*other*/) noexcept { 148 : return *this; 149 : } 150 : 151 : intrusive_ptr_target(const intrusive_ptr_target& /*other*/) noexcept 152 : : intrusive_ptr_target() {} 153 : 154 : intrusive_ptr_target& operator=( 155 : const intrusive_ptr_target& /*other*/) noexcept { 156 : return *this; 157 : } 158 : 159 : private: 160 : /** 161 : * This is called when refcount reaches zero. 162 : * You can override this to release expensive resources. 163 : * There might still be weak references, so your object might not get 164 : * destructed yet, but you can assume the object isn't used anymore, 165 : * i.e. no more calls to methods or accesses to members (we just can't 166 : * destruct it yet because we need the weakcount accessible). 167 : * 168 : * If there are no weak references (i.e. your class is about to be 169 : * destructed), this function WILL NOT be called. 170 : */ 171 0 : virtual void release_resources() {} 172 : }; 173 : 174 : namespace detail { 175 : template <class TTarget> 176 : struct intrusive_target_default_null_type final { 177 149296920 : static constexpr TTarget* singleton() noexcept { 178 149296920 : return nullptr; 179 : } 180 : }; 181 : 182 : template <class TTarget, class ToNullType, class FromNullType> 183 378768 : TTarget* assign_ptr_(TTarget* rhs) { 184 378768 : if (FromNullType::singleton() == rhs) { 185 0 : return ToNullType::singleton(); 186 : } else { 187 378768 : return rhs; 188 : } 189 : } 190 : 191 : // Increment needs to be acquire-release to make use_count() and 192 : // unique() reliable. 193 440756 : inline size_t atomic_refcount_increment(std::atomic<size_t>& refcount) { 194 881512 : return refcount.fetch_add(1, std::memory_order_acq_rel) + 1; 195 : } 196 : 197 : // weak_use_count() is only used for testing, so we don't need it to 198 : // be reliable. Relaxed should be fine. 199 : inline size_t atomic_weakcount_increment(std::atomic<size_t>& weakcount) { 200 : return weakcount.fetch_add(1, std::memory_order_relaxed) + 1; 201 : } 202 : 203 : // Both decrements need to be acquire-release for correctness. See 204 : // e.g. std::shared_ptr implementation. 205 212684736 : inline size_t atomic_refcount_decrement(std::atomic<size_t>& refcount) { 206 425369472 : return refcount.fetch_sub(1, std::memory_order_acq_rel) - 1; 207 : } 208 : 209 0 : inline size_t atomic_weakcount_decrement(std::atomic<size_t>& weakcount) { 210 0 : return weakcount.fetch_sub(1, std::memory_order_acq_rel) - 1; 211 : } 212 : 213 : } // namespace detail 214 : 215 : template <class TTarget, class NullType> 216 : class weak_intrusive_ptr; 217 : 218 : template < 219 : class TTarget, 220 : class NullType = detail::intrusive_target_default_null_type<TTarget>> 221 : class intrusive_ptr final { 222 : private: 223 : // the following static assert would be nice to have but it requires 224 : // the target class T to be fully defined when intrusive_ptr<T> is instantiated 225 : // this is a problem for classes that contain pointers to themselves 226 : // static_assert( 227 : // std::is_base_of<intrusive_ptr_target, TTarget>::value, 228 : // "intrusive_ptr can only be used for classes that inherit from 229 : // intrusive_ptr_target."); 230 : #ifndef _WIN32 231 : // This static_assert triggers on MSVC 232 : // error C2131: expression did not evaluate to a constant 233 : static_assert( 234 : NullType::singleton() == NullType::singleton(), 235 : "NullType must have a constexpr singleton() method"); 236 : #endif 237 : static_assert( 238 : std::is_base_of< 239 : TTarget, 240 : typename std::remove_pointer<decltype(NullType::singleton())>::type>:: 241 : value, 242 : "NullType::singleton() must return a element_type* pointer"); 243 : 244 : TTarget* target_; 245 : 246 : template <typename T> 247 : friend struct ExclusivelyOwnedTensorTraits; 248 : template <class TTarget2, class NullType2> 249 : friend class intrusive_ptr; 250 : friend class weak_intrusive_ptr<TTarget, NullType>; 251 : 252 : // Make pybind11::class_ be a friend class of intrusive_ptr, so that custom 253 : // smart holder in pybind11 could access the private constructor of 254 : // intrusive_ptr(T*) which took the ownership of the object. This is required 255 : // by customer holder macro PYBIND11_DECLARE_HOLDER_TYPE, where it uses 256 : // intrusive_ptr(TTarget*) to initialize and take ownership of the object. For 257 : // details, see 258 : // https://pybind11.readthedocs.io/en/stable/advanced/smart_ptrs.html#custom-smart-pointers 259 : template <typename, typename...> 260 : friend class pybind11::class_; 261 : 262 441642 : void retain_() { 263 441642 : if (target_ != NullType::singleton()) { 264 : size_t new_refcount = 265 440756 : detail::atomic_refcount_increment(target_->refcount_); 266 440756 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 267 : new_refcount != 1, 268 : "intrusive_ptr: Cannot increase refcount after it reached zero."); 269 : } 270 441642 : } 271 : 272 710681430 : void reset_() noexcept { 273 923366166 : if (target_ != NullType::singleton() && 274 212684736 : detail::atomic_refcount_decrement(target_->refcount_) == 0) { 275 : // See comment above about weakcount. As long as refcount>0, 276 : // weakcount is one larger than the actual number of weak references. 277 : // So we need to decrement it here. 278 211864418 : bool should_delete = 279 211864418 : target_->weakcount_.load(std::memory_order_acquire) == 1; 280 211864418 : if (!should_delete) { 281 : // justification for const_cast: release_resources is basically a 282 : // destructor and a destructor always mutates the object, even for const 283 : // objects. NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) 284 0 : const_cast<std::remove_const_t<TTarget>*>(target_)->release_resources(); 285 0 : should_delete = 286 0 : detail::atomic_weakcount_decrement(target_->weakcount_) == 0; 287 : } 288 211864418 : if (should_delete) { 289 211864418 : delete target_; 290 : } 291 : } 292 710681430 : } 293 : 294 : // raw pointer constructors are not public because we shouldn't make 295 : // intrusive_ptr out of raw pointers except from inside the make_intrusive(), 296 : // reclaim() and weak_intrusive_ptr::lock() implementations. 297 : 298 : // This constructor will increase the ref counter for you. 299 : // This constructor will be used by the make_intrusive(), and also pybind11, 300 : // which wrap the intrusive_ptr holder around the raw pointer and incref 301 : // correspondingly (pybind11 requires raw pointer constructor to incref by 302 : // default). 303 37304192 : explicit intrusive_ptr(TTarget* target) 304 37304192 : : intrusive_ptr(target, raw::DontIncreaseRefcount{}) { 305 37304192 : if (target_ != NullType::singleton()) { 306 : // We just created result.target_, so we know no other thread has 307 : // access to it, so we know we needn't care about memory ordering. 308 : // (On x86_64, a store with memory_order_relaxed generates a plain old 309 : // `mov`, whereas an atomic increment does a lock-prefixed `add`, which is 310 : // much more expensive: https://godbolt.org/z/eKPzj8.) 311 37304192 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 312 : target_->refcount_ == 0 && target_->weakcount_ == 0, 313 : "intrusive_ptr: Newly-created target had non-zero refcounts. Does its " 314 : "constructor do something strange like incref or create an " 315 : "intrusive_ptr from `this`?"); 316 37304192 : target_->refcount_.store(1, std::memory_order_relaxed); 317 37304192 : target_->weakcount_.store(1, std::memory_order_relaxed); 318 : } 319 37304192 : } 320 : 321 : public: 322 : using element_type = TTarget; 323 : 324 37370160 : intrusive_ptr() noexcept 325 37370160 : : intrusive_ptr(NullType::singleton(), raw::DontIncreaseRefcount{}) {} 326 : 327 : intrusive_ptr(std::nullptr_t) noexcept 328 : : intrusive_ptr(NullType::singleton(), raw::DontIncreaseRefcount{}) {} 329 : 330 : // This constructor will not increase the ref counter for you. 331 : // We use the tagged dispatch mechanism to explicitly mark this constructor 332 : // to not increase the refcount 333 179950728 : explicit intrusive_ptr(TTarget* target, raw::DontIncreaseRefcount) noexcept 334 179950728 : : target_(target) {} 335 : 336 : explicit intrusive_ptr(std::unique_ptr<TTarget> rhs) noexcept 337 : : intrusive_ptr(rhs.release()) {} 338 : 339 772197404 : intrusive_ptr(intrusive_ptr&& rhs) noexcept : target_(rhs.target_) { 340 772197404 : rhs.target_ = NullType::singleton(); 341 772197404 : } 342 : 343 : template <class From, class FromNullType> 344 : /* implicit */ intrusive_ptr(intrusive_ptr<From, FromNullType>&& rhs) noexcept 345 : : target_( 346 : detail::assign_ptr_<TTarget, NullType, FromNullType>(rhs.target_)) { 347 : static_assert( 348 : std::is_convertible<From*, TTarget*>::value, 349 : "Type mismatch. intrusive_ptr move constructor got pointer of wrong type."); 350 : rhs.target_ = FromNullType::singleton(); 351 : } 352 : 353 62874 : intrusive_ptr(const intrusive_ptr& rhs) : target_(rhs.target_) { 354 62874 : retain_(); 355 62874 : } 356 : 357 : template <class From, class FromNullType> 358 378768 : /* implicit */ intrusive_ptr(const intrusive_ptr<From, FromNullType>& rhs) 359 378768 : : target_( 360 378768 : detail::assign_ptr_<TTarget, NullType, FromNullType>(rhs.target_)) { 361 : static_assert( 362 : std::is_convertible<From*, TTarget*>::value, 363 : "Type mismatch. intrusive_ptr copy constructor got pointer of wrong type."); 364 378768 : retain_(); 365 378768 : } 366 : 367 710681430 : ~intrusive_ptr() noexcept { 368 710681430 : reset_(); 369 710681430 : } 370 : 371 435080 : intrusive_ptr& operator=(intrusive_ptr&& rhs) & noexcept { 372 435080 : return operator=<TTarget, NullType>(std::move(rhs)); 373 : } 374 : 375 : template <class From, class FromNullType> 376 435080 : intrusive_ptr& operator=(intrusive_ptr<From, FromNullType>&& rhs) & noexcept { 377 : static_assert( 378 : std::is_convertible<From*, TTarget*>::value, 379 : "Type mismatch. intrusive_ptr move assignment got pointer of wrong type."); 380 435080 : intrusive_ptr tmp = std::move(rhs); 381 435080 : swap(tmp); 382 870160 : return *this; 383 435080 : } 384 : 385 1820 : intrusive_ptr& operator=(const intrusive_ptr& rhs) & noexcept { 386 1820 : return operator=<TTarget, NullType>(rhs); 387 : } 388 : 389 : template <class From, class FromNullType> 390 1820 : intrusive_ptr& operator=(const intrusive_ptr<From, NullType>& rhs) & { 391 : static_assert( 392 : std::is_convertible<From*, TTarget*>::value, 393 : "Type mismatch. intrusive_ptr copy assignment got pointer of wrong type."); 394 1820 : intrusive_ptr tmp = rhs; 395 1820 : swap(tmp); 396 1820 : return *this; 397 1820 : } 398 : 399 105695220 : TTarget* get() const noexcept { 400 105695220 : return target_; 401 : } 402 : 403 0 : TTarget& operator*() const noexcept { 404 0 : return *target_; 405 : } 406 : 407 121964308 : TTarget* operator->() const noexcept { 408 : // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDelete) 409 121964308 : return target_; 410 : } 411 : 412 466880 : operator bool() const noexcept { 413 466880 : return target_ != NullType::singleton(); 414 : } 415 : 416 : void reset() noexcept { 417 : reset_(); 418 : target_ = NullType::singleton(); 419 : } 420 : 421 436900 : void swap(intrusive_ptr& rhs) noexcept { 422 436900 : TTarget* tmp = target_; 423 436900 : target_ = rhs.target_; 424 436900 : rhs.target_ = tmp; 425 436900 : } 426 : 427 : // We do a lot of null-pointer checks in our code, good to have this be cheap. 428 : bool defined() const noexcept { 429 : return target_ != NullType::singleton(); 430 : } 431 : 432 458920 : size_t use_count() const noexcept { 433 458920 : if (target_ == NullType::singleton()) { 434 0 : return 0; 435 : } 436 917840 : return target_->refcount_.load(std::memory_order_acquire); 437 : } 438 : 439 : size_t weak_use_count() const noexcept { 440 : if (target_ == NullType::singleton()) { 441 : return 0; 442 : } 443 : return target_->weakcount_.load(std::memory_order_acquire); 444 : } 445 : 446 : bool unique() const noexcept { 447 : return use_count() == 1; 448 : } 449 : 450 : /** 451 : * Returns an owning (!) pointer to the underlying object and makes the 452 : * intrusive_ptr instance invalid. That means the refcount is not decreased. 453 : * You *must* put the returned pointer back into a intrusive_ptr using 454 : * intrusive_ptr::reclaim(ptr) to properly destruct it. 455 : * This is helpful for C APIs. 456 : */ 457 : TTarget* release() noexcept { 458 : // NOLINTNEXTLINE(clang-analyzer-core.uninitialized.Assign) 459 : TTarget* result = target_; 460 : target_ = NullType::singleton(); 461 : return result; 462 : } 463 : 464 : /** 465 : * Takes an owning pointer to TTarget* and creates an intrusive_ptr that takes 466 : * over ownership. That means the refcount is not increased. 467 : * This is the counter-part to intrusive_ptr::release() and the pointer 468 : * passed in *must* have been created using intrusive_ptr::release(). 469 : */ 470 105276376 : static intrusive_ptr reclaim(TTarget* owning_ptr) { 471 315829128 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 472 : owning_ptr == NullType::singleton() || 473 : owning_ptr->refcount_.load() == 0 || owning_ptr->weakcount_.load(), 474 : "TTarget violates the invariant that refcount > 0 => weakcount > 0"); 475 105276376 : return intrusive_ptr(owning_ptr, raw::DontIncreaseRefcount{}); 476 : } 477 : 478 : /** 479 : * Takes an owning pointer to TTarget* and creates an intrusive_ptr 480 : * representing a new reference, i.e. the raw pointer retains 481 : * ownership. 482 : */ 483 : static intrusive_ptr reclaim_copy(TTarget* owning_ptr) { 484 : auto ret = reclaim(owning_ptr); 485 : ret.retain_(); 486 : return ret; 487 : } 488 : 489 : /** 490 : * Allocate a heap object with args and wrap it inside a intrusive_ptr and 491 : * incref. This is a helper function to let make_intrusive() access private 492 : * intrusive_ptr constructors. 493 : */ 494 : template <class... Args> 495 37304192 : static intrusive_ptr make(Args&&... args) { 496 37304192 : return intrusive_ptr(new TTarget(std::forward<Args>(args)...)); 497 : } 498 : 499 : /** 500 : * Turn a new instance of TTarget (e.g., literally allocated 501 : * using new TTarget(...) into an intrusive_ptr. If possible, 502 : * use intrusive_ptr::make instead which statically guarantees 503 : * that the allocation was done properly. 504 : * 505 : * At the moment, the only reason this method exists is because 506 : * pybind11 holder types expect to be able to allocate in 507 : * this way (because pybind11 handles the new allocation itself). 508 : */ 509 : static intrusive_ptr unsafe_steal_from_new(TTarget* raw_ptr) { 510 : return intrusive_ptr(raw_ptr); 511 : } 512 : 513 : /** 514 : * Turn an instance of TTarget that should not be reference counted 515 : * (e.g., allocated into an arena with placement new) into an 516 : * intrusive_ptr. This is gratuitously unsafe and should only be 517 : * used if you can guarantee that the pointer will not escape and be 518 : * refcounted as normal. 519 : * 520 : * `expected_decrefs` is a debugging parameter: it indicates the 521 : * number of strong owners the intrusive_ptr_target in question is 522 : * expected to get. In most use cases, this will likely be 1. 523 : * 524 : * The reason this method exists is for manually sharing 525 : * StorageImpls across Tensors in the static runtime. It needs 526 : * access to private intrusive_ptr members so that the refcounts can 527 : * be initialized to custom values. 528 : */ 529 : static intrusive_ptr unsafe_adapt_non_heap_allocated( 530 : TTarget* raw_ptr, 531 : size_t expected_decrefs) { 532 : intrusive_ptr result(raw_ptr, raw::DontIncreaseRefcount{}); 533 : // INT_MAX is impractically huge for a reference count, while 534 : // being in no danger of overflowing size_t. We actually only need to 535 : // initialize the refcount to 2 -- we are just doing an unbalanced 536 : // incref to prevent the non-heap-allocated target from being 537 : // freed, and we are optimizing that incref by directly 538 : // initializing the refcounts rather than doing an expensive 539 : // atomic increment. The reason to use INT_MAX is to accommodate 540 : // the debug assertions in ~intrusive_ptr_target. 541 : #ifdef NDEBUG 542 : expected_decrefs = 0; 543 : #endif 544 : result.target_->refcount_.store( 545 : INT_MAX + expected_decrefs, std::memory_order_relaxed); 546 : result.target_->weakcount_.store(INT_MAX, std::memory_order_relaxed); 547 : return result; 548 : } 549 : 550 : /** 551 : * Turn a **non-owning raw pointer** to an intrusive_ptr. It is 552 : * the moral equivalent of enable_shared_from_this on a shared pointer. 553 : * 554 : * This method is only valid for objects that are already live. If 555 : * you are looking for the moral equivalent of unique_ptr<T>(T*) 556 : * constructor, see steal_from_new. 557 : * 558 : * TODO: https://github.com/pytorch/pytorch/issues/56482 559 : */ 560 : static intrusive_ptr unsafe_reclaim_from_nonowning(TTarget* raw_ptr) { 561 : // See Note [Stack allocated intrusive_ptr_target safety] 562 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 563 : raw_ptr == NullType::singleton() || raw_ptr->refcount_.load() > 0, 564 : "intrusive_ptr: Can only reclaim pointers that are owned by someone"); 565 : auto ptr = reclaim(raw_ptr); // doesn't increase refcount 566 : ptr.retain_(); 567 : return ptr; 568 : } 569 : }; 570 : 571 : template < 572 : class TTarget, 573 : class NullType = detail::intrusive_target_default_null_type<TTarget>, 574 : class... Args> 575 37304192 : inline intrusive_ptr<TTarget, NullType> make_intrusive(Args&&... args) { 576 37304192 : return intrusive_ptr<TTarget, NullType>::make(std::forward<Args>(args)...); 577 : } 578 : 579 : template <class TTarget, class NullType> 580 : inline void swap( 581 : intrusive_ptr<TTarget, NullType>& lhs, 582 : intrusive_ptr<TTarget, NullType>& rhs) noexcept { 583 : lhs.swap(rhs); 584 : } 585 : 586 : // To allow intrusive_ptr inside std::map or std::set, we need operator< 587 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 588 : inline bool operator<( 589 : const intrusive_ptr<TTarget1, NullType1>& lhs, 590 : const intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 591 : return lhs.get() < rhs.get(); 592 : } 593 : 594 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 595 : inline bool operator==( 596 : const intrusive_ptr<TTarget1, NullType1>& lhs, 597 : const intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 598 : return lhs.get() == rhs.get(); 599 : } 600 : 601 : template <class TTarget1, class NullType1> 602 : inline bool operator==( 603 : const intrusive_ptr<TTarget1, NullType1>& lhs, 604 : std::nullptr_t) noexcept { 605 : return lhs.get() == nullptr; 606 : } 607 : 608 : template <class TTarget2, class NullType2> 609 : inline bool operator==( 610 : std::nullptr_t, 611 : const intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 612 : return nullptr == rhs.get(); 613 : } 614 : 615 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 616 : inline bool operator!=( 617 : const intrusive_ptr<TTarget1, NullType1>& lhs, 618 : const intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 619 : return !operator==(lhs, rhs); 620 : } 621 : 622 : template <class TTarget1, class NullType1> 623 : inline bool operator!=( 624 : const intrusive_ptr<TTarget1, NullType1>& lhs, 625 : std::nullptr_t) noexcept { 626 : return !operator==(lhs, nullptr); 627 : } 628 : 629 : template <class TTarget2, class NullType2> 630 : inline bool operator!=( 631 : std::nullptr_t, 632 : const intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 633 : return !operator==(nullptr, rhs); 634 : } 635 : template <typename T> 636 : struct MaybeOwnedTraits<c10::intrusive_ptr<T>> { 637 : using owned_type = c10::intrusive_ptr<T>; 638 : using borrow_type = c10::intrusive_ptr<T>; 639 : 640 : static borrow_type createBorrow(const owned_type& from) { 641 : return borrow_type::reclaim(from.get()); 642 : } 643 : 644 : static void assignBorrow(borrow_type& lhs, const borrow_type& rhs) { 645 : lhs.release(); 646 : lhs = borrow_type::reclaim(rhs.get()); 647 : } 648 : 649 : static void destroyBorrow(borrow_type& toDestroy) { 650 : toDestroy.release(); 651 : } 652 : 653 : static const owned_type& referenceFromBorrow(const borrow_type& borrow) { 654 : return borrow; 655 : } 656 : 657 : static const owned_type* pointerFromBorrow(const borrow_type& borrow) { 658 : return &borrow; 659 : } 660 : 661 : static bool debugBorrowIsValid(const borrow_type& /*borrow*/) { 662 : return true; 663 : } 664 : }; 665 : 666 : template < 667 : typename TTarget, 668 : class NullType = detail::intrusive_target_default_null_type<TTarget>> 669 : class weak_intrusive_ptr final { 670 : private: 671 : static_assert( 672 : std::is_base_of<intrusive_ptr_target, TTarget>::value, 673 : "intrusive_ptr can only be used for classes that inherit from intrusive_ptr_target."); 674 : #ifndef _WIN32 675 : // This static_assert triggers on MSVC 676 : // error C2131: expression did not evaluate to a constant 677 : static_assert( 678 : NullType::singleton() == NullType::singleton(), 679 : "NullType must have a constexpr singleton() method"); 680 : #endif 681 : static_assert( 682 : std::is_base_of< 683 : TTarget, 684 : typename std::remove_pointer<decltype(NullType::singleton())>::type>:: 685 : value, 686 : "NullType::singleton() must return a element_type* pointer"); 687 : 688 : TTarget* target_; 689 : 690 : template <class TTarget2, class NullType2> 691 : friend class weak_intrusive_ptr; 692 : 693 : void retain_() { 694 : if (target_ != NullType::singleton()) { 695 : size_t new_weakcount = 696 : detail::atomic_weakcount_increment(target_->weakcount_); 697 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 698 : new_weakcount != 1, 699 : "weak_intrusive_ptr: Cannot increase weakcount after it reached zero."); 700 : } 701 : } 702 : 703 : void reset_() noexcept { 704 : if (target_ != NullType::singleton() && 705 : detail::atomic_weakcount_decrement(target_->weakcount_) == 0) { 706 : // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDelete) 707 : delete target_; 708 : } 709 : target_ = NullType::singleton(); 710 : } 711 : 712 : constexpr explicit weak_intrusive_ptr(TTarget* target) : target_(target) {} 713 : 714 : public: 715 : using element_type = TTarget; 716 : 717 : explicit weak_intrusive_ptr(const intrusive_ptr<TTarget, NullType>& ptr) 718 : : weak_intrusive_ptr(ptr.get()) { 719 : retain_(); 720 : } 721 : 722 : weak_intrusive_ptr(weak_intrusive_ptr&& rhs) noexcept : target_(rhs.target_) { 723 : rhs.target_ = NullType::singleton(); 724 : } 725 : 726 : template <class From, class FromNullType> 727 : /* implicit */ weak_intrusive_ptr( 728 : weak_intrusive_ptr<From, FromNullType>&& rhs) noexcept 729 : : target_( 730 : detail::assign_ptr_<TTarget, NullType, FromNullType>(rhs.target_)) { 731 : static_assert( 732 : std::is_convertible<From*, TTarget*>::value, 733 : "Type mismatch. weak_intrusive_ptr move constructor got pointer of wrong type."); 734 : rhs.target_ = FromNullType::singleton(); 735 : } 736 : 737 : weak_intrusive_ptr(const weak_intrusive_ptr& rhs) : target_(rhs.target_) { 738 : retain_(); 739 : } 740 : 741 : template <class From, class FromNullType> 742 : /* implicit */ weak_intrusive_ptr( 743 : const weak_intrusive_ptr<From, FromNullType>& rhs) 744 : : target_( 745 : detail::assign_ptr_<TTarget, NullType, FromNullType>(rhs.target_)) { 746 : static_assert( 747 : std::is_convertible<From*, TTarget*>::value, 748 : "Type mismatch. weak_intrusive_ptr copy constructor got pointer of wrong type."); 749 : retain_(); 750 : } 751 : 752 : ~weak_intrusive_ptr() noexcept { 753 : reset_(); 754 : } 755 : 756 : weak_intrusive_ptr& operator=(weak_intrusive_ptr&& rhs) & noexcept { 757 : return operator=<TTarget, NullType>(std::move(rhs)); 758 : } 759 : 760 : template <class From, class FromNullType> 761 : weak_intrusive_ptr& operator=( 762 : weak_intrusive_ptr<From, FromNullType>&& rhs) & noexcept { 763 : static_assert( 764 : std::is_convertible<From*, TTarget*>::value, 765 : "Type mismatch. weak_intrusive_ptr move assignment got pointer of wrong type."); 766 : weak_intrusive_ptr tmp = std::move(rhs); 767 : swap(tmp); 768 : return *this; 769 : } 770 : 771 : weak_intrusive_ptr& operator=(const weak_intrusive_ptr& rhs) & noexcept { 772 : return operator=<TTarget, NullType>(rhs); 773 : } 774 : 775 : weak_intrusive_ptr& operator=( 776 : const intrusive_ptr<TTarget, NullType>& rhs) & noexcept { 777 : weak_intrusive_ptr tmp(rhs); 778 : swap(tmp); 779 : return *this; 780 : } 781 : 782 : template <class From, class FromNullType> 783 : weak_intrusive_ptr& operator=( 784 : const weak_intrusive_ptr<From, NullType>& rhs) & { 785 : static_assert( 786 : std::is_convertible<From*, TTarget*>::value, 787 : "Type mismatch. weak_intrusive_ptr copy assignment got pointer of wrong type."); 788 : weak_intrusive_ptr tmp = rhs; 789 : swap(tmp); 790 : return *this; 791 : } 792 : 793 : void reset() noexcept { 794 : reset_(); 795 : } 796 : 797 : void swap(weak_intrusive_ptr& rhs) noexcept { 798 : TTarget* tmp = target_; 799 : target_ = rhs.target_; 800 : rhs.target_ = tmp; 801 : } 802 : 803 : // NB: This should ONLY be used by the std::hash implementation 804 : // for weak_intrusive_ptr. Another way you could do this is 805 : // friend std::hash<weak_intrusive_ptr>, but this triggers two 806 : // bugs: 807 : // 808 : // (1) It triggers an nvcc bug, where std::hash in a friend class 809 : // declaration gets preprocessed into hash, which then cannot 810 : // actually be found. The error in this case looks like: 811 : // 812 : // error: no template named 'hash'; did you mean 'std::hash'? 813 : // 814 : // (2) On OS X, std::hash is declared as a struct, not a class. 815 : // This twings: 816 : // 817 : // error: class 'hash' was previously declared as a struct 818 : // [-Werror,-Wmismatched-tags] 819 : // 820 : // Both of these are work-aroundable, but on the whole, I decided 821 : // it would be simpler and easier to make work if we just expose 822 : // an unsafe getter for target_ 823 : // 824 : TTarget* _unsafe_get_target() const noexcept { 825 : return target_; 826 : } 827 : 828 : size_t use_count() const noexcept { 829 : if (target_ == NullType::singleton()) { 830 : return 0; 831 : } 832 : return target_->refcount_.load( 833 : std::memory_order_acquire); // refcount, not weakcount! 834 : } 835 : 836 : size_t weak_use_count() const noexcept { 837 : if (target_ == NullType::singleton()) { 838 : return 0; 839 : } 840 : return target_->weakcount_.load(std::memory_order_acquire); 841 : } 842 : 843 : bool expired() const noexcept { 844 : return use_count() == 0; 845 : } 846 : 847 : intrusive_ptr<TTarget, NullType> lock() const noexcept { 848 : if (expired()) { 849 : return intrusive_ptr<TTarget, NullType>(); 850 : } else { 851 : auto refcount = target_->refcount_.load(std::memory_order_seq_cst); 852 : do { 853 : if (refcount == 0) { 854 : // Object already destructed, no strong references left anymore. 855 : // Return nullptr. 856 : return intrusive_ptr<TTarget, NullType>(); 857 : } 858 : } while ( 859 : !target_->refcount_.compare_exchange_weak(refcount, refcount + 1)); 860 : return intrusive_ptr<TTarget, NullType>( 861 : target_, raw::DontIncreaseRefcount{}); 862 : } 863 : } 864 : 865 : /** 866 : * Returns an owning (but still only weakly referenced) pointer to the 867 : * underlying object and makes the weak_intrusive_ptr instance invalid. 868 : * That means the weakcount is not decreased. 869 : * You *must* put the returned pointer back into a weak_intrusive_ptr using 870 : * weak_intrusive_ptr::reclaim(ptr) to properly destruct it. 871 : * This is helpful for C APIs. 872 : */ 873 : TTarget* release() noexcept { 874 : TTarget* result = target_; 875 : target_ = NullType::singleton(); 876 : return result; 877 : } 878 : 879 : /** 880 : * Takes an owning (but must be weakly referenced) pointer to TTarget* and 881 : * creates a weak_intrusive_ptr that takes over ownership. 882 : * This means that the weakcount is not increased. 883 : * This is the counter-part to weak_intrusive_ptr::release() and the pointer 884 : * passed in *must* have been created using weak_intrusive_ptr::release(). 885 : */ 886 : static weak_intrusive_ptr reclaim(TTarget* owning_weak_ptr) { 887 : // See Note [Stack allocated intrusive_ptr_target safety] 888 : // if refcount > 0, weakcount must be >1 for weak references to exist. 889 : // see weak counting explanation at top of this file. 890 : // if refcount == 0, weakcount only must be >0. 891 : TORCH_INTERNAL_ASSERT_DEBUG_ONLY( 892 : owning_weak_ptr == NullType::singleton() || 893 : owning_weak_ptr->weakcount_.load() > 1 || 894 : (owning_weak_ptr->refcount_.load() == 0 && 895 : owning_weak_ptr->weakcount_.load() > 0), 896 : "weak_intrusive_ptr: Can only weak_intrusive_ptr::reclaim() owning pointers that were created using weak_intrusive_ptr::release()."); 897 : return weak_intrusive_ptr(owning_weak_ptr); 898 : } 899 : 900 : /** 901 : * Takes a pointer to TTarget* (may be weak or strong) and creates a 902 : * new weak_intrusive_ptr representing a new weak reference, i.e. 903 : * the raw pointer retains ownership. 904 : */ 905 : static weak_intrusive_ptr reclaim_copy(TTarget* owning_ptr) { 906 : auto ret = reclaim(owning_ptr); 907 : ret.retain_(); 908 : return ret; 909 : } 910 : 911 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 912 : friend bool operator<( 913 : const weak_intrusive_ptr<TTarget1, NullType1>& lhs, 914 : const weak_intrusive_ptr<TTarget2, NullType2>& rhs) noexcept; 915 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 916 : friend bool operator==( 917 : const weak_intrusive_ptr<TTarget1, NullType1>& lhs, 918 : const weak_intrusive_ptr<TTarget2, NullType2>& rhs) noexcept; 919 : }; 920 : 921 : template <class TTarget, class NullType> 922 : inline void swap( 923 : weak_intrusive_ptr<TTarget, NullType>& lhs, 924 : weak_intrusive_ptr<TTarget, NullType>& rhs) noexcept { 925 : lhs.swap(rhs); 926 : } 927 : 928 : // To allow weak_intrusive_ptr inside std::map or std::set, we need operator< 929 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 930 : inline bool operator<( 931 : const weak_intrusive_ptr<TTarget1, NullType1>& lhs, 932 : const weak_intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 933 : return lhs.target_ < rhs.target_; 934 : } 935 : 936 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 937 : inline bool operator==( 938 : const weak_intrusive_ptr<TTarget1, NullType1>& lhs, 939 : const weak_intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 940 : return lhs.target_ == rhs.target_; 941 : } 942 : 943 : template <class TTarget1, class NullType1, class TTarget2, class NullType2> 944 : inline bool operator!=( 945 : const weak_intrusive_ptr<TTarget1, NullType1>& lhs, 946 : const weak_intrusive_ptr<TTarget2, NullType2>& rhs) noexcept { 947 : return !operator==(lhs, rhs); 948 : } 949 : 950 : // Alias for documentary purposes, to more easily distinguish 951 : // weak raw intrusive pointers from intrusive pointers. 952 : using weak_intrusive_ptr_target = intrusive_ptr_target; 953 : 954 : // This namespace provides some methods for working with 955 : // raw pointers that subclass intrusive_ptr_target. They are not provided 956 : // as methods on intrusive_ptr_target, because ideally you would not need these 957 : // methods at all (use smart pointers), but if you are dealing with legacy code 958 : // that still needs to pass around raw pointers, you may find these quite 959 : // useful. 960 : // 961 : // An important usage note: some functions are only valid if you have a 962 : // strong raw pointer to the object, while others are only valid if you 963 : // have a weak raw pointer to the object. ONLY call intrusive_ptr namespace 964 : // functions on strong pointers, and weak_intrusive_ptr namespace functions 965 : // on weak pointers. If you mix it up, you may get an assert failure. 966 : namespace raw { 967 : 968 : namespace intrusive_ptr { 969 : 970 : // WARNING: Unlike the reclaim() API, it is NOT valid to pass 971 : // NullType::singleton to this function 972 0 : inline void incref(intrusive_ptr_target* self) { 973 0 : if (self) { 974 0 : detail::atomic_refcount_increment(self->refcount_); 975 : } 976 0 : } 977 : 978 : // WARNING: Unlike the reclaim() API, it is NOT valid to pass 979 : // NullType::singleton to this function 980 0 : inline void decref(intrusive_ptr_target* self) { 981 : // Let it die 982 0 : c10::intrusive_ptr<intrusive_ptr_target>::reclaim(self); 983 : // NB: Caller still has 'self' pointer, but it's now invalid. 984 : // If you want more safety, used the actual c10::intrusive_ptr class 985 0 : } 986 : 987 : template <typename T> 988 : inline T* make_weak(T* self) { 989 : // NB: 'this' is a strong pointer, but we return a weak pointer 990 : auto ptr = c10::intrusive_ptr<T>::reclaim(self); 991 : c10::weak_intrusive_ptr<T> wptr(ptr); 992 : ptr.release(); 993 : return wptr.release(); 994 : } 995 : 996 : inline size_t use_count(intrusive_ptr_target* self) { 997 : auto ptr = c10::intrusive_ptr<intrusive_ptr_target>::reclaim(self); 998 : auto r = ptr.use_count(); 999 : ptr.release(); 1000 : return r; 1001 : } 1002 : 1003 : } // namespace intrusive_ptr 1004 : 1005 : namespace weak_intrusive_ptr { 1006 : 1007 : inline void incref(weak_intrusive_ptr_target* self) { 1008 : detail::atomic_weakcount_increment(self->weakcount_); 1009 : } 1010 : 1011 : inline void decref(weak_intrusive_ptr_target* self) { 1012 : // Let it die 1013 : c10::weak_intrusive_ptr<intrusive_ptr_target>::reclaim(self); 1014 : // NB: You still "have" the 'self' pointer, but it's now invalid. 1015 : // If you want more safety, used the actual c10::weak_intrusive_ptr class 1016 : } 1017 : 1018 : template <typename T> 1019 : inline T* lock(T* self) { 1020 : auto wptr = c10::weak_intrusive_ptr<T>::reclaim(self); 1021 : auto ptr = wptr.lock(); 1022 : wptr.release(); 1023 : return ptr.release(); 1024 : } 1025 : 1026 : // This gives the STRONG refcount of a WEAK pointer 1027 : inline size_t use_count(weak_intrusive_ptr_target* self) { 1028 : auto wptr = c10::weak_intrusive_ptr<intrusive_ptr_target>::reclaim(self); 1029 : auto r = wptr.use_count(); 1030 : wptr.release(); 1031 : return r; 1032 : } 1033 : 1034 : } // namespace weak_intrusive_ptr 1035 : 1036 : } // namespace raw 1037 : 1038 : } // namespace c10 1039 : 1040 : namespace std { 1041 : // To allow intrusive_ptr and weak_intrusive_ptr inside std::unordered_map or 1042 : // std::unordered_set, we need std::hash 1043 : template <class TTarget, class NullType> 1044 : struct hash<c10::intrusive_ptr<TTarget, NullType>> { 1045 : size_t operator()(const c10::intrusive_ptr<TTarget, NullType>& x) const { 1046 : return std::hash<TTarget*>()(x.get()); 1047 : } 1048 : }; 1049 : template <class TTarget, class NullType> 1050 : struct hash<c10::weak_intrusive_ptr<TTarget, NullType>> { 1051 : size_t operator()(const c10::weak_intrusive_ptr<TTarget, NullType>& x) const { 1052 : return std::hash<TTarget*>()(x._unsafe_get_target()); 1053 : } 1054 : }; 1055 : } // namespace std |
![]() |
Generated by: LCOV version 2.0-1 |
</html>