!33737 [MS][LITE][parallel predict] fix pack weight

Merge pull request !33737 from yefeng/cherry-pick-1651150559

mindspore/lite/src/cxx_api/model_pool/model_pool.cc

@@ -367,7 +367,7 @@ Status ModelPool::Init(const std::string &model_path, const std::shared_ptr<Runn
     int numa_node_id = model_pool_context[i]->numa_id;
     auto ret = lite::PackWeightManager::GetInstance()->InitByBuf(graph_buf, size, numa_node_id);
     MS_CHECK_FALSE_MSG(ret != kSuccess, kLiteError, "InitWeightManagerByBuf failed.");
-    auto new_model_buf = lite::PackWeightManager::GetInstance()->GetNumaModelBuf(numa_node_id);
+    auto new_model_buf = lite::PackWeightManager::GetInstance()->GetNumaModelBuf(graph_buf, numa_node_id);
     MS_CHECK_TRUE_MSG(new_model_buf != nullptr, kLiteError, "get model buf is nullptr from PackWeightManager");
     model_worker = std::make_shared<ModelWorker>();
     if (model_worker == nullptr) {
@@ -641,6 +641,7 @@ Status ModelPool::Predict(const std::vector<MSTensor> &inputs, std::vector<MSTen
     // split batch
     auto status = PredictBySplitBatch(inputs, outputs, before, after, max_wait_worker_node_id);
     if (status != kSuccess) {
+      predict_task_mutex_.unlock();
       MS_LOG(ERROR) << "do split batch failed. ret=" << status;
       return kLiteError;
     }
@@ -665,6 +666,7 @@ Status ModelPool::Predict(const std::vector<MSTensor> &inputs, std::vector<MSTen
     auto predict_task = std::make_shared<PredictTask>(&inputs, outputs, before, after);
     if (predict_task == nullptr) {
       MS_LOG(ERROR) << "predict_task is nullptr.";
+      predict_task_mutex_.unlock();
       return kLiteNullptr;
     }
     predict_task_queue_->PushPredictTask(predict_task, max_wait_worker_node_id);

mindspore/lite/src/cxx_api/model_pool/model_worker.cc

@@ -159,6 +159,7 @@ Status ModelWorker::Predict(const std::vector<MSTensor> &inputs, std::vector<MST
   auto model_input = model_->GetInputs();
   if (model_input.size() != inputs.size()) {
     MS_LOG(ERROR) << "model input size is: " << model_input.size() << ", but get input size is: " << inputs.size();
+    available_ = true;
     return kLiteError;
   }
   auto resize_pair = GetModelResize(model_input, inputs);
@@ -168,6 +169,7 @@ Status ModelWorker::Predict(const std::vector<MSTensor> &inputs, std::vector<MST
     auto status = model_->Resize(model_->GetInputs(), dims);
     if (status != kSuccess) {
       MS_LOG(ERROR) << "model pool resize failed.";
+      available_ = true;
       return kLiteError;
     }
   }
@@ -187,6 +189,7 @@ Status ModelWorker::Predict(const std::vector<MSTensor> &inputs, std::vector<MST
   auto status = model_->Predict(model_input, &model_output, before, after);
   if (status != kSuccess) {
     MS_LOG(ERROR) << "model predict failed.";
+    available_ = true;
     return status;
   }
   for (size_t i = 0; i < model_input.size(); i++) {
@@ -195,6 +198,7 @@ Status ModelWorker::Predict(const std::vector<MSTensor> &inputs, std::vector<MST
   if (need_copy_output_) {
     status = CopyOutputTensor(model_output, outputs);
     if (status != kSuccess) {
+      available_ = true;
       return kLiteError;
     }
   } else {

mindspore/lite/src/pack_weight.cc

@@ -18,7 +18,9 @@ namespace mindspore::lite {
 
 STATUS PackWeight::InitWeightManagerByBuf(const char *model_buf, size_t model_size, int numa_id) {
   MS_CHECK_TRUE_MSG(model_buf != nullptr, RET_ERROR, "model buf is nullptr in pack weight manager.");
-  if (numa_model_buf_.find(numa_id) != numa_model_buf_.end()) {
+  if (model_buf_map_.find(model_buf) != model_buf_map_.end() &&
+      find(numa_model_buf_[model_buf].begin(), numa_model_buf_[model_buf].end(), numa_id) !=
+        numa_model_buf_[model_buf].end()) {
     MS_LOG(DEBUG) << "same numa id, use same model buf.";
     return RET_OK;
   }
@@ -34,7 +36,7 @@ STATUS PackWeight::InitWeightManagerByBuf(const char *model_buf, size_t model_si
     return RET_ERROR;
   }
   memcpy(new_model_buf, model_buf, model_size);
-  numa_model_buf_.insert(std::make_pair(numa_id, new_model_buf));
+  numa_model_buf_[model_buf] = {numa_id};
   auto *model_const_weight = new (std::nothrow) ModelConstWeight();
   if (model_const_weight == nullptr) {
     MS_LOG(ERROR) << "model const weight is nullptr.";
@@ -43,15 +45,18 @@ STATUS PackWeight::InitWeightManagerByBuf(const char *model_buf, size_t model_si
   model_const_weight->numa_id = numa_id;
   buf_model_weight_[new_model_buf] = model_const_weight;
   buf_model_weight_[new_model_buf]->allocator = allocator;
+  model_buf_map_.insert(std::make_pair(model_buf, new_model_buf));
   return RET_OK;
 }
 
-char *PackWeight::GetNumaModelBuf(int numa_id) {
-  if (numa_model_buf_.find(numa_id) == numa_model_buf_.end()) {
+char *PackWeight::GetNumaModelBuf(const char *model_buf, int numa_id) {
+  if (model_buf_map_.find(model_buf) == model_buf_map_.end() ||
+      find(numa_model_buf_[model_buf].begin(), numa_model_buf_[model_buf].end(), numa_id) ==
+        numa_model_buf_[model_buf].end()) {
     MS_LOG(ERROR) << "can not find numa id in saved model buf.";
     return nullptr;
   }
-  return numa_model_buf_[numa_id];
+  return model_buf_map_[model_buf];
 }
 
 STATUS PackWeight::StoreOriginTensorData(const char *model_buf, const void *origin_tensor_data) {

mindspore/lite/src/pack_weight.h

@@ -40,7 +40,7 @@ class PackWeight {
   PackWeight() = default;
   ~PackWeight();
   STATUS InitWeightManagerByBuf(const char *model_buf, size_t model_size, int numa_id = -1);
-  char *GetNumaModelBuf(int numa_id);
+  char *GetNumaModelBuf(const char *model_buf, int numa_id);
   STATUS StoreOriginTensorData(const char *model_buf, const void *origin_tensor_data);
   void *GetPackData(const void *tensor_data, const size_t size, bool *is_packed);
 
@@ -49,7 +49,8 @@ class PackWeight {
 
   std::mutex mtx_weight_;
   std::unordered_map<const char *, ModelConstWeight *> buf_model_weight_;
-  std::unordered_map<int, char *> numa_model_buf_;
+  std::unordered_map<const char *, std::vector<int>> numa_model_buf_;
+  std::unordered_map<const char *, char *> model_buf_map_;
 };
 }  // namespace mindspore::lite
 #endif  // MINDSPORE_LITE_SRC_PACK_WEIGHT_H_

mindspore/lite/src/pack_weight_manager.cc

@@ -43,9 +43,9 @@ STATUS PackWeightManager::InitByBuf(const char *model_buf, size_t model_size, in
   return RET_OK;
 }
 
-char *PackWeightManager::GetNumaModelBuf(int numa_id) {
+char *PackWeightManager::GetNumaModelBuf(const char *model_buf, int numa_id) {
 #ifdef SHARING_MODEL_WEIGHT
-  return pack_weight_->GetNumaModelBuf(numa_id);
+  return pack_weight_->GetNumaModelBuf(model_buf, numa_id);
 #endif
   return nullptr;
 }

mindspore/lite/src/pack_weight_manager.h

@@ -29,7 +29,7 @@ class PackWeightManager {
   static PackWeightManager *GetInstance();
   ~PackWeightManager() = default;
   STATUS InitByBuf(const char *model_buf, size_t model_size, int numa_id = -1);
-  char *GetNumaModelBuf(int numa_id);
+  char *GetNumaModelBuf(const char *model_buf, int numa_id);
   STATUS StoreOriginTensorData(Model *model);
   void *GetPackData(const void *tensor_data, const size_t size, bool *is_packed);
   void Free(void *tensor_data);

mindspore/lite/src/runtime/kernel/cpu/fp32/convolution_slidewindow_fp32.cc

@@ -213,7 +213,6 @@ int ConvolutionSWCPUKernel::MallocWeightBiasData() {
       MS_LOG(ERROR) << "malloc packed weight failed.";
       return RET_NULL_PTR;
     }
-    memset(packed_weight_, 0, pack_weight_size * sizeof(float));
   }
 
   if (in_tensors_.size() == kInputSize2) {