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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
|
#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*/,
SymInt size_bytes,
Allocator* allocator = nullptr,
bool resizable = false)
: 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 = nullptr,
bool resizable = false)
: 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));
}
// Mimic create_legacy, but without requiring a newly-created StorageImpl.
void reset_legacy() {
TORCH_CHECK(resizable() && allocator());
set_nbytes(0);
set_data_ptr_noswap(allocator()->allocate(0));
}
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);
}
void set_nbytes(c10::SymInt 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();
}
SymInt sym_nbytes() const {
return storage_impl_->sym_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));
}
void set_data_ptr_noswap(at::DataPtr&& data_ptr) const {
return storage_impl_.get()->set_data_ptr_noswap(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();
}
c10::weak_intrusive_ptr<StorageImpl> getWeakStorageImpl() const {
return c10::weak_intrusive_ptr<StorageImpl>(storage_impl_);
}
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()) {
TORCH_CHECK(
false,
"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()) {
TORCH_CHECK(
false,
"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
|
