dynamic memory pool support multi-thread

mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.cc

@@ -28,6 +28,7 @@ DynamicMemPoolBestFit::~DynamicMemPoolBestFit() {
 
 DeviceMemPtr DynamicMemPoolBestFit::AllocTensorMem(size_t size) {
   size_t align_size = AlignMemorySize(size);
+  std::lock_guard<std::mutex> locker(mutex_);
   // Find the idle memory buf by tensor size, if not find, then add new memory block and memory buf.
   DeviceMemPtr device_addr = FindIdleMemBuf(align_size);
   if (!device_addr) {
@@ -44,6 +45,7 @@ std::vector<DeviceMemPtr> DynamicMemPoolBestFit::AllocContinuousTensorMem(size_t
   if (!device_addr) {
     return device_addr_list;
   }
+  std::lock_guard<std::mutex> locker(mutex_);
   // Remove the pre-alloc memory.
   auto mem_block = FindMemBlock(device_addr);
   MS_EXCEPTION_IF_NULL(mem_block);
@@ -189,9 +191,10 @@ DynamicMemBlockPtr DynamicMemPoolBestFit::FindMemBlock(const DeviceMemPtr &devic
 
 void DynamicMemPoolBestFit::FreeTensorMem(const DeviceMemPtr &device_addr) {
   MS_EXCEPTION_IF_NULL(device_addr);
+  std::lock_guard<std::mutex> locker(mutex_);
   auto mem_block = FindMemBlock(device_addr);
   if (mem_block == nullptr) {
-    // May be destory the memory pool first, then destory the address, so this is normal case.
+    // May be destroy the memory pool first, then destroy the address, so this is normal case.
     MS_LOG(DEBUG) << "Can't find the mem_block of the device address[" << device_addr << "].";
     return;
   }
@@ -263,6 +266,7 @@ void DynamicMemPoolBestFit::EraseIdleMemBuf(size_t size, const DeviceMemPtr &dev
 }
 
 void DynamicMemPoolBestFit::ReleaseDeviceRes() {
+  std::lock_guard<std::mutex> locker(mutex_);
   MS_LOG(INFO) << "The dynamic memory pool total size is " << total_mem_statistics_ << ", total used size is "
                << total_used_mem_statistics_ << ", used peak size is " << used_mem_peak_statistics_ << ".";
   for (auto iter = global_mem_block_list_.begin(); iter != global_mem_block_list_.end(); ++iter) {
@@ -276,6 +280,7 @@ void DynamicMemPoolBestFit::ReleaseDeviceRes() {
 }
 
 void DynamicMemPoolBestFit::DumpDynamicMemPoolInfo() {
+  std::lock_guard<std::mutex> locker(mutex_);
   MS_LOG(INFO) << "Start dump dynamic memory pool info.";
   DeviceAddrMapMemBuf mem_block_map;
   DynamicMemBufPtr mem_buf;

mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h

@@ -22,6 +22,8 @@
 #include <vector>
 #include <algorithm>
 #include <utility>
+#include <thread>
+#include <mutex>
 
 namespace mindspore {
 namespace device {
@@ -88,7 +90,10 @@ class DynamicMemPoolBestFit {
   // Display the information of memory block and memory buf.
   void DumpDynamicMemPoolInfo();
   // Get the map of global idle mem buf and size.
-  SizeMapMemBuf global_idle_mem_buf_map() const { return global_idle_mem_buf_map_; }
+  SizeMapMemBuf global_idle_mem_buf_map() {
+    std::lock_guard<std::mutex> locker(mutex_);
+    return global_idle_mem_buf_map_;
+  }
 
   // Get the related memory statistics information.
   size_t total_mem_statistics() const { return total_mem_statistics_; }
@@ -118,7 +123,7 @@ class DynamicMemPoolBestFit {
   bool IsDivide(size_t tensor_size, size_t mem_buf_size) const;
   // Divide the memory buf by alloc size.
   void DivideMemBuf(size_t size, const DynamicMemBufPtr &mem_buf);
-  // Find the memory block by deivce address.
+  // Find the memory block by device address.
   DynamicMemBlockPtr FindMemBlock(const DeviceMemPtr &device_addr);
   // The Comparator of memory block by device address, because memory blocks are arranged in order by device address.
   static bool CmpMemBlock(const DeviceMemPtr &device_addr, const DynamicMemBlockPtr &mem_block);
@@ -137,6 +142,9 @@ class DynamicMemPoolBestFit {
   size_t total_mem_statistics_{0};
   size_t total_used_mem_statistics_{0};
   size_t used_mem_peak_statistics_{0};
+
+  // Support multi-thread.
+  std::mutex mutex_;
 };
 }  // namespace device
 }  // namespace mindspore