Home | History | Annotate | Download | only in core 
      1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #ifndef TENSORFLOW_LIB_CORE_REFCOUNT_H_
     17 #define TENSORFLOW_LIB_CORE_REFCOUNT_H_
     18 
     19 #include <atomic>
     20 #include "tensorflow/core/platform/logging.h"
     21 
     22 namespace tensorflow {
     23 namespace core {
     24 
     25 class RefCounted {
     26  public:
     27   // Initial reference count is one.
     28   RefCounted();
     29 
     30   // Increments reference count by one.
     31   void Ref() const;
     32 
     33   // Decrements reference count by one.  If the count remains
     34   // positive, returns false.  When the count reaches zero, returns
     35   // true and deletes this, in which case the caller must not access
     36   // the object afterward.
     37   bool Unref() const;
     38 
     39   // Return whether the reference count is one.
     40   // If the reference count is used in the conventional way, a
     41   // reference count of 1 implies that the current thread owns the
     42   // reference and no other thread shares it.
     43   // This call performs the test for a reference count of one, and
     44   // performs the memory barrier needed for the owning thread
     45   // to act on the object, knowing that it has exclusive access to the
     46   // object.
     47   bool RefCountIsOne() const;
     48 
     49  protected:
     50   // Make destructor protected so that RefCounted objects cannot
     51   // be instantiated directly. Only subclasses can be instantiated.
     52   virtual ~RefCounted();
     53 
     54  private:
     55   mutable std::atomic_int_fast32_t ref_;
     56 
     57   RefCounted(const RefCounted&) = delete;
     58   void operator=(const RefCounted&) = delete;
     59 };
     60 
     61 // Helper class to unref an object when out-of-scope.
     62 class ScopedUnref {
     63  public:
     64   explicit ScopedUnref(RefCounted* o) : obj_(o) {}
     65   ~ScopedUnref() {
     66     if (obj_) obj_->Unref();
     67   }
     68 
     69  private:
     70   RefCounted* obj_;
     71 
     72   ScopedUnref(const ScopedUnref&) = delete;
     73   void operator=(const ScopedUnref&) = delete;
     74 };
     75 
     76 // Inlined routines, since these are performance critical
     77 inline RefCounted::RefCounted() : ref_(1) {}
     78 
     79 inline RefCounted::~RefCounted() { DCHECK_EQ(ref_.load(), 0); }
     80 
     81 inline void RefCounted::Ref() const {
     82   DCHECK_GE(ref_.load(), 1);
     83   ref_.fetch_add(1, std::memory_order_relaxed);
     84 }
     85 
     86 inline bool RefCounted::Unref() const {
     87   DCHECK_GT(ref_.load(), 0);
     88   // If ref_==1, this object is owned only by the caller. Bypass a locked op
     89   // in that case.
     90   if (RefCountIsOne() || ref_.fetch_sub(1) == 1) {
     91     // Make DCHECK in ~RefCounted happy
     92     DCHECK((ref_.store(0), true));
     93     delete this;
     94     return true;
     95   } else {
     96     return false;
     97   }
     98 }
     99 
    100 inline bool RefCounted::RefCountIsOne() const {
    101   return (ref_.load(std::memory_order_acquire) == 1);
    102 }
    103 
    104 }  // namespace core
    105 }  // namespace tensorflow
    106 
    107 #endif  // TENSORFLOW_LIB_CORE_REFCOUNT_H_
    108