aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/lib/core/refcount.cc
blob: 3ed8c58eb8d5723731a79f81addb39860b4235b7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/platform/logging.h"

namespace tensorflow {
namespace core {

RefCounted::RefCounted() : ref_(1) {}

RefCounted::~RefCounted() { DCHECK_EQ(ref_.load(), 0); }

void RefCounted::Ref() const {
  DCHECK_GE(ref_.load(), 1);
  ref_.fetch_add(1, std::memory_order_relaxed);
}

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 (ref_.load(std::memory_order_acquire) == 1 || ref_.fetch_sub(1) == 1) {
    // Make DCHECK in ~RefCounted happy
    DCHECK((ref_.store(0), true));
    delete this;
    return true;
  } else {
    return false;
  }
}

bool RefCounted::RefCountIsOne() const {
  return (ref_.load(std::memory_order_acquire) == 1);
}

}  // namespace core
}  // namespace tensorflow