1 #include "mindir_memory_manager.h"
2 #include <iostream>
3 #include "securec.h"
4 #include "utils.h"
5 #include "log.h"
6
7 namespace mindspore {
8 namespace lite {
9 namespace {
10 template <typename KEY_TYPE, typename VALUE_TYPE>
ClearMap(std::map<KEY_TYPE,VALUE_TYPE> & map)11 void ClearMap(std::map<KEY_TYPE, VALUE_TYPE> &map) {
12 for (auto iter = map.begin(); iter != map.end();) {
13 if (iter->second != nullptr) {
14 free(iter->second);
15 map.erase(iter++);
16 } else {
17 iter++;
18 }
19 }
20 }
21 } // namespace
GetInstance()22 MindIRMemoryManager *MindIRMemoryManager::GetInstance() {
23 static MindIRMemoryManager instance;
24 return &instance;
25 }
26
CopyFbbToNewMemory(flatbuffers::FlatBufferBuilder & fbb)27 void *MindIRMemoryManager::CopyFbbToNewMemory(flatbuffers::FlatBufferBuilder &fbb) {
28 auto buff = reinterpret_cast<uint8_t *>(malloc(fbb.GetSize()));
29 if (buff == nullptr) {
30 MS_LOG(ERROR) << "malloc memory for primitive failed!";
31 fbb.Clear();
32 return nullptr;
33 }
34 auto ret = memcpy_s(buff, fbb.GetSize(), fbb.GetBufferPointer(), fbb.GetSize());
35 if (ret != EOK) {
36 MS_LOG(ERROR) << "memcpy failed, ret:" << ret;
37 fbb.Clear();
38 return nullptr;
39 }
40 fbb.Clear();
41 return buff;
42 }
CreateTensorFromBuilder(flatbuffers::FlatBufferBuilder & fbb_new,schema::Tensor * tensor)43 void *MindIRMemoryManager::CreateTensorFromBuilder(flatbuffers::FlatBufferBuilder &fbb_new, schema::Tensor *tensor) {
44 std::lock_guard<std::mutex> lck(mutex);
45 if (tensor != nullptr) {
46 // find primitive exist
47 if (tensor_map.find(tensor) != tensor_map.end()) {
48 // if find, then delete
49 void *flatbuffer_ptr = tensor_map[tensor];
50 if (flatbuffer_ptr != nullptr) {
51 free(flatbuffer_ptr);
52 tensor_map[tensor] = nullptr;
53 tensor_map.erase(tensor_map.find(tensor));
54 }
55 }
56 }
57 // then copy fbb
58 auto new_memory_ptr = CopyFbbToNewMemory(fbb_new);
59 auto tensor_root = flatbuffers::GetMutableRoot<schema::Tensor>(new_memory_ptr);
60 tensor_map[tensor_root] = new_memory_ptr;
61 return new_memory_ptr;
62 }
63
CreatePrimitiveFromBuilder(flatbuffers::FlatBufferBuilder & fbb_new,schema::Primitive * primitive)64 void *MindIRMemoryManager::CreatePrimitiveFromBuilder(flatbuffers::FlatBufferBuilder &fbb_new,
65 schema::Primitive *primitive) {
66 std::lock_guard<std::mutex> lck(mutex);
67 if (primitive != nullptr) {
68 // find primitive exist
69 if (primitive_map.find(primitive) != primitive_map.end()) {
70 // if find, then delete
71 void *flatbuffer_ptr = primitive_map[primitive];
72 if (flatbuffer_ptr != nullptr) {
73 free(flatbuffer_ptr);
74 primitive_map[primitive] = nullptr;
75 primitive_map.erase(primitive_map.find(primitive));
76 }
77 }
78 }
79 // then copy fbb
80 auto new_memory_ptr = CopyFbbToNewMemory(fbb_new);
81 auto primitive_root = flatbuffers::GetMutableRoot<schema::Primitive>(new_memory_ptr);
82 primitive_map[primitive_root] = new_memory_ptr;
83 return new_memory_ptr;
84 }
85
DeletePrimitive(schema::Primitive * primitive)86 void MindIRMemoryManager::DeletePrimitive(schema::Primitive *primitive) {
87 std::lock_guard<std::mutex> lck(mutex);
88 if (primitive == nullptr) {
89 MS_LOG(ERROR) << "primitive is nullptr, no need to delete.";
90 return;
91 }
92 if (primitive_map.find(primitive) != primitive_map.end()) {
93 // if find, then delete
94 void *flatbuffer_ptr = primitive_map[primitive];
95 if (flatbuffer_ptr != nullptr) {
96 free(flatbuffer_ptr);
97 primitive_map[primitive] = nullptr;
98 primitive_map.erase(primitive_map.find(primitive));
99 }
100 }
101 }
102
DeleteTensor(schema::Tensor * tensor)103 void MindIRMemoryManager::DeleteTensor(schema::Tensor *tensor) {
104 std::lock_guard<std::mutex> lck(mutex);
105 if (tensor == nullptr) {
106 MS_LOG(ERROR) << "tensor is nullptr, no need to delete.";
107 return;
108 }
109 if (tensor != nullptr) {
110 // find primitive exist
111 if (tensor_map.find(tensor) != tensor_map.end()) {
112 // if find, then delete
113 void *flatbuffer_ptr = tensor_map[tensor];
114 if (flatbuffer_ptr != nullptr) {
115 free(flatbuffer_ptr);
116 tensor_map[tensor] = nullptr;
117 tensor_map.erase(tensor_map.find(tensor));
118 }
119 }
120 }
121 }
122
ClearAllMemory()123 void MindIRMemoryManager::ClearAllMemory() {
124 std::lock_guard<std::mutex> lck(mutex);
125 ClearMap(primitive_map);
126 ClearMap(tensor_map);
127 }
128 } // namespace lite
129 } // namespace mindspore