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
36
37
38
39
40
41
42
|
#include <c10/core/impl/COWDeleter.h>
#include <c10/util/Exception.h>
#include <mutex>
namespace c10::impl {
void cow::cow_deleter(void* ctx) {
static_cast<cow::COWDeleterContext*>(ctx)->decrement_refcount();
}
cow::COWDeleterContext::COWDeleterContext(
std::unique_ptr<void, DeleterFnPtr> data)
: data_(std::move(data)) {
// We never wrap a COWDeleterContext.
TORCH_INTERNAL_ASSERT(data_.get_deleter() != cow::cow_deleter);
}
auto cow::COWDeleterContext::increment_refcount() -> void {
auto refcount = ++refcount_;
TORCH_INTERNAL_ASSERT(refcount > 1);
}
auto cow::COWDeleterContext::decrement_refcount()
-> std::variant<NotLastReference, LastReference> {
auto refcount = --refcount_;
TORCH_INTERNAL_ASSERT(refcount >= 0, refcount);
if (refcount == 0) {
std::unique_lock lock(mutex_);
auto result = std::move(data_);
lock.unlock();
delete this;
return {std::move(result)};
}
return std::shared_lock(mutex_);
}
cow::COWDeleterContext::~COWDeleterContext() {
TORCH_INTERNAL_ASSERT(refcount_ == 0);
}
} // namespace c10::impl
|
