1From decd3c1abfadace6022cc1baf1fa9505e1af1ca2 Mon Sep 17 00:00:00 2001 2From: chengfeng27 <chengfeng27@huawei.com> 3Date: Fri, 20 Sep 2024 11:49:09 +0800 4Subject: [PATCH] GetNNTensor unordered_map crash 5 6--- 7 .../src/litert/cxx_api/model/model_impl.cc | 6 +- 8 .../litert/delegate/nnrt/nnrt_allocator.cc | 114 +++++++++++++----- 9 .../src/litert/delegate/nnrt/nnrt_allocator.h | 14 ++- 10 .../litert/delegate/nnrt/nnrt_model_kernel.cc | 16 +-- 11 mindspore/lite/src/litert/lite_session.cc | 4 + 12 mindspore/lite/src/tensor.cc | 8 ++ 13 6 files changed, 114 insertions(+), 48 deletions(-) 14 15diff --git a/mindspore/lite/src/litert/cxx_api/model/model_impl.cc b/mindspore/lite/src/litert/cxx_api/model/model_impl.cc 16index 6a73a927..b7fa3c65 100644 17--- a/mindspore/lite/src/litert/cxx_api/model/model_impl.cc 18+++ b/mindspore/lite/src/litert/cxx_api/model/model_impl.cc 19@@ -321,7 +321,11 @@ Status ModelImpl::Build() { 20 21 static void ResetTensorData(std::vector<void *> old_data, const std::vector<lite::Tensor *> &tensors) { 22 for (size_t j = 0; j < old_data.size(); j++) { 23- tensors.at(j)->set_data(old_data.at(j)); 24+ auto tensor = tensors.at(j); 25+ if (old_data.at(j) != tensor->data()) { 26+ tensor->FreeData(); 27+ } 28+ tensor->set_data(old_data.at(j)); 29 } 30 } 31 32diff --git a/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.cc b/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.cc 33index 4910343f..c160315e 100644 34--- a/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.cc 35+++ b/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.cc 36@@ -25,53 +25,100 @@ 37 38 namespace mindspore { 39 namespace lite { 40-NNRTAllocator::~NNRTAllocator() { 41+std::shared_ptr<NNRTAllocator> NNRTAllocator::GetInstance() { 42+ static std::shared_ptr<NNRTAllocator> instance(new (std::nothrow) NNRTAllocator()); 43+ return instance; 44+} 45+ 46+void NNRTAllocator::ClearFreeList() { 47 std::lock_guard<std::mutex> locker(mutex_); 48- for (auto &it : allocated_list_) { 49+ for (auto &it : free_list_) { 50 auto membuf = it.second; 51+ if (membuf == nullptr) { 52+ MS_LOG(ERROR) << "membuf in free_list_ is nullptr."; 53+ continue; 54+ } 55 OH_NNTensor_Destroy(&membuf->tensor_); 56 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_); 57 delete membuf; 58+ membuf = nullptr; 59+ } 60+ free_list_.clear(); 61+} 62+ 63+void NNRTAllocator::FreeAllocatedTensor(void *data) { 64+ std::lock_guard<std::mutex> locker(mutex_); 65+ auto iter = allocated_list_.find(data); 66+ if (iter == allocated_list_.end()) { 67+ return; 68+ } 69+ auto membuf = iter->second; 70+ if (membuf == nullptr) { 71+ MS_LOG(ERROR) << "membuf in allocated_list_ is nullptr, data: " << data; 72+ return; 73+ } 74+ membuf->ref_count_ = 0; 75+ (void)allocated_list_.erase(iter); 76+ OH_NNTensor_Destroy(&membuf->tensor_); 77+ OH_NNTensorDesc_Destroy(&membuf->tensor_desc_); 78+ delete membuf; 79+ membuf = nullptr; 80+} 81+ 82+NNRTAllocator::~NNRTAllocator() { 83+ std::lock_guard<std::mutex> locker(mutex_); 84+ for (auto &it : allocated_list_) { 85+ auto membuf = it.second; 86+ if (membuf != nullptr) { 87+ MS_LOG(ERROR) << "NN_Tensor is not released, may lead to memory leak, data ptr: " << membuf->data << ", size: " 88+ << membuf->size; 89+ } 90 } 91- allocated_list_.clear(); 92 93 for (auto &it : free_list_) { 94 auto membuf = it.second; 95- OH_NNTensor_Destroy(&membuf->tensor_); 96- OH_NNTensorDesc_Destroy(&membuf->tensor_desc_); 97- delete membuf; 98+ if (membuf != nullptr) { 99+ MS_LOG(ERROR) << "NN_Tensor is not released, may lead to memory leak, data ptr: " << membuf->data << ", size: " 100+ << membuf->size; 101+ } 102 } 103- free_list_.clear(); 104 } 105 106-NN_TensorDesc *NNRTAllocator::CreateNNRtTensorDesc(const std::vector<int> &shape, const TypeId data_type, 107- const Format format, const std::string &name) { 108- auto tensor_desc = OH_NNTensorDesc_Create(); 109- if (tensor_desc == nullptr) { 110- MS_LOG(ERROR) << "OH_NNTensorDesc_Create failed, i = " << index_; 111- return nullptr; 112- } 113+OH_NN_ReturnCode NNRTAllocator::SetTensorDesc(NN_TensorDesc *tensor_desc, const std::vector<int> &shape, 114+ const TypeId data_type, const Format format, const std::string &name) { 115 OH_NN_ReturnCode ret = OH_NNTensorDesc_SetShape(tensor_desc, shape.data(), shape.size()); 116 if (ret != OH_NN_SUCCESS) { 117- MS_LOG(ERROR) << "OH_NNTensorDesc_SetShape failed, i = " << index_ << ", shape: " << shape; 118- OH_NNTensorDesc_Destroy(&tensor_desc); 119- return nullptr; 120+ MS_LOG(ERROR) << "OH_NNTensorDesc_SetShape failed, shape: " << shape << ", name: " << name; 121+ return ret; 122 } 123 ret = OH_NNTensorDesc_SetDataType(tensor_desc, CastToNNRtDataType(data_type)); 124 if (ret != OH_NN_SUCCESS) { 125- MS_LOG(ERROR) << "OH_NNTensorDesc_SetDataType failed, i = " << index_ << ", data_type: " << data_type; 126- OH_NNTensorDesc_Destroy(&tensor_desc); 127- return nullptr; 128+ MS_LOG(ERROR) << "OH_NNTensorDesc_SetDataType failed, data_type: " << data_type << ", name: " << name; 129+ return ret; 130 } 131 ret = OH_NNTensorDesc_SetFormat(tensor_desc, CastToNNRtFormat(format)); 132 if (ret != OH_NN_SUCCESS) { 133- MS_LOG(ERROR) << "OH_NNTensorDesc_SetFormat failed, i = " << index_ << ", format: " << format; 134- OH_NNTensorDesc_Destroy(&tensor_desc); 135- return nullptr; 136+ MS_LOG(ERROR) << "OH_NNTensorDesc_SetFormat failed, format: " << format << ", name: " << name; 137+ return ret; 138 } 139 ret = OH_NNTensorDesc_SetName(tensor_desc, name.c_str()); 140 if (ret != OH_NN_SUCCESS) { 141- MS_LOG(ERROR) << "OH_NNTensorDesc_SetName failed, i = " << index_ << ", name: " << name; 142+ MS_LOG(ERROR) << "OH_NNTensorDesc_SetName failed, name: " << name; 143+ return ret; 144+ } 145+ return ret; 146+} 147+ 148+NN_TensorDesc *NNRTAllocator::CreateNNRtTensorDesc(const std::vector<int> &shape, const TypeId data_type, 149+ const Format format, const std::string &name) { 150+ auto tensor_desc = OH_NNTensorDesc_Create(); 151+ if (tensor_desc == nullptr) { 152+ MS_LOG(ERROR) << "OH_NNTensorDesc_Create failed, name: " << name; 153+ return nullptr; 154+ } 155+ OH_NN_ReturnCode ret = SetTensorDesc(tensor_desc, shape, data_type, format, name); 156+ if (ret != OH_NN_SUCCESS) { 157+ MS_LOG(ERROR) << "SetTensorDesc failed, name: " << name; 158 OH_NNTensorDesc_Destroy(&tensor_desc); 159 return nullptr; 160 } 161@@ -82,12 +129,17 @@ void *NNRTAllocator::MallocByDesc(size_t size, const std::vector<int> &shape, co 162 const Format format, const std::string &name) { 163 std::lock_guard<std::mutex> locker(mutex_); 164 auto iter = free_list_.lower_bound(size); 165- if (iter != free_list_.end()) { 166+ if (iter != free_list_.end() && (size == iter->second->size)) { 167 auto membuf = iter->second; 168- membuf->ref_count_ = 0; 169- (void)free_list_.erase(iter); 170- allocated_list_[membuf->data] = membuf; 171- return membuf->data; 172+ OH_NN_ReturnCode ret = SetTensorDesc(membuf->tensor_desc_, shape, data_type, format, name); 173+ if (ret != OH_NN_SUCCESS) { 174+ MS_LOG(ERROR) << "SetTensorDesc failed, name: " << name; 175+ } else { 176+ membuf->ref_count_ = 0; 177+ (void)free_list_.erase(iter); 178+ allocated_list_[membuf->data] = membuf; 179+ return membuf->data; 180+ } 181 } 182 183 auto membuf = new (std::nothrow) MemBuf(); 184@@ -104,14 +156,14 @@ void *NNRTAllocator::MallocByDesc(size_t size, const std::vector<int> &shape, co 185 } 186 membuf->tensor_ = OH_NNTensor_Create(device_id_, membuf->tensor_desc_); 187 if (membuf->tensor_ == nullptr) { 188- MS_LOG(ERROR) << "OH_NNTensor_CreateWithSize failed, i = " << index_; 189+ MS_LOG(ERROR) << "OH_NNTensor_CreateWithSize failed, name: " << name; 190 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_); 191 delete membuf; 192 return nullptr; 193 } 194 membuf->data = OH_NNTensor_GetDataBuffer(membuf->tensor_); 195 if (membuf->data == nullptr) { 196- MS_LOG(ERROR) << "OH_NNTensor_GetDataBuffer failed, i = " << index_; 197+ MS_LOG(ERROR) << "OH_NNTensor_GetDataBuffer failed, name: " << name; 198 OH_NNTensor_Destroy(&membuf->tensor_); 199 OH_NNTensorDesc_Destroy(&membuf->tensor_desc_); 200 delete membuf; 201diff --git a/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.h b/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.h 202index ef27f307..9d7fa3c0 100644 203--- a/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.h 204+++ b/mindspore/lite/src/litert/delegate/nnrt/nnrt_allocator.h 205@@ -31,31 +31,35 @@ struct OH_NNExecutor; 206 207 namespace mindspore { 208 namespace lite { 209-enum MemoryCategory { NNRT_INPUT, NNRT_OUTPUT }; 210 211 class NNRTAllocator : public Allocator { 212 public: 213- NNRTAllocator(OH_NNExecutor *executor, int index, size_t device_id, MemoryCategory memory_category) 214- : index_(index), device_id_(device_id), memory_category_(memory_category), executor_(executor) {} 215+ NNRTAllocator() {} 216 ~NNRTAllocator() override; 217- 218+ static std::shared_ptr<NNRTAllocator> GetInstance(); 219 void *Malloc(size_t size) override; 220 void *MallocByDesc(size_t size, const std::vector<int> &shape, const TypeId data_type, const Format format, 221 const std::string &name); 222 NN_TensorDesc *CreateNNRtTensorDesc(const std::vector<int> &shape, const TypeId data_type, const Format format, 223 const std::string &name); 224+ OH_NN_ReturnCode SetTensorDesc(NN_TensorDesc *tensor_desc, const std::vector<int> &shape, const TypeId data_type, 225+ const Format format, const std::string &name); 226 void Free(void *ptr) override; 227 int RefCount(void *ptr) override; 228 int SetRefCount(void *ptr, int ref_count) override; 229 int DecRefCount(void *ptr, int ref_count) override; 230 int IncRefCount(void *ptr, int ref_count) override; 231 NN_Tensor *GetNNTensor(void *ptr) { 232+ std::lock_guard<std::mutex> locker(mutex_); 233 auto iter = allocated_list_.find(ptr); 234 if (iter != allocated_list_.end()) { 235 return iter->second->tensor_; 236 } 237 return nullptr; 238 } 239+ void SetDeviceId(size_t id) { device_id_ = id; } 240+ void ClearFreeList(); 241+ void FreeAllocatedTensor(void *data); 242 243 private: 244 struct MemBuf { 245@@ -66,9 +70,7 @@ class NNRTAllocator : public Allocator { 246 size_t size{0}; 247 }; 248 249- int index_{0}; 250 size_t device_id_{0}; 251- MemoryCategory memory_category_{NNRT_INPUT}; 252 OH_NNExecutor *executor_{nullptr}; 253 std::mutex mutex_; 254 // <membuf->memory_->data, membuf> 255diff --git a/mindspore/lite/src/litert/delegate/nnrt/nnrt_model_kernel.cc b/mindspore/lite/src/litert/delegate/nnrt/nnrt_model_kernel.cc 256index 1411020b..26c4633a 100644 257--- a/mindspore/lite/src/litert/delegate/nnrt/nnrt_model_kernel.cc 258+++ b/mindspore/lite/src/litert/delegate/nnrt/nnrt_model_kernel.cc 259@@ -24,20 +24,16 @@ constexpr auto kDynamicDims = "DynamicDims"; 260 } 261 262 int NNRTModelKernel::Prepare() { 263+ auto nnrt_allocator = lite::NNRTAllocator::GetInstance(); 264+ if (nnrt_allocator == nullptr) { 265+ MS_LOG(ERROR) << "Get NNRTAllocator failed"; 266+ return lite::RET_NULL_PTR; 267+ } 268+ nnrt_allocator->SetDeviceId(nnrt_device_info_.device_id_); 269 for (size_t i = 0; i < inputs_.size(); i++) { 270- auto nnrt_allocator = std::make_shared<lite::NNRTAllocator>(oh_nn_executor_, i, nnrt_device_info_.device_id_, lite::NNRT_INPUT); 271- if (nnrt_allocator == nullptr) { 272- MS_LOG(ERROR) << "Create NNRTAllocator failed"; 273- return lite::RET_NULL_PTR; 274- } 275 inputs_[i].SetAllocator(nnrt_allocator); 276 } 277 for (size_t i = 0; i < outputs_.size(); i++) { 278- auto nnrt_allocator = std::make_shared<lite::NNRTAllocator>(oh_nn_executor_, i, nnrt_device_info_.device_id_, lite::NNRT_OUTPUT); 279- if (nnrt_allocator == nullptr) { 280- MS_LOG(ERROR) << "Create NNRTAllocator failed"; 281- return lite::RET_NULL_PTR; 282- } 283 outputs_[i].SetAllocator(nnrt_allocator); 284 } 285 return lite::RET_OK; 286diff --git a/mindspore/lite/src/litert/lite_session.cc b/mindspore/lite/src/litert/lite_session.cc 287index f635c8d2..7502ec27 100644 288--- a/mindspore/lite/src/litert/lite_session.cc 289+++ b/mindspore/lite/src/litert/lite_session.cc 290@@ -69,6 +69,7 @@ 291 #include "src/litert/runtime_packed_node_pass.h" 292 #ifdef SUPPORT_NNRT 293 #include "src/litert/delegate/nnrt/nnrt_delegate.h" 294+#include "src/litert/delegate/nnrt/nnrt_allocator.h" 295 #endif 296 297 using AbstractBaseModel = mindspore::infer::AbstractBaseModel; 298@@ -1280,6 +1281,9 @@ LiteSession::~LiteSession() { 299 } 300 delete (model_); 301 model_ = nullptr; 302+#ifdef SUPPORT_NNRT 303+ NNRTAllocator::GetInstance()->ClearFreeList(); 304+#endif 305 is_running_.store(false); 306 } 307 308diff --git a/mindspore/lite/src/tensor.cc b/mindspore/lite/src/tensor.cc 309index 9d9a1491..6e97750a 100644 310--- a/mindspore/lite/src/tensor.cc 311+++ b/mindspore/lite/src/tensor.cc 312@@ -112,8 +112,16 @@ Tensor *Tensor::CopyTensor(const Tensor &src_tensor, bool copy_data, AllocatorPt 313 } 314 315 Tensor::~Tensor() { 316+#ifdef SUPPORT_NNRT 317+ void *allocated_data = this->tensor_c_.data_; 318+#endif 319 FreeData(); 320 this->tensor_c_.data_ = nullptr; 321+#ifdef SUPPORT_NNRT 322+ if (this->own_data_ && IS_NNRT_ALLOCATOR(allocator_)) { 323+ NNRTAllocator::GetInstance()->FreeAllocatedTensor(allocated_data); 324+ } 325+#endif 326 } 327 328 bool Tensor::operator==(const Tensor &tensor) { 329-- 3302.17.1 331 332