153 lines
3.5 KiB
C++
153 lines
3.5 KiB
C++
// Copyright (c) 2021 CINN 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.
|
|
|
|
#pragma once
|
|
#include <atomic>
|
|
#include <string>
|
|
#include <type_traits>
|
|
|
|
namespace cinn {
|
|
namespace common {
|
|
|
|
class RefCount {
|
|
public:
|
|
using value_type = int32_t;
|
|
RefCount() = default;
|
|
|
|
value_type Inc() { return ++count_; }
|
|
value_type Dec() { return --count_; }
|
|
bool is_zero() const { return 0 == count_; }
|
|
std::string to_string() { return std::to_string(count_.load()); }
|
|
int32_t val() const { return count_; }
|
|
|
|
private:
|
|
std::atomic<value_type> count_{0};
|
|
};
|
|
|
|
class Object;
|
|
/**
|
|
* The templated methods are used to unify the way to get the RefCount instance
|
|
* in client classes.
|
|
*/
|
|
template <typename T>
|
|
RefCount& ref_count(const T* t) {
|
|
static_assert(std::is_base_of<Object, T>::value, "T is not a Object");
|
|
return t->__ref_count__;
|
|
}
|
|
template <typename T>
|
|
void Destroy(const T* t) {
|
|
delete t;
|
|
}
|
|
|
|
template <typename T>
|
|
struct Shared {
|
|
using object_ptr = T*;
|
|
|
|
Shared() = default;
|
|
Shared(T* p) : p_(p) { // NOLINT
|
|
if (p) IncRef(p);
|
|
}
|
|
Shared(const Shared& other) : p_(other.p_) { IncRef(p_); }
|
|
Shared(Shared&& other) : p_(other.p_) { other.p_ = nullptr; }
|
|
Shared<T>& operator=(const Shared<T>& other);
|
|
|
|
//! Reset to another pointer \p x.
|
|
void Reset(T* x = nullptr);
|
|
|
|
//! Access the pointer in various ways.
|
|
// @{
|
|
inline T* get() const { return p_; }
|
|
inline T& operator*() const { return *p_; }
|
|
inline T* operator->() const { return p_; }
|
|
inline T* self() { return p_; }
|
|
inline const T* self() const { return p_; }
|
|
// @}
|
|
|
|
inline bool same_as(const Shared& other) const { return p_ == other.p_; }
|
|
inline bool defined() const { return p_; }
|
|
inline bool operator<(const Shared& other) const { return p_ < other.p_; }
|
|
inline Shared<T>& operator=(T* x);
|
|
inline bool operator==(const Shared& other) const { return p_ == other.p_; }
|
|
|
|
~Shared();
|
|
|
|
private:
|
|
//! Increase the share count.
|
|
void IncRef(T* p);
|
|
|
|
//! Decrease the share count.
|
|
void DecRef(T* p);
|
|
|
|
protected:
|
|
T* p_{};
|
|
};
|
|
|
|
template <typename T>
|
|
void Shared<T>::IncRef(T* p) {
|
|
if (p) {
|
|
ref_count(p).Inc();
|
|
}
|
|
}
|
|
template <typename T>
|
|
void Shared<T>::DecRef(T* p) {
|
|
if (p) {
|
|
if (ref_count(p).Dec() == 0) {
|
|
Destroy(p);
|
|
}
|
|
}
|
|
}
|
|
template <typename T>
|
|
Shared<T>& Shared<T>::operator=(const Shared<T>& other) {
|
|
if (other.p_ == p_) return *this;
|
|
// Other can be inside of something owned by this, so we should be careful to
|
|
// incref other before we decref ourselves.
|
|
T* tmp = other.p_;
|
|
IncRef(tmp);
|
|
DecRef(p_);
|
|
p_ = tmp;
|
|
return *this;
|
|
}
|
|
|
|
template <typename T, typename... Args>
|
|
T* make_shared(Args&&... args) {
|
|
return new T(args...);
|
|
}
|
|
|
|
template <typename T>
|
|
Shared<T>& Shared<T>::operator=(T* x) {
|
|
if (p_ == x) return *this;
|
|
|
|
T* tmp = x;
|
|
IncRef(tmp);
|
|
DecRef(p_);
|
|
p_ = tmp;
|
|
return *this;
|
|
}
|
|
|
|
template <typename T>
|
|
Shared<T>::~Shared() {
|
|
DecRef(p_);
|
|
p_ = nullptr;
|
|
}
|
|
|
|
template <typename T>
|
|
void Shared<T>::Reset(T* x) {
|
|
if (x) IncRef(x);
|
|
DecRef(p_);
|
|
p_ = x;
|
|
}
|
|
|
|
} // namespace common
|
|
} // namespace cinn
|