1 /**
2 * Copyright 2023 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include <memory>
18 #include <atomic>
19 #include <unordered_map>
20 #include <map>
21 #include <mutex>
22 #include "src/litert/delegate/nnrt/nnrt_allocator.h"
23 #include "src/litert/delegate/nnrt/nnrt_utils.h"
24 #include "src/common/log.h"
25
26 namespace mindspore {
27 namespace lite {
GetInstance()28 std::shared_ptr<NNRTAllocator> NNRTAllocator::GetInstance() {
29 static std::shared_ptr<NNRTAllocator> instance(new (std::nothrow) NNRTAllocator());
30 return instance;
31 }
32
ClearFreeList()33 void NNRTAllocator::ClearFreeList() {
34 std::lock_guard<std::mutex> locker(mutex_);
35 for (auto &it : free_list_) {
36 auto membuf = it.second;
37 if (membuf == nullptr) {
38 MS_LOG(ERROR) << "membuf in free_list_ is nullptr.";
39 continue;
40 }
41 OH_NNTensor_Destroy(&membuf->tensor_);
42 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_);
43 delete membuf;
44 membuf = nullptr;
45 }
46 free_list_.clear();
47 }
48
FreeAllocatedTensor(void * data,lite::Tensor * tensor)49 void NNRTAllocator::FreeAllocatedTensor(void *data, lite::Tensor *tensor) {
50 std::lock_guard<std::mutex> locker(mutex_);
51 auto iter = allocated_list_.find(data);
52 if (iter == allocated_list_.end()) {
53 return;
54 }
55
56 auto iter_tensor = allocated_lite_tensors_.find(data);
57 if (iter_tensor == allocated_lite_tensors_.end() || iter_tensor->second != tensor) {
58 MS_LOG(INFO) << "data: " << data << " not belong to lite tensor: " << tensor << ", skip free";
59 return;
60 }
61
62 auto membuf = iter->second;
63 if (membuf == nullptr) {
64 MS_LOG(ERROR) << "membuf in allocated_list_ is nullptr, data: " << data;
65 return;
66 }
67 membuf->ref_count_ = 0;
68 (void)allocated_list_.erase(iter);
69 (void)allocated_lite_tensors_.erase(data);
70 OH_NNTensor_Destroy(&membuf->tensor_);
71 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_);
72 delete membuf;
73 membuf = nullptr;
74 data = nullptr;
75 }
76
~NNRTAllocator()77 NNRTAllocator::~NNRTAllocator() {
78 std::lock_guard<std::mutex> locker(mutex_);
79 for (auto &it : allocated_list_) {
80 auto membuf = it.second;
81 if (membuf != nullptr) {
82 MS_LOG(ERROR) << "NN_Tensor is not released, may lead to memory leak, data ptr: " << membuf->data << ", size: "
83 << membuf->size;
84 }
85 }
86
87 for (auto &it : free_list_) {
88 auto membuf = it.second;
89 if (membuf != nullptr) {
90 MS_LOG(ERROR) << "NN_Tensor is not released, may lead to memory leak, data ptr: " << membuf->data << ", size: "
91 << membuf->size;
92 }
93 }
94 }
95
SetTensorDesc(NN_TensorDesc * tensor_desc,const std::vector<int> & shape,const TypeId data_type,const Format format,const std::string & name)96 OH_NN_ReturnCode NNRTAllocator::SetTensorDesc(NN_TensorDesc *tensor_desc, const std::vector<int> &shape,
97 const TypeId data_type, const Format format, const std::string &name) {
98 OH_NN_ReturnCode ret = OH_NNTensorDesc_SetShape(tensor_desc, shape.data(), shape.size());
99 if (ret != OH_NN_SUCCESS) {
100 MS_LOG(ERROR) << "OH_NNTensorDesc_SetShape failed, shape: " << shape << ", name: " << name;
101 return ret;
102 }
103 ret = OH_NNTensorDesc_SetDataType(tensor_desc, CastToNNRtDataType(data_type));
104 if (ret != OH_NN_SUCCESS) {
105 MS_LOG(ERROR) << "OH_NNTensorDesc_SetDataType failed, data_type: " << data_type << ", name: " << name;
106 return ret;
107 }
108 ret = OH_NNTensorDesc_SetFormat(tensor_desc, CastToNNRtFormat(format));
109 if (ret != OH_NN_SUCCESS) {
110 MS_LOG(ERROR) << "OH_NNTensorDesc_SetFormat failed, format: " << format << ", name: " << name;
111 return ret;
112 }
113 ret = OH_NNTensorDesc_SetName(tensor_desc, name.c_str());
114 if (ret != OH_NN_SUCCESS) {
115 MS_LOG(ERROR) << "OH_NNTensorDesc_SetName failed, name: " << name;
116 return ret;
117 }
118 return ret;
119 }
120
CreateNNRtTensorDesc(const std::vector<int> & shape,const TypeId data_type,const Format format,const std::string & name)121 NN_TensorDesc *NNRTAllocator::CreateNNRtTensorDesc(const std::vector<int> &shape, const TypeId data_type,
122 const Format format, const std::string &name) {
123 auto tensor_desc = OH_NNTensorDesc_Create();
124 if (tensor_desc == nullptr) {
125 MS_LOG(ERROR) << "OH_NNTensorDesc_Create failed, name: " << name;
126 return nullptr;
127 }
128 OH_NN_ReturnCode ret = SetTensorDesc(tensor_desc, shape, data_type, format, name);
129 if (ret != OH_NN_SUCCESS) {
130 MS_LOG(ERROR) << "SetTensorDesc failed, name: " << name;
131 OH_NNTensorDesc_Destroy(&tensor_desc);
132 return nullptr;
133 }
134 return tensor_desc;
135 }
136
MallocByDesc(size_t size,const std::vector<int> & shape,const TypeId data_type,const Format format,const std::string & name)137 void *NNRTAllocator::MallocByDesc(size_t size, const std::vector<int> &shape, const TypeId data_type,
138 const Format format, const std::string &name) {
139 std::lock_guard<std::mutex> locker(mutex_);
140 auto iter = free_list_.lower_bound(size);
141 if (iter != free_list_.end() && (size == iter->second->size)) {
142 auto membuf = iter->second;
143 OH_NN_ReturnCode ret = SetTensorDesc(membuf->tensor_desc_, shape, data_type, format, name);
144 if (ret != OH_NN_SUCCESS) {
145 MS_LOG(ERROR) << "SetTensorDesc failed, name: " << name;
146 } else {
147 membuf->ref_count_ = 0;
148 (void)free_list_.erase(iter);
149 allocated_list_[membuf->data] = membuf;
150 return membuf->data;
151 }
152 }
153
154 auto membuf = new (std::nothrow) MemBuf();
155 if (membuf == nullptr) {
156 MS_LOG(ERROR) << "new Membuf failed.";
157 return nullptr;
158 }
159 membuf->ref_count_ = 0;
160 membuf->tensor_desc_ = CreateNNRtTensorDesc(shape, data_type, format, name);
161 if (membuf->tensor_desc_ == nullptr) {
162 MS_LOG(ERROR) << "create NN_TensorDesc failed.";
163 delete membuf;
164 return nullptr;
165 }
166 membuf->tensor_ = OH_NNTensor_Create(device_id_, membuf->tensor_desc_);
167 if (membuf->tensor_ == nullptr) {
168 MS_LOG(ERROR) << "OH_NNTensor_CreateWithSize failed, name: " << name;
169 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_);
170 delete membuf;
171 return nullptr;
172 }
173 membuf->data = OH_NNTensor_GetDataBuffer(membuf->tensor_);
174 if (membuf->data == nullptr) {
175 MS_LOG(ERROR) << "OH_NNTensor_GetDataBuffer failed, name: " << name;
176 OH_NNTensor_Destroy(&membuf->tensor_);
177 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_);
178 delete membuf;
179 return nullptr;
180 }
181
182 membuf->size = size;
183 allocated_list_[membuf->data] = membuf;
184 return membuf->data;
185 }
186
Malloc(size_t size)187 void *NNRTAllocator::Malloc(size_t size) {
188 MS_LOG(ERROR) << "NNRt Allocator is not support malloc by size.";
189 return nullptr;
190 }
191
Free(void * ptr)192 void NNRTAllocator::Free(void *ptr) {
193 if (ptr == nullptr) {
194 return;
195 }
196
197 std::lock_guard<std::mutex> locker(mutex_);
198 auto iter = allocated_list_.find(ptr);
199 if (iter == allocated_list_.end()) {
200 return;
201 }
202 auto membuf = iter->second;
203 membuf->ref_count_ = 0;
204 (void)allocated_list_.erase(iter);
205 (void)allocated_lite_tensors_.erase(ptr);
206 (void)free_list_.insert(std::make_pair(membuf->size, membuf));
207 }
208
RefCount(void * ptr)209 int NNRTAllocator::RefCount(void *ptr) {
210 if (ptr == nullptr) {
211 return NNRT_ALLOCATION;
212 }
213 std::lock_guard<std::mutex> locker(mutex_);
214 auto iter = allocated_list_.find(ptr);
215 if (iter != allocated_list_.end()) {
216 auto membuf = iter->second;
217 int ref_count = std::atomic_load(&membuf->ref_count_);
218 return ref_count;
219 }
220 return -1;
221 }
222
SetRefCount(void * ptr,int ref_count)223 int NNRTAllocator::SetRefCount(void *ptr, int ref_count) {
224 if (ptr == nullptr) {
225 return -1;
226 }
227 std::lock_guard<std::mutex> locker(mutex_);
228 auto iter = allocated_list_.find(ptr);
229 if (iter != allocated_list_.end()) {
230 auto membuf = iter->second;
231 std::atomic_store(&membuf->ref_count_, ref_count);
232 return ref_count;
233 }
234 return -1;
235 }
236
DecRefCount(void * ptr,int ref_count)237 int NNRTAllocator::DecRefCount(void *ptr, int ref_count) {
238 if (ptr == nullptr) {
239 return -1;
240 }
241 std::lock_guard<std::mutex> locker(mutex_);
242 auto iter = allocated_list_.find(ptr);
243 if (iter != allocated_list_.end()) {
244 auto membuf = iter->second;
245 std::atomic_fetch_sub(&membuf->ref_count_, ref_count);
246 return membuf->ref_count_;
247 }
248 return -1;
249 }
250
IncRefCount(void * ptr,int ref_count)251 int NNRTAllocator::IncRefCount(void *ptr, int ref_count) {
252 if (ptr == nullptr) {
253 return -1;
254 }
255 std::lock_guard<std::mutex> locker(mutex_);
256 auto iter = allocated_list_.find(ptr);
257 if (iter != allocated_list_.end()) {
258 auto membuf = iter->second;
259 std::atomic_fetch_add(&membuf->ref_count_, ref_count);
260 return membuf->ref_count_;
261 }
262 return -1;
263 }
264 } // namespace lite
265 } // namespace mindspore