!29217 【MS】【LITE】support dynamic memory reused

Merge pull request !29217 from chenjianping/master_dev

mindspore/lite/CMakeLists.txt

@@ -40,6 +40,7 @@ option(MSLITE_ENABLE_RUNTIME_GLOG "enable runtime glog" off)
 option(MSLITE_ENABLE_COVERAGE "enable code coverage" off)
 option(MSLITE_ENABLE_SHARING_MEM_WITH_OPENGL "enable sharing memory with OpenGL" off)
 option(MSLITE_ENABLE_SERVING "enable serving" off)
+option(MSLITE_ENABLE_SERVER_INFERENCE "enable inference on server" off)
 
 #Option that can be configured through manually
 option(ENABLE_VERBOSE "" off)
@@ -144,6 +145,9 @@ endif()
 if(DEFINED ENV{MSLITE_ENABLE_SERVING})
     set(MSLITE_ENABLE_SERVING $ENV{MSLITE_ENABLE_SERVING})
 endif()
+if(DEFINED ENV{MSLITE_ENABLE_SERVER_INFERENCE})
+    set(MSLITE_ENABLE_SERVER_INFERENCE $ENV{MSLITE_ENABLE_SERVER_INFERENCE})
+endif()
 
 if(MACHINE_LINUX_ARM64)
     add_compile_definitions(MACHINE_LINUX_ARM64)

mindspore/lite/src/CMakeLists.txt

