diff --git a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.cc b/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.cc index 604a0b6555..7a9388ca71 100644 --- a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.cc +++ b/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 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 DynamicMemPoolBestFit::AllocContinuousTensorMem(size_t if (!device_addr) { return device_addr_list; } + std::lock_guard 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 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 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 locker(mutex_); MS_LOG(INFO) << "Start dump dynamic memory pool info."; DeviceAddrMapMemBuf mem_block_map; DynamicMemBufPtr mem_buf; diff --git a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h b/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h index 644737cbd5..62dfc0aee8 100644 --- a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h +++ b/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h @@ -22,6 +22,8 @@ #include #include #include +#include +#include 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 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