forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
blob.h
209 lines (184 loc) · 5.29 KB
/
blob.h
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
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
#pragma once
#include <cstddef>
#include <sstream>
#include <type_traits>
#include <typeinfo>
#include <vector>
#include <c10/util/intrusive_ptr.h>
#include <c10/util/typeid.h>
#include <c10/macros/Macros.h>
namespace caffe2 {
class Tensor;
/**
* @brief Blob is a general container that hosts a typed pointer.
*
* A Blob hosts a pointer as well as its type, and takes charge of deleting it
* properly when the blob is deallocated or re-allocated with a new type. A blob
* could contain anything, although the most common case is to contain a Tensor.
*/
class TORCH_API Blob final : public c10::intrusive_ptr_target {
public:
/**
* Initializes an empty Blob.
*/
Blob() noexcept : meta_(), pointer_(nullptr), has_ownership_(false) {}
~Blob() {
Reset();
}
Blob(Blob&& other) noexcept : Blob() {
swap(other);
}
Blob& operator=(Blob&& other) noexcept {
Blob(std::move(other)).swap(*this);
return *this;
}
/**
* Checks if the content stored in the blob is of type T.
*/
template <class T>
bool IsType() const noexcept {
return meta_.Match<T>();
}
/**
* Returns the meta info of the blob.
*/
const TypeMeta meta() const noexcept {
return meta_;
}
/**
* Returns a printable typename of the blob.
*/
c10::string_view TypeName() const noexcept {
return meta_.name();
}
/**
* @brief Gets the const reference of the stored object. The code checks if
* the stored object is of the desired type.
*/
// TODO(jerryzh): add a Get(DeviceType) function?
template <class T>
const T& Get() const {
TORCH_INTERNAL_ASSERT(
IsType<T>(),
"wrong type for the Blob instance. Blob contains ",
meta_.name(),
" while caller expects ",
TypeMeta::TypeName<T>());
// TODO: after we add Get<Tensor>(DeviceType)
// and changed all the callsites, we can add
// a static assert here to enforce T != Tensor
// NOLINTNEXTLINE(clang-analyzer-core.uninitialized.UndefReturn)
return *static_cast<const T*>(pointer_);
}
const void* GetRaw() const noexcept {
return pointer_;
}
void* GetRaw() noexcept {
return pointer_;
}
/**
* @brief Gets a mutable pointer to the stored object.
*
* If the current object is not of the right type, a new object is created
* and the old object is freed. Note that type T should have a default
* constructor. Otherwise, create the object yourself first, and use
* Reset().
*/
template <class T>
T* GetMutable() {
static_assert(
std::is_default_constructible<T>::value,
"GetMutable can't be called with non-default-constructible types. "
"Try using specialized methods");
if (IsType<T>()) {
return static_cast<T*>(pointer_);
} else {
// TODO Re-enable logging
// VLOG(1) << "Create new mutable object " << TypeMeta::TypeName<T>();
return Reset<T>(new T());
}
}
template <class T>
T* GetMutableOrNull() {
if (IsType<T>()) {
return static_cast<T*>(pointer_);
} else {
return nullptr;
}
}
/**
* Sets the underlying object to the allocated one. The Blob then takes over
* the ownership of the passed in pointer. If there is already an object in
* the Blob, the old object is freed.
*
* This is used when the underlying class T does not have a default ctor, or
* complex initializations needs to be done outside the blob.
*/
template <class T>
T* Reset(T* allocated) {
free_();
meta_ = TypeMeta::Make<T>();
pointer_ = static_cast<void*>(allocated);
has_ownership_ = true;
return allocated;
}
/**
* Sets the underlying object to the allocated one, but does not take over
* the ownership of the passed in pointer. If there is already an object in
* the Blob, the old object is freed.
*
* Unlike Reset, this does not take over the ownership of the pointer and the
* caller is responsible for making sure that the lifetime of the allocated
* blob outlasts the lifetime of any access to this blob, until another Reset
* call is made or the blob is destructed.
*/
template <class T>
typename std::remove_const<T>::type* ShareExternal(
typename std::remove_const<T>::type* allocated) {
return static_cast<T*>(ShareExternal(
static_cast<void*>(allocated),
TypeMeta::Make<typename std::remove_const<T>::type>()));
}
void* ShareExternal(void* allocated, const TypeMeta meta) {
free_();
meta_ = meta;
pointer_ = allocated;
has_ownership_ = false;
return allocated;
}
/**
* Resets the Blob to an empty one.
*/
void Reset() {
free_();
pointer_ = nullptr;
meta_ = TypeMeta();
has_ownership_ = false;
}
/**
* @brief Swaps the underlying storage of two blobs.
*/
void swap(Blob& rhs) {
using std::swap;
swap(meta_, rhs.meta_);
swap(pointer_, rhs.pointer_);
swap(has_ownership_, rhs.has_ownership_);
}
private:
void free_() {
if (has_ownership_ && pointer_ != nullptr) {
(*meta_.deleteFn())(pointer_);
}
}
TypeMeta meta_;
void* pointer_;
bool has_ownership_;
C10_DISABLE_COPY_AND_ASSIGN(Blob);
};
inline void swap(Blob& lhs, Blob& rhs) {
lhs.swap(rhs);
}
inline std::ostream& operator<<(std::ostream& out, const Blob& v) {
return out << "Blob[" << v.TypeName() << "]";
}
} // namespace caffe2