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 <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 // Return whether the reference count is one.
42 // If the reference count is used in the conventional way, a
43 // reference count of 1 implies that the current thread owns the
44 // reference and no other thread shares it.
45 // This call performs the test for a reference count of one, and
46 // performs the memory barrier needed for the owning thread
47 // to act on the object, knowing that it has exclusive access to the
48 // object.
49 bool RefCountIsOne() const;
50
51 protected:
52 // Make destructor protected so that RefCounted objects cannot
53 // be instantiated directly. Only subclasses can be instantiated.
54 virtual ~RefCounted();
55
56 private:
57 mutable std::atomic_int_fast32_t ref_;
58
59 RefCounted(const RefCounted&) = delete;
60 void operator=(const RefCounted&) = delete;
61 };
62
63 // A deleter class to form a std::unique_ptr that unrefs objects.
64 struct RefCountDeleter {
operatorRefCountDeleter65 void operator()(tensorflow::core::RefCounted* o) const { o->Unref(); }
66 };
67
68 // A unique_ptr that unrefs the owned object on destruction.
69 template <typename T>
70 using RefCountPtr = std::unique_ptr<T, RefCountDeleter>;
71
72 // Helper class to unref an object when out-of-scope.
73 class ScopedUnref {
74 public:
ScopedUnref(RefCounted * o)75 explicit ScopedUnref(RefCounted* o) : obj_(o) {}
~ScopedUnref()76 ~ScopedUnref() {
77 if (obj_) obj_->Unref();
78 }
79
80 private:
81 RefCounted* obj_;
82
83 ScopedUnref(const ScopedUnref&) = delete;
84 void operator=(const ScopedUnref&) = delete;
85 };
86
87 // Inlined routines, since these are performance critical
RefCounted()88 inline RefCounted::RefCounted() : ref_(1) {}
89
~RefCounted()90 inline RefCounted::~RefCounted() { DCHECK_EQ(ref_.load(), 0); }
91
Ref()92 inline void RefCounted::Ref() const {
93 DCHECK_GE(ref_.load(), 1);
94 ref_.fetch_add(1, std::memory_order_relaxed);
95 }
96
Unref()97 inline bool RefCounted::Unref() const {
98 DCHECK_GT(ref_.load(), 0);
99 // If ref_==1, this object is owned only by the caller. Bypass a locked op
100 // in that case.
101 if (RefCountIsOne() || ref_.fetch_sub(1) == 1) {
102 // Make DCHECK in ~RefCounted happy
103 DCHECK((ref_.store(0), true));
104 delete this;
105 return true;
106 } else {
107 return false;
108 }
109 }
110
RefCountIsOne()111 inline bool RefCounted::RefCountIsOne() const {
112 return (ref_.load(std::memory_order_acquire) == 1);
113 }
114
115 } // namespace core
116 } // namespace tensorflow
117
118 #endif // TENSORFLOW_LIB_CORE_REFCOUNT_H_
119