forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tensor.cc
327 lines (292 loc) · 9.97 KB
/
tensor.cc
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
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
#include "caffe2/core/tensor.h"
#include "caffe2/core/tensor_int8.h"
#include "caffe2/core/blob_stats.h"
#if defined(EXPOSE_C2_OPS) || \
!defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
#include <ATen/core/grad_mode.h>
#include "ATen/core/Tensor.h"
#endif
namespace caffe2 {
CAFFE_KNOWN_TYPE(Tensor);
TensorPrinter::TensorPrinter(
const std::string& tensor_name,
const std::string& file_name,
int limit)
: to_file_(!file_name.empty()),
limit_(limit ? limit : k_limit_default_),
tensor_name_(tensor_name) {
if (to_file_) {
// We will output to file instead of printing on screen.
// We will write each individual tensor to its individual file.
log_file_.reset(new std::ofstream(
file_name, std::ofstream::out | std::ofstream::trunc));
CAFFE_ENFORCE(
log_file_->good(),
"Failed to open TensorPrinter file ",
file_name,
". rdstate() = ",
log_file_->rdstate());
}
}
TensorPrinter::~TensorPrinter() {
if (log_file_.get()) {
log_file_->close();
}
}
void TensorPrinter::PrintMeta(const Tensor& tensor) {
if (to_file_) {
(*log_file_) << MetaStr(tensor) << std::endl;
} else {
LOG(INFO) << MetaStr(tensor);
}
}
std::string TensorPrinter::MetaStr(const Tensor& tensor) {
std::stringstream meta_stream;
meta_stream << "Tensor " << tensor_name_ << " of type "
<< tensor.dtype().name() << ". Dims: (";
for (const auto dim : tensor.sizes()) {
meta_stream << dim << ",";
}
meta_stream << "): ";
return meta_stream.str();
}
TypeMeta GetTensorType(const void* c) {
const Tensor* tc = static_cast<const Tensor*>(c);
return tc->dtype();
}
TypeMeta GetInt8TensorType(const void* c) {
const int8::Int8TensorCPU* int8_tensor =
static_cast<const int8::Int8TensorCPU*>(c);
return (int8_tensor->t).dtype();
}
// TODO(jerryzh): Remove
static CaffeMap<TypeIdentifier, TypeCall> type_call_registry_{
{TypeMeta::Id<Tensor>(), GetTensorType},
{TypeMeta::Id<int8::Int8TensorCPU>(), GetInt8TensorType},
};
TypeCall GetTypeCallFunction(TypeIdentifier id) {
auto f = type_call_registry_.find(id);
if (f == type_call_registry_.end()) {
return nullptr;
}
return f->second;
}
void RegisterTypeCallFunction(TypeIdentifier id, TypeCall c) {
type_call_registry_[id] = c;
}
int GetGPUIDForPointer(const void* ptr);
vector<int64_t>
GetTensorInfo(const void* c, size_t* capacity, DeviceOption* device) {
CHECK(capacity);
const Tensor* tc = static_cast<const Tensor*>(c);
CHECK(tc);
CHECK(tc->unsafeGetTensorImpl());
CHECK(tc->unsafeGetTensorImpl()->storage().unsafeGetStorageImpl());
*capacity = tc->storage().capacity();
ExtractDeviceOption(device, tc->GetDevice());
return tc->sizes().vec();
}
vector<int64_t>
GetInt8TensorInfo(const void* c, size_t* capacity, DeviceOption* device) {
const int8::Int8TensorCPU* int8_tensor =
static_cast<const int8::Int8TensorCPU*>(c);
return GetTensorInfo(&(int8_tensor->t), capacity, device);
}
// since we only have one tensor, probably need to remove this at some point?
static CaffeMap<TypeIdentifier, TensorInfoCall> tensor_info_call_registry_{
{TypeMeta::Id<Tensor>(), GetTensorInfo},
{TypeMeta::Id<int8::Int8TensorCPU>(), GetInt8TensorInfo},
};
// TODO: Remove this code in a separate diff, since we only have one
// GetTensorInfo function now
TensorInfoCall GetTensorInfoFunction(TypeIdentifier id) {
auto f = tensor_info_call_registry_.find(id);
if (f == tensor_info_call_registry_.end()) {
return nullptr;
}
return f->second;
}
void RegisterTensorInfoFunction(TypeIdentifier id, TensorInfoCall c) {
tensor_info_call_registry_[id] = c;
}
void TensorVectorResize(
std::vector<Tensor>& tensors,
int size,
DeviceType type) {
tensors.reserve(size);
for (auto i = 0; i < size; ++i) {
tensors.emplace_back(type);
}
}
Tensor empty(at::IntArrayRef dims, at::TensorOptions options) {
// TODO: merge this with at::empty after Tensor is merged
auto tensor = Tensor(dims, options.device());
tensor.raw_mutable_data(options.dtype());
return tensor;
}
void ReinitializeTensor(
Tensor* tensor,
at::IntArrayRef dims,
at::TensorOptions options) {
CAFFE_ENFORCE(options.device_opt() != c10::nullopt);
if (*tensor) {
// Note: we don't compare device_id here because of the purpose of
// ReinitializeTensor: https://github.com/pytorch/pytorch/pull/13147
// In the original code, we don't have device_id defined, therefore, we
// should not include device_id in the comparison
if (tensor->GetDeviceType() == options.device().type()) {
if (tensor->sizes() != dims) {
// Resize when the dims doesn't match
tensor->Resize(dims);
}
if (tensor->dtype() == options.dtype()) {
tensor->raw_mutable_data();
} else {
C10_LOG_FIRST_N(WARNING, 1)
<< "Changing the data type of Tensor is discouraged."
<< " Attempt to change data type from: " << tensor->dtype()
<< " to: " << options.dtype();
// create a new Tensor when the data_type doesn't match
*tensor = caffe2::empty(dims, options);
}
return;
}
// create a new Tensor when device doesn't match
}
VLOG(1) << "Create new mutable object " << TypeMeta::TypeName<Tensor>()
<< " dims: " << dims;
*tensor = caffe2::empty(dims, options);
}
void ReinitializeAndCopyFrom(
Tensor* t,
at::TensorOptions options,
const Tensor& src,
bool async) {
auto device_type = options.device().type();
CAFFE_ENFORCE(t != nullptr, "Target tensor ptr is null.");
if (!*t || device_type != t->GetDeviceType()) {
*t = Tensor(device_type);
}
CAFFE_ENFORCE(
!t->dtype_initialized() || t->dtype() == src.dtype(),
"We don't allow a change of data type in ReinitializeAndCopyFrom. Attempt to "
" change from: ",
t->dtype(),
" to: ",
src.dtype());
t->CopyFrom(src, async);
}
void Tensor::enforce_invariants() {
if (impl_.get() == nullptr) {
throw std::runtime_error("TensorImpl with nullptr is not supported");
}
// TODO: only check `!impl_->requires_grad()` after Variable and Tensor are
// merged
#if !defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
CAFFE_ENFORCE(
!(impl_->requires_grad() && at::GradMode::is_enabled()),
"Caffe2 tensor wrapper doesn't support autograd variables that require grad");
#endif
CAFFE_ENFORCE_EQ(
impl_->layout(),
at::kStrided,
"Caffe2 tensor wrapper supports only regular non-sparse tensors");
CAFFE_ENFORCE(
impl_->is_contiguous(),
"Caffe2 tensor wrapper supports only contiguous tensors");
}
void Tensor::CopyFrom(const Tensor& src, bool async) {
// TODO: only check `!impl_->requires_grad()` after Variable and Tensor are
// merged
#if defined(EXPOSE_C2_OPS) || \
!defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
AT_ASSERT(!(impl_->requires_grad() && at::GradMode::is_enabled()));
#endif
AT_ASSERTM(
src.impl_->is_contiguous(),
"Right now only copy of contiguous source Tensor is supported.");
AT_ASSERTM(
src.impl_->storage_initialized(),
"Cannot copy from an uninitialized Tensor");
if (src.impl_.get() == impl_.get()) {
return;
}
// Test if we need to allocate a new storage
// Uninitialized storages are guaranteed to be uniquely owned,
// so we don't need to swap in dst case.
// If the dtype changed, we need to reallocate storage.
if (impl_->dtype() != src.impl_->dtype()) {
// NB: copy preserves device_type
// This storage will get initialized by the mutable_data call below.
impl_->set_storage(
at::Storage::create_legacy(impl_->device_type(), src.impl_->dtype()));
}
impl_->Resize(src.impl_->sizes());
if (impl_->numel() > 0) {
if (impl_->dtype().copy()) {
AT_ASSERTM(
impl_->device_type() == ::at::DeviceType::CPU,
"In CopyFrom source and dest tensors must both be CPU for "
"non-POD copy, but dest tensor was ",
impl_->device_type());
AT_ASSERTM(
src.impl_->device_type() == ::at::DeviceType::CPU,
"In CopyFrom source and dest tensors must both be CPU for "
"non-POD copy, but src tensor was ",
src.impl_->device_type());
impl_->dtype().copy()(
src.impl_->data(),
impl_->raw_mutable_data(impl_->dtype()),
impl_->numel());
} else {
// The following copy uses the current (thread local) stream for copying
// and also takes the GPU id from the device() field passed in.
//
// TODO: Potentially more enforcements are necessary to avoid accidental
// switch to sync copy if the currently set device is wrong.
//
// Specifically, we might need to switch to a different context device
// here explicitly to avoid relying on user synchronizing things
// properly.
//
// note: raw_mutable_data initializes device here
void* new_data = impl_->raw_mutable_data(impl_->dtype());
at::CopyBytes(
impl_->numel() * impl_->itemsize(),
src.impl_->data(),
src.impl_->device(),
new_data,
impl_->device(),
async);
}
}
}
#if defined(EXPOSE_C2_OPS) || \
!defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
Tensor::Tensor(at::Tensor tensor) : impl_(std::move(tensor.impl_)) {
enforce_invariants();
}
Tensor::operator at::Tensor() const& {
return at::Tensor::wrap_tensor_impl(impl_);
}
Tensor::operator at::Tensor() && {
return at::Tensor::wrap_tensor_impl(std::move(impl_));
}
#endif
namespace {
struct TensorStatGetter : BlobStatGetter {
size_t sizeBytes(const Blob& blob) const override {
const auto& tensor = blob.Get<Tensor>();
auto nbytes = tensor.nbytes();
if (nbytes > 0 && tensor.IsType<std::string>()) {
const auto* data = tensor.data<std::string>();
for (int i = 0; i < tensor.numel(); ++i) {
nbytes += data[i].size();
}
}
return nbytes;
}
};
REGISTER_BLOB_STAT_GETTER(Tensor, TensorStatGetter);
} // namespace
} // namespace caffe2