!25606 [MSLITE] opt allocator resize bug

Merge pull request !25606 from ling/bug

mindspore/lite/src/lite_session.cc

@@ -975,6 +975,13 @@ int LiteSession::Resize(const std::vector<mindspore::tensor::MSTensor *> &inputs
     is_running_.store(false);
     return ret;
   }
+
+  if (RuntimeAllocatorInit() != RET_OK) {
+    MS_LOG(ERROR) << "Runtime allocator in resize failed.";
+    is_running_.store(false);
+    return RET_ERROR;
+  }
+
   is_running_.store(false);
   return RET_OK;
 }
@@ -1054,10 +1061,6 @@ int LiteSession::RuntimeAllocatorValid() {
     MS_LOG(DEBUG) << "Not support runtime allocator in runtime-infershape.";
     return RET_ERROR;
   }
-  if (is_control_flow_ == true) {
-    MS_LOG(DEBUG) << "Not support runtime allocator in control flow model.";
-    return RET_ERROR;
-  }
   if (kernels_.size() != 1) {
     MS_LOG(DEBUG) << "Not support runtime allocator in random subgraph sort";
     return RET_ERROR;
@@ -1166,8 +1169,11 @@ int LiteSession::RuntimeAllocatorInit() {
   if (RuntimeAllocatorValid() != RET_OK) {
     return RET_OK;
   }
-
-  runtime_allocator_ = std::shared_ptr<RuntimeAllocator>(new (std::nothrow) RuntimeAllocator());
+  if (runtime_allocator_ == nullptr) {
+    runtime_allocator_ = std::shared_ptr<RuntimeAllocator>(new (std::nothrow) RuntimeAllocator());
+  } else {
+    runtime_allocator_->Clear(context_->allocator);
+  }
   if (runtime_allocator_ == nullptr) {
     MS_LOG(ERROR) << "RuntimeAllocator is null.";
     return RET_ERROR;
@@ -1357,5 +1363,4 @@ int lite::LiteSession::CreateSessionByPath(const std::string &model_path, sessio
   (reinterpret_cast<lite::LiteSession *>(session))->set_model(model);
   return RET_OK;
 }
-
 }  // namespace mindspore

mindspore/lite/src/runtime/runtime_allocator.cc

@@ -79,6 +79,21 @@ void RuntimeAllocator::SetDataOffset(lite::Tensor *tensor, size_t offset) {
   return;
 }
 
+void RuntimeAllocator::Clear(AllocatorPtr default_allocator) {
+  total_size_ = 0;
+  for (auto iter : offset_map_) {
+    iter.first->set_allocator(default_allocator);
+    iter.first->set_data(nullptr);
+  }
+  if (data_ != nullptr) {
+    free(data_);
+    data_ = nullptr;
+  }
+  offset_map_.clear();
+  free_list_.clear();
+  used_list_.clear();
+}
+
 void RuntimeAllocator::MallocTensorData(lite::Tensor *tensor) {
   size_t size = tensor->Size();
   size_t offset = FindMinFree(size);

mindspore/lite/src/runtime/runtime_allocator.h

@@ -44,6 +44,7 @@ class RuntimeAllocator : public Allocator {
   void FreeTensorData(lite::Tensor *tensor);
   void *MallocOptData();
   const std::unordered_map<lite::Tensor *, size_t> &GetOffsetMap() const { return offset_map_; }
+  void Clear(AllocatorPtr default_allocator);
 
  private:
   size_t FindMinFree(size_t size);