diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 6aa4db99310..8a7bed67048 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -2131,6 +2131,7 @@ LIB_INTERNAL_PUBLIC_HEADERS = [ "//tensorflow/core/platform:platform.h", "//tensorflow/core/platform:monitoring.h", "//tensorflow/core/platform:protobuf_internal.h", + "//tensorflow/core/platform:refcount.h", "//tensorflow/core/platform:setround.h", "//tensorflow/core/platform:snappy.h", "//tensorflow/core/platform:tensor_coding.h", diff --git a/tensorflow/core/lib/core/BUILD b/tensorflow/core/lib/core/BUILD index 38d15136d80..a37b1c41fce 100644 --- a/tensorflow/core/lib/core/BUILD +++ b/tensorflow/core/lib/core/BUILD @@ -95,7 +95,7 @@ cc_library( cc_library( name = "refcount", hdrs = ["refcount.h"], - deps = ["//tensorflow/core/platform:logging"], + deps = ["//tensorflow/core/platform:refcount"], ) cc_library( diff --git a/tensorflow/core/lib/core/refcount.h b/tensorflow/core/lib/core/refcount.h index 7b9bb4c9862..5f85dcb2784 100644 --- a/tensorflow/core/lib/core/refcount.h +++ b/tensorflow/core/lib/core/refcount.h @@ -16,103 +16,6 @@ limitations under the License. #ifndef TENSORFLOW_LIB_CORE_REFCOUNT_H_ #define TENSORFLOW_LIB_CORE_REFCOUNT_H_ -#include -#include - -#include "tensorflow/core/platform/logging.h" - -namespace tensorflow { -namespace core { - -class RefCounted { - public: - // Initial reference count is one. - RefCounted(); - - // Increments reference count by one. - void Ref() const; - - // Decrements reference count by one. If the count remains - // positive, returns false. When the count reaches zero, returns - // true and deletes this, in which case the caller must not access - // the object afterward. - bool Unref() const; - - // Return whether the reference count is one. - // If the reference count is used in the conventional way, a - // reference count of 1 implies that the current thread owns the - // reference and no other thread shares it. - // This call performs the test for a reference count of one, and - // performs the memory barrier needed for the owning thread - // to act on the object, knowing that it has exclusive access to the - // object. - bool RefCountIsOne() const; - - protected: - // Make destructor protected so that RefCounted objects cannot - // be instantiated directly. Only subclasses can be instantiated. - virtual ~RefCounted(); - - private: - mutable std::atomic_int_fast32_t ref_; - - RefCounted(const RefCounted&) = delete; - void operator=(const RefCounted&) = delete; -}; - -// A deleter class to form a std::unique_ptr that unrefs objects. -struct RefCountDeleter { - void operator()(tensorflow::core::RefCounted* o) const { o->Unref(); } -}; - -// A unique_ptr that unrefs the owned object on destruction. -template -using RefCountPtr = std::unique_ptr; - -// Helper class to unref an object when out-of-scope. -class ScopedUnref { - public: - explicit ScopedUnref(const RefCounted* o) : obj_(o) {} - ~ScopedUnref() { - if (obj_) obj_->Unref(); - } - - private: - const RefCounted* obj_; - - ScopedUnref(const ScopedUnref&) = delete; - void operator=(const ScopedUnref&) = delete; -}; - -// Inlined routines, since these are performance critical -inline RefCounted::RefCounted() : ref_(1) {} - -inline RefCounted::~RefCounted() { DCHECK_EQ(ref_.load(), 0); } - -inline void RefCounted::Ref() const { - DCHECK_GE(ref_.load(), 1); - ref_.fetch_add(1, std::memory_order_relaxed); -} - -inline bool RefCounted::Unref() const { - DCHECK_GT(ref_.load(), 0); - // If ref_==1, this object is owned only by the caller. Bypass a locked op - // in that case. - if (RefCountIsOne() || ref_.fetch_sub(1) == 1) { - // Make DCHECK in ~RefCounted happy - DCHECK((ref_.store(0), true)); - delete this; - return true; - } else { - return false; - } -} - -inline bool RefCounted::RefCountIsOne() const { - return (ref_.load(std::memory_order_acquire) == 1); -} - -} // namespace core -} // namespace tensorflow +#include "tensorflow/core/platform/refcount.h" #endif // TENSORFLOW_LIB_CORE_REFCOUNT_H_ diff --git a/tensorflow/core/platform/BUILD b/tensorflow/core/platform/BUILD index 0df8010e319..6ffaffba100 100644 --- a/tensorflow/core/platform/BUILD +++ b/tensorflow/core/platform/BUILD @@ -382,6 +382,14 @@ cc_library( ], ) +cc_library( + name = "refcount", + hdrs = ["refcount.h"], + deps = [ + ":logging", + ], +) + cc_library( name = "regexp", hdrs = ["regexp.h"], diff --git a/tensorflow/core/platform/refcount.h b/tensorflow/core/platform/refcount.h new file mode 100644 index 00000000000..b58ca319138 --- /dev/null +++ b/tensorflow/core/platform/refcount.h @@ -0,0 +1,118 @@ +/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_CORE_PLATFORM_REFCOUNT_H_ +#define TENSORFLOW_CORE_PLATFORM_REFCOUNT_H_ + +#include +#include + +#include "tensorflow/core/platform/logging.h" + +namespace tensorflow { +namespace core { + +class RefCounted { + public: + // Initial reference count is one. + RefCounted(); + + // Increments reference count by one. + void Ref() const; + + // Decrements reference count by one. If the count remains + // positive, returns false. When the count reaches zero, returns + // true and deletes this, in which case the caller must not access + // the object afterward. + bool Unref() const; + + // Return whether the reference count is one. + // If the reference count is used in the conventional way, a + // reference count of 1 implies that the current thread owns the + // reference and no other thread shares it. + // This call performs the test for a reference count of one, and + // performs the memory barrier needed for the owning thread + // to act on the object, knowing that it has exclusive access to the + // object. + bool RefCountIsOne() const; + + protected: + // Make destructor protected so that RefCounted objects cannot + // be instantiated directly. Only subclasses can be instantiated. + virtual ~RefCounted(); + + private: + mutable std::atomic_int_fast32_t ref_; + + RefCounted(const RefCounted&) = delete; + void operator=(const RefCounted&) = delete; +}; + +// A deleter class to form a std::unique_ptr that unrefs objects. +struct RefCountDeleter { + void operator()(tensorflow::core::RefCounted* o) const { o->Unref(); } +}; + +// A unique_ptr that unrefs the owned object on destruction. +template +using RefCountPtr = std::unique_ptr; + +// Helper class to unref an object when out-of-scope. +class ScopedUnref { + public: + explicit ScopedUnref(const RefCounted* o) : obj_(o) {} + ~ScopedUnref() { + if (obj_) obj_->Unref(); + } + + private: + const RefCounted* obj_; + + ScopedUnref(const ScopedUnref&) = delete; + void operator=(const ScopedUnref&) = delete; +}; + +// Inlined routines, since these are performance critical +inline RefCounted::RefCounted() : ref_(1) {} + +inline RefCounted::~RefCounted() { DCHECK_EQ(ref_.load(), 0); } + +inline void RefCounted::Ref() const { + DCHECK_GE(ref_.load(), 1); + ref_.fetch_add(1, std::memory_order_relaxed); +} + +inline bool RefCounted::Unref() const { + DCHECK_GT(ref_.load(), 0); + // If ref_==1, this object is owned only by the caller. Bypass a locked op + // in that case. + if (RefCountIsOne() || ref_.fetch_sub(1) == 1) { + // Make DCHECK in ~RefCounted happy + DCHECK((ref_.store(0), true)); + delete this; + return true; + } else { + return false; + } +} + +inline bool RefCounted::RefCountIsOne() const { + return (ref_.load(std::memory_order_acquire) == 1); +} + +} // namespace core +} // namespace tensorflow + +#endif // TENSORFLOW_CORE_PLATFORM_REFCOUNT_H_