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
|
#pragma once
#include <ATen/core/ivalue.h>
namespace torch {
namespace jit {
// Used in torch.package and TorchScript serialization to coordinate
// sharing of storages between models. Also used to create deterministic
// naming for storages.
class TORCH_API SerializationStorageContext {
public:
explicit SerializationStorageContext() = default;
SerializationStorageContext operator=(const SerializationStorageContext&) =
delete;
SerializationStorageContext(const SerializationStorageContext&) = delete;
uint64_t getOrAddStorage(c10::Storage storage) {
if (!hasStorage(storage)) {
uint64_t size = storage_id_map_.size();
storage_id_map_[storage] = size;
}
return storage_id_map_[storage];
}
bool hasStorage(c10::Storage storage) {
return storage_id_map_.find(storage) != storage_id_map_.end();
}
~SerializationStorageContext() = default;
private:
class StorageSerializationHash {
public:
size_t operator()(const c10::Storage& storage) const {
return std::hash<void*>()(
reinterpret_cast<void*>(storage.unsafeGetStorageImpl()));
}
};
class StorageSerializationEqual {
public:
bool operator()(const c10::Storage& lhs, const c10::Storage& rhs) const {
return lhs.unsafeGetStorageImpl() == rhs.unsafeGetStorageImpl();
}
};
std::unordered_map<
c10::Storage,
uint64_t,
StorageSerializationHash,
StorageSerializationEqual>
storage_id_map_;
};
// Used in torch.package and TorchScript deserialization to coordinate
// sharing of storages between models.
class TORCH_API DeserializationStorageContext {
public:
explicit DeserializationStorageContext() = default;
DeserializationStorageContext operator=(
const DeserializationStorageContext&) = delete;
DeserializationStorageContext(const DeserializationStorageContext&) = delete;
void addStorage(const std::string& name, c10::Storage storage) {
TORCH_INTERNAL_ASSERT(!hasStorage(name));
name_storage_map_.insert({name, storage});
}
bool hasStorage(const std::string& name) {
return name_storage_map_.find(name) != name_storage_map_.end();
}
c10::Storage getStorage(const std::string& name) {
TORCH_INTERNAL_ASSERT(hasStorage(name));
return name_storage_map_.find(name)->second;
}
~DeserializationStorageContext() = default;
private:
std::unordered_map<std::string, c10::Storage> name_storage_map_;
};
} // namespace jit
} // namespace torch
|
