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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
#pragma once
#include <c10/core/StorageImpl.h>
namespace c10 {
struct C10_API Storage {
public:
struct use_byte_size_t {};
Storage() {}
Storage(c10::intrusive_ptr<StorageImpl> ptr) : storage_impl_(std::move(ptr)) {}
// Allocates memory buffer using given allocator and creates a storage with it
Storage(
use_byte_size_t use_byte_size,
size_t size_bytes,
Allocator* allocator,
bool resizable)
: storage_impl_(c10::make_intrusive<StorageImpl>(
StorageImpl::use_byte_size_t(),
size_bytes,
allocator,
resizable)) {}
// Creates storage with pre-allocated memory buffer. Allocator is given for
// potential future reallocations, however it can be nullptr if the storage
// is non-resizable
Storage(
use_byte_size_t use_byte_size,
size_t size_bytes,
at::DataPtr data_ptr,
at::Allocator* allocator,
bool resizable)
: storage_impl_(c10::make_intrusive<StorageImpl>(
StorageImpl::use_byte_size_t(),
size_bytes,
std::move(data_ptr),
allocator,
resizable)) {}
// Legacy constructor for partially initialized (dtype or memory) storages
// that can be temporarily created with Caffe2 APIs. See the note on top of
// TensorImpl.h for details.
static Storage create_legacy(at::Device device) {
auto allocator = GetAllocator(device.type());
return Storage(c10::make_intrusive<StorageImpl>(
StorageImpl::use_byte_size_t(),
0,
allocator->allocate(0), // materialize a non-default Device.
allocator,
true));
}
template <typename T>
T* data() const { return storage_impl_->data<T>(); }
template <typename T>
T* unsafe_data() const { return storage_impl_->unsafe_data<T>(); }
// TODO: remove later
void set_nbytes(size_t size_bytes) const {
storage_impl_.get()->set_nbytes(size_bytes);
}
bool resizable() const {
return storage_impl_->resizable();
}
size_t nbytes() const {
return storage_impl_->nbytes();
}
// get() use here is to get const-correctness
void* data() const {
return storage_impl_.get()->data();
}
at::DataPtr& data_ptr() {
return storage_impl_->data_ptr();
}
const at::DataPtr& data_ptr() const {
return storage_impl_->data_ptr();
}
// Returns the previous data_ptr
at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) const {
return storage_impl_.get()->set_data_ptr(std::move(data_ptr));
};
DeviceType device_type() const {
return storage_impl_->device_type();
}
at::Allocator* allocator() const {
return storage_impl_.get()->allocator();
}
at::Device device() const {
return storage_impl_->device();
}
StorageImpl* unsafeReleaseStorageImpl() {
return storage_impl_.release();
}
StorageImpl* unsafeGetStorageImpl() const noexcept {
return storage_impl_.get();
}
operator bool() const {
return storage_impl_;
}
size_t use_count() const {
return storage_impl_.use_count();
}
inline bool unique() const {
return storage_impl_.unique();
}
bool is_alias_of(const Storage& other) const {
return storage_impl_ == other.storage_impl_;
}
void UniqueStorageShareExternalPointer(
void* src,
size_t capacity,
DeleterFnPtr d = nullptr) {
if (!storage_impl_.unique()) {
AT_ERROR(
"UniqueStorageShareExternalPointer can only be called when use_count == 1");
}
storage_impl_->UniqueStorageShareExternalPointer(src, capacity, d);
}
void UniqueStorageShareExternalPointer(
at::DataPtr&& data_ptr,
size_t capacity) {
if (!storage_impl_.unique()) {
AT_ERROR(
"UniqueStorageShareExternalPointer can only be called when use_count == 1");
}
storage_impl_->UniqueStorageShareExternalPointer(
std::move(data_ptr), capacity);
}
protected:
c10::intrusive_ptr<StorageImpl> storage_impl_;
};
} // namespace c10
|