@@ -99,6 +99,7 @@ set(LITE_SRC
         ${CMAKE_CURRENT_SOURCE_DIR}/common/lite_utils.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/common/prim_util.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/common/tensor_util.cc
+        ${CMAKE_CURRENT_SOURCE_DIR}/runtime/allocator.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/runtime/inner_allocator.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/runtime/runtime_allocator.cc
         ${CMAKE_CURRENT_SOURCE_DIR}/runtime/infer_manager.cc
@@ -121,6 +122,14 @@ set(LITE_SRC
         ${CMAKE_CURRENT_SOURCE_DIR}/pack_weight_manager.cc
         )
 
+if(MSLITE_ENABLE_SERVER_INFERENCE)
+set(LITE_SRC
+        ${LITE_SRC}
+        ${CMAKE_CURRENT_SOURCE_DIR}/runtime/dynamic_mem_allocator.cc
+        ${CMAKE_CURRENT_SOURCE_DIR}/runtime/dynamic_mem_manager.cc
+        )
+endif()
+
 if(MSLITE_ENABLE_CONTROLFLOW)
     file(GLOB CONTROL_FLOW_SRC
             ${CMAKE_CURRENT_SOURCE_DIR}/control_flow/*.cc

mindspore/lite/src/runtime/allocator.cc

@@ -0,0 +1,30 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifdef MSLITE_ENABLE_SERVER_INFERENCE
+#include "src/runtime/dynamic_mem_allocator.h"
+#else
+#include "src/runtime/inner_allocator.h"
+#endif
+
+namespace mindspore {
+std::shared_ptr<Allocator> Allocator::Create() {
+#ifdef MSLITE_ENABLE_SERVER_INFERENCE
+  return std::make_shared<DynamicMemAllocator>();
+#else
+  return std::make_shared<DefaultAllocator>();
+#endif
+}
+}  // namespace mindspore

mindspore/lite/src/runtime/dynamic_mem_allocator.cc

@@ -0,0 +1,54 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "src/runtime/dynamic_mem_allocator.h"
+#include <memory>
+#include "src/runtime/dynamic_mem_manager.h"
+#include "src/common/utils.h"
+#include "src/common/log_adapter.h"
+
+namespace mindspore {
+void *DynamicMemAllocator::Malloc(size_t size) { return DynamicMemManager::GetInstance()->Malloc(size); }
+
+void DynamicMemAllocator::Free(void *ptr) { DynamicMemManager::GetInstance()->Free(ptr); }
+
+int DynamicMemAllocator::RefCount(void *ptr) {
+  if (ptr == nullptr) {
+    return -1;
+  }
+  return DynamicMemManager::GetInstance()->RefCount(ptr);
+}
+
+int DynamicMemAllocator::SetRefCount(void *ptr, int ref_count) {
+  if (ptr == nullptr) {
+    return -1;
+  }
+  return DynamicMemManager::GetInstance()->SetRefCount(ptr, ref_count);
+}
+
+int DynamicMemAllocator::IncRefCount(void *ptr, int ref_count) {
+  if (ptr == nullptr) {
+    return -1;
+  }
+  return DynamicMemManager::GetInstance()->IncRefCount(ptr, ref_count);
+}
+
+int DynamicMemAllocator::DecRefCount(void *ptr, int ref_count) {
+  if (ptr == nullptr) {
+    return -1;
+  }
+  return DynamicMemManager::GetInstance()->DecRefCount(ptr, ref_count);
+}
+}  // namespace mindspore

mindspore/lite/src/runtime/dynamic_mem_allocator.h

@@ -0,0 +1,39 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_LITE_SRC_RUNTIME_DYNAMIC_MEM_ALLOCATOR_H_
+#define MINDSPORE_LITE_SRC_RUNTIME_DYNAMIC_MEM_ALLOCATOR_H_
+
+#include <mutex>
+#include <map>
+#include <unordered_map>
+#include "include/api/allocator.h"
+
+namespace mindspore {
+class DynamicMemAllocator : public Allocator {
+ public:
+  DynamicMemAllocator() = default;
+  virtual ~DynamicMemAllocator() = default;
+  void *Malloc(size_t size) override;
+  void Free(void *ptr) override;
+  int RefCount(void *ptr) override;
+  int SetRefCount(void *ptr, int ref_count) override;
+  int IncRefCount(void *ptr, int ref_count) override;
+  int DecRefCount(void *ptr, int ref_count) override;
+};
+}  // namespace mindspore
+
+#endif  // MINDSPORE_LITE_SRC_RUNTIME_DYNAMIC_MEM_ALLOCATOR_H_

mindspore/lite/src/runtime/dynamic_mem_manager.cc

@@ -0,0 +1,264 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "src/runtime/dynamic_mem_manager.h"
+#include "src/common/log_adapter.h"
+#include "src/common/utils.h"
+
+namespace mindspore {
+namespace {
+// Alloc memory aligned according to 64 bytes.
+static constexpr size_t kMemAlginSize = 64;
+
+// The minimum unit size (512M) of memory block used for dynamic extend.
+static constexpr size_t kAllocUnitSize = 536870912;
+
+static constexpr size_t kBlockSize = 1024;
+// invalid block index
+static constexpr int64_t kInvalidIndex = -1;
+
+size_t Rounded(size_t size) { return (size + kMemAlginSize - 1) & (~(kMemAlginSize - 1)); }
+
+void *Allocate(size_t allocate_size) {
+  if (allocate_size > lite::GetMaxMallocSize()) {
+    MS_LOG(ERROR) << "MallocData out of max_size, size: " << allocate_size;
+    return nullptr;
+  }
+  void *data = nullptr;
+#ifdef _WIN32
+  data = _aligned_malloc(allocate_size, kMemAlginSize);
+#else
+  auto ret = posix_memalign(&data, kMemAlginSize, allocate_size);
+  if (UNLIKELY(ret != 0)) {
+    MS_LOG(ERROR) << "posix_memalign failed!ret: " << ret;
+    return nullptr;
+  }
+#endif
+  if (UNLIKELY(data == nullptr)) {
+    MS_LOG(ERROR) << "malloc data failed!";
+    return nullptr;
+  }
+  return data;
+}
+}  // namespace
+
+DynamicMemManager::DynamicMemManager() {
+  blocks_.resize(kBlockSize);
+  garbage_block_ = kInvalidIndex;
+  auto *block = GetBlock();
+  block->data_ = Allocate(kAllocUnitSize);
+  if (UNLIKELY(block->data_ == nullptr)) {
+    return;
+  }
+  all_datas_.emplace_back(block->data_);
+  block->size_ = kAllocUnitSize;
+  free_blocks_.emplace(kAllocUnitSize, block->index_);
+}
+
+Block *DynamicMemManager::GetBlock() {
+  Block *block;
+  if (garbage_block_ != kInvalidIndex) {
+    block = &blocks_[garbage_block_];
+    garbage_block_ = blocks_[garbage_block_].next_index_;
+  } else {
+    if (block_count_ >= blocks_.size()) {
+      blocks_.resize(blocks_.size() + kBlockSize);
+    }
+    blocks_[block_count_].index_ = block_count_;
+    block = &blocks_[block_count_++];
+  }
+  block->used_ = false;
+  block->ref_count_ = 0;
+  block->pre_index_ = kInvalidIndex;
+  block->next_index_ = kInvalidIndex;
+  return block;
+}
+
+void DynamicMemManager::AddGarbageBlock(const int64_t index) {
+  blocks_[index].next_index_ = garbage_block_;
+  garbage_block_ = index;
+}
+
+// malloc memory for data storage
+void *DynamicMemManager::Malloc(size_t size) {
+  auto rounded_size = Rounded(size);
+  std::lock_guard<std::mutex> locker(mutex_);
+  auto iter = free_blocks_.lower_bound(rounded_size);
+  if (iter != free_blocks_.end()) {
+    auto index = iter->second;
+    free_blocks_.erase(iter);
+    auto *block = &blocks_[index];
+    block->used_ = true;
+    datas_.emplace(block->data_, index);
+    if (block->size_ > rounded_size) {
+      Block *block_next = GetBlock();
+      block_next->size_ = block->size_ - rounded_size;
+      block->size_ = rounded_size;
+      block_next->data_ = static_cast<int8_t *>(block->data_) + rounded_size;
+      block_next->pre_index_ = index;
+      auto next_index = block->next_index_;
+      block_next->next_index_ = next_index;
+      if (next_index != kInvalidIndex) {
+        blocks_[next_index].pre_index_ = block_next->index_;
+      }
+      block->next_index_ = block_next->index_;
+      free_blocks_.emplace(block_next->size_, block_next->index_);
+    }
+    return block->data_;
+  }
+  // todo kAllocUnitSize can be replaced by config
+  auto allocate_size = rounded_size < kAllocUnitSize ? kAllocUnitSize : rounded_size;
+  void *data = Allocate(allocate_size);
+  if (UNLIKELY(data == nullptr)) {
+    return nullptr;
+  }
+  all_datas_.emplace_back(data);
+  Block *block = GetBlock();
+  block->size_ = rounded_size;
+  block->data_ = data;
+  block->used_ = true;
+  datas_.emplace(data, block->index_);
+  if (allocate_size > rounded_size) {
+    Block *block_next = GetBlock();
+    block_next->data_ = static_cast<int8_t *>(data) + rounded_size;
+    block_next->size_ = allocate_size - rounded_size;
+    block_next->pre_index_ = block->index_;
+    block->next_index_ = block_next->index_;
+    free_blocks_.emplace(block_next->size_, block_next->index_);
+  }
+  return data;
+}
+
+// return memory to the memory pool
+void DynamicMemManager::Free(void *ptr) {
+  if (UNLIKELY(ptr == nullptr)) {
+    return;
+  }
+  std::lock_guard<std::mutex> locker(mutex_);
+  auto iter = datas_.find(ptr);
+  if (iter == datas_.end()) {
+    return;
+  }
+
+  auto index = iter->second;
+  datas_.erase(iter);
+  Block *block = &blocks_[index];
+  auto next_index = block->next_index_;
+  if (next_index != kInvalidIndex && !blocks_[next_index].used_) {
+    EraseFreeBlock(next_index);
+    block->size_ += blocks_[next_index].size_;
+    auto next_next_index = blocks_[next_index].next_index_;
+    if (next_next_index != kInvalidIndex) {
+      blocks_[next_next_index].pre_index_ = block->index_;
+    }
+    block->next_index_ = next_next_index;
+    block->used_ = false;
+    block->ref_count_ = 0;
+    free_blocks_.emplace(block->size_, block->index_);
+    AddGarbageBlock(next_index);
+  }
+  auto pre_index = block->pre_index_;
+  if (pre_index != kInvalidIndex && !blocks_[pre_index].used_) {
+    EraseFreeBlock(pre_index);
+    if (!block->used_) {
+      EraseFreeBlock(index);
+    }
+    blocks_[pre_index].size_ += block->size_;
+    next_index = block->next_index_;
+    blocks_[pre_index].next_index_ = next_index;
+    if (next_index != kInvalidIndex) {
+      blocks_[next_index].pre_index_ = pre_index;
+    }
+    block->used_ = false;
+    block->ref_count_ = 0;
+    free_blocks_.emplace(blocks_[pre_index].size_, pre_index);
+    AddGarbageBlock(index);
+  }
+  if (block->used_) {
+    block->used_ = false;
+    block->ref_count_ = 0;
+    free_blocks_.emplace(block->size_, block->index_);
+  }
+}
+
+void DynamicMemManager::EraseFreeBlock(const int64_t index) {
+  auto range = free_blocks_.equal_range(blocks_[index].size_);
+  for (auto item = range.first; item != range.second; ++item) {
+    if (item->second == index) {
+      free_blocks_.erase(item);
+      break;
+    }
+  }
+}
+
+DynamicMemManager::~DynamicMemManager() {
+  MS_LOG(DEBUG) << "~DynamicMemManager() begin.";
+  for (auto &&data : all_datas_) {
+#ifdef _WIN32
+    _aligned_free(data);
+#else
+    free(data);
+#endif
+    data = nullptr;
+  }
+  free_blocks_.clear();
+  all_datas_.clear();
+  blocks_.clear();
+  MS_LOG(DEBUG) << "~DynamicMemManager() end.";
+}
+
+int DynamicMemManager::SetRefCount(void *ptr, int ref_count) {
+  std::lock_guard<std::mutex> locker(mutex_);
+  auto iter = datas_.find(ptr);
+  if (iter != datas_.end()) {
+    auto index = iter->second;
+    blocks_[index].ref_count_ = ref_count;
+    return ref_count;
+  }
+  return -1;
+}
+
+int DynamicMemManager::IncRefCount(void *ptr, int ref_count) {
+  std::lock_guard<std::mutex> locker(mutex_);
+  auto iter = datas_.find(ptr);
+  if (iter != datas_.end()) {
+    auto index = iter->second;
+    blocks_[index].ref_count_ += ref_count;
+    return blocks_[index].ref_count_;
+  }
+  return -1;
+}
+
+int DynamicMemManager::DecRefCount(void *ptr, int ref_count) {
+  std::lock_guard<std::mutex> locker(mutex_);
+  auto iter = datas_.find(ptr);
+  if (iter != datas_.end()) {
+    auto index = iter->second;
+    blocks_[index].ref_count_ -= ref_count;
+    return blocks_[index].ref_count_;
+  }
+  return -1;
+}
+
+int DynamicMemManager::RefCount(void *ptr) {
+  std::lock_guard<std::mutex> locker(mutex_);
+  auto iter = datas_.find(ptr);
+  if (iter != datas_.end()) {
+    return blocks_[iter->second].ref_count_;
+  }
+  return -1;
+}
+}  // namespace mindspore

mindspore/lite/src/runtime/dynamic_mem_manager.h

@@ -0,0 +1,77 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef MINDSPORE_LITE_SRC_RUNTIME_DYNAMIC_MEM_MANAGER_H_
+#define MINDSPORE_LITE_SRC_RUNTIME_DYNAMIC_MEM_MANAGER_H_
+
+#include <memory>
+#include <vector>
+#include <mutex>
+#include <map>
+#include <unordered_map>
+#include <deque>
+
+#define UNLIKELY(x) __builtin_expect(!!(x), 0)
+
+namespace mindspore {
+struct Block {
+  // used_ may be true when ref_count_ == 0
+  bool used_ = false;
+  void *data_ = nullptr;
+  size_t size_;
+  int64_t index_;
+  int64_t ref_count_ = 0;
+  int64_t pre_index_ = -1;
+  int64_t next_index_ = -1;
+};
+
+class DynamicMemManager {
+ public:
+  static DynamicMemManager *GetInstance() {
+    static DynamicMemManager instance;
+    return &instance;
+  }
+
+  virtual ~DynamicMemManager();
+
+  void *Malloc(size_t size);
+  void Free(void *ptr);
+  int SetRefCount(void *ptr, int ref_count);
+  int IncRefCount(void *ptr, int ref_count);
+  int DecRefCount(void *ptr, int ref_count);
+  int RefCount(void *ptr);
+
+ private:
+  DynamicMemManager();
+  Block *GetBlock();
+  void EraseFreeBlock(const int64_t index);
+  void AddGarbageBlock(const int64_t index);
+
+ private:
+  // all data blocks
+  size_t block_count_ = 0;
+  int64_t garbage_block_;
+  std::mutex mutex_;
+  std::vector<Block> blocks_;
+  // key: data size, value: Block index
+  std::multimap<size_t, int64_t> free_blocks_;
+  // key: data addr, value: Block index
+  std::unordered_map<void *, int64_t> datas_;
+  std::vector<void *> all_datas_;
+};
+}  // namespace mindspore
+
+#endif  // MINDSPORE_LITE_SRC_RUNTIME_DYNAMIC_MEM_MANAGER_H_

mindspore/lite/src/runtime/inner_allocator.cc

@@ -20,8 +20,6 @@
 #include "src/common/utils.h"
 
 namespace mindspore {
-std::shared_ptr<Allocator> Allocator::Create() { return std::make_shared<DefaultAllocator>(); }
-
 DefaultAllocator::DefaultAllocator(size_t aligned_size) { aligned_size_ = aligned_size; }
 
 DefaultAllocator::~DefaultAllocator() { Clear(); }

mindspore/lite/tools/converter/CMakeLists.txt

@@ -95,8 +95,11 @@ set(LITE_SRC ${API_SRC}
         ${SRC_DIR}/common/log.cc
         ${SRC_DIR}/common/prim_util.cc
         ${SRC_DIR}/common/tensor_util.cc
+        ${SRC_DIR}/runtime/allocator.cc
         ${SRC_DIR}/runtime/inner_allocator.cc
         ${SRC_DIR}/runtime/runtime_allocator.cc
+        ${SRC_DIR}/runtime/dynamic_mem_allocator.cc
+        ${SRC_DIR}/runtime/dynamic_mem_manager.cc
         ${SRC_DIR}/runtime/infer_manager.cc
         ${SRC_DIR}/runtime/runtime_pass.cc
         ${SRC_DIR}/inner_context.cc
@@ -122,6 +125,13 @@ set(LITE_SRC ${API_SRC}
         ${SRC_DIR}/pack_weight_manager.cc
         ${SRC_DIR}/delegate/tensorrt/distribution/distribution_base.cc
         )
+if(MSLITE_ENABLE_SERVER_INFERENCE)
+set(LITE_SRC
+        ${LITE_SRC}
+        ${SRC_DIR}/runtime/dynamic_mem_allocator.cc
+        ${SRC_DIR}/runtime/dynamic_mem_manager.cc
+        )
+endif()
 
 if(MSLITE_ENABLE_CONTROLFLOW)
     file(GLOB CONTROL_FLOW_SRC