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_CORE_PLATFORM_REFCOUNT_H_
17 #define TENSORFLOW_CORE_PLATFORM_REFCOUNT_H_
18
19 #include <atomic>
20 #include <memory>
21
22 #include "tensorflow/core/platform/logging.h"
23
24 namespace tensorflow {
25 namespace core {
26
27 class RefCounted {
28 public:
29 // Initial reference count is one.
30 RefCounted();
31
32 // Increments reference count by one.
33 void Ref() const;
34
35 // Decrements reference count by one. If the count remains
36 // positive, returns false. When the count reaches zero, returns
37 // true and deletes this, in which case the caller must not access
38 // the object afterward.
39 bool Unref() const;
40
41 // Gets the current reference count.
42 int_fast32_t RefCount() const;
43
44 // Return whether the reference count is one.
45 // If the reference count is used in the conventional way, a
46 // reference count of 1 implies that the current thread owns the
47 // reference and no other thread shares it.
48 // This call performs the test for a reference count of one, and
49 // performs the memory barrier needed for the owning thread
50 // to act on the object, knowing that it has exclusive access to the
51 // object.
52 bool RefCountIsOne() const;
53
54 protected:
55 // Make destructor protected so that RefCounted objects cannot
56 // be instantiated directly. Only subclasses can be instantiated.
57 virtual ~RefCounted();
58
59 private:
60 mutable std::atomic_int_fast32_t ref_;
61
62 RefCounted(const RefCounted&) = delete;
63 void operator=(const RefCounted&) = delete;
64 };
65
66 // A deleter class to form a std::unique_ptr that unrefs objects.
67 struct RefCountDeleter {
operatorRefCountDeleter68 void operator()(tensorflow::core::RefCounted* o) const { o->Unref(); }
69 };
70
71 // A unique_ptr that unrefs the owned object on destruction.
72 template <typename T>
73 using RefCountPtr = std::unique_ptr<T, RefCountDeleter>;
74
75 // Helper class to unref an object when out-of-scope.
76 class ScopedUnref {
77 public:
ScopedUnref(const RefCounted * o)78 explicit ScopedUnref(const RefCounted* o) : obj_(o) {}
~ScopedUnref()79 ~ScopedUnref() {
80 if (obj_) obj_->Unref();
81 }
82
83 private:
84 const RefCounted* obj_;
85
86 ScopedUnref(const ScopedUnref&) = delete;
87 void operator=(const ScopedUnref&) = delete;
88 };
89
90 // Inlined routines, since these are performance critical
RefCounted()91 inline RefCounted::RefCounted() : ref_(1) {}
92
~RefCounted()93 inline RefCounted::~RefCounted() { DCHECK_EQ(ref_.load(), 0); }
94
Ref()95 inline void RefCounted::Ref() const {
96 DCHECK_GE(ref_.load(), 1);
97 ref_.fetch_add(1, std::memory_order_relaxed);
98 }
99
Unref()100 inline bool RefCounted::Unref() const {
101 DCHECK_GT(ref_.load(), 0);
102 // If ref_==1, this object is owned only by the caller. Bypass a locked op
103 // in that case.
104 if (RefCountIsOne() || ref_.fetch_sub(1) == 1) {
105 // Make DCHECK in ~RefCounted happy
106 DCHECK((ref_.store(0), true));
107 delete this;
108 return true;
109 } else {
110 return false;
111 }
112 }
113
RefCount()114 inline int_fast32_t RefCounted::RefCount() const {
115 return ref_.load(std::memory_order_acquire);
116 }
117
RefCountIsOne()118 inline bool RefCounted::RefCountIsOne() const {
119 return (ref_.load(std::memory_order_acquire) == 1);
120 }
121
122 } // namespace core
123 } // namespace tensorflow
124
125 #endif // TENSORFLOW_CORE_PLATFORM_REFCOUNT_H_
126