diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.cc b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.cc index 6d11900a46..dec0122189 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.cc +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_manager.cc @@ -24,7 +24,6 @@ namespace ascend { constexpr uint64_t kAscendDeviceMemGB = 30; constexpr uint64_t kMemSizeGB = 30; constexpr uint64_t kAscendDeviceMemSize = (kAscendDeviceMemGB << kMemSizeGB); -constexpr uint64_t kReservedMemorySize = 10 * 1024 * 1024; void AscendMemoryManager::MallocDeviceMemory() { auto context_mem = GetDeviceMemSizeFromContext(); @@ -35,8 +34,9 @@ void AscendMemoryManager::MallocDeviceMemory() { MS_EXCEPTION(DeviceProcessError) << "rtMalloc mem size[" << device_mem_size_ << "] fail, ret[" << ret << "]"; } - dynamic_mem_offset_ = device_mem_size_ - kReservedMemorySize; + AscendMemoryPool::GetInstance().set_device_mem_size(device_mem_size_); AscendMemoryPool::GetInstance().set_device_mem_pool_base(device_mem_base_); + AscendMemoryPool::GetInstance().set_device_mem_pool_offset(device_mem_size_); AscendMemoryPool::GetInstance().set_graph_dynamic_mem_offset(dynamic_mem_offset_); } @@ -69,18 +69,11 @@ void AscendMemoryManager::FreeDeviceMemory() { } device_mem_base_ = nullptr; } - if (device_mem_pool_base_ != nullptr) { - auto ret = rtFree(device_mem_pool_base_); - if (ret != RT_ERROR_NONE) { - MS_LOG(ERROR) << "rtFree mem size[" << device_mem_pool_size_ << "] fail, ret[" << ret << "]"; - } - device_mem_pool_base_ = nullptr; - } } void AscendMemoryManager::ResetDynamicMemory() { total_dynamic_size_ = 0; - dynamic_mem_offset_ = device_mem_size_ - kReservedMemorySize; + dynamic_mem_offset_ = 0; AscendMemoryPool::GetInstance().set_graph_dynamic_mem_offset(dynamic_mem_offset_); } @@ -99,7 +92,7 @@ uint8_t *AscendMemoryManager::MallocStaticMem(size_t size, bool communication_me auto device_mem_pool_offset = AscendMemoryPool::GetInstance().device_mem_pool_offset(); MS_LOG(INFO) << "Malloc Memory: Static, total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_ - << "] memory pool[" << device_mem_pool_offset << "])" + << "] memory pool[" << device_mem_size_ - device_mem_pool_offset << "])" << " malloc [" << align_size << "] communication_mem: " << communication_mem; if (communication_mem) { @@ -123,15 +116,11 @@ uint8_t *AscendMemoryManager::MallocDynamicMem(size_t size, bool communication_m MS_LOG(INFO) << "Malloc Memory: Dynamic, total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_ << "] memory pool[" << device_mem_pool_offset << "])" << " malloc [" << align_size << "] communication_mem: " << communication_mem; - - if (dynamic_mem_offset_ < align_size) { - MS_LOG(EXCEPTION) << "Out of memory!!! total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_ - << "]) malloc [" << align_size << "] failed!"; - } - auto new_offset = dynamic_mem_offset_ - align_size; - if (new_offset <= device_mem_pool_offset) { + auto offset = dynamic_mem_offset_; + auto new_offset = dynamic_mem_offset_ + align_size; + if (new_offset >= device_mem_pool_offset) { MS_LOG(EXCEPTION) << "Out of memory!!! total[" << device_mem_size_ << "] (dynamic[" << total_dynamic_size_ - << "] memory pool[" << device_mem_pool_offset << "])" + << "] memory pool[" << device_mem_size_ - device_mem_pool_offset << "])" << " malloc [" << align_size << "] failed!"; } total_dynamic_size_ += align_size; @@ -139,9 +128,9 @@ uint8_t *AscendMemoryManager::MallocDynamicMem(size_t size, bool communication_m AscendMemoryPool::GetInstance().set_graph_dynamic_mem_offset(dynamic_mem_offset_); if (communication_mem) { // create protect area [kMemAlignSize -- data -- kMemAlignSize] - return device_mem_base_ + dynamic_mem_offset_ + kMemAlignSize; + return device_mem_base_ + offset + kMemAlignSize; } else { - return device_mem_base_ + dynamic_mem_offset_; + return device_mem_base_ + offset; } } } // namespace ascend diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc index 4715a72a20..6d711dacd9 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc @@ -25,13 +25,13 @@ size_t AscendMemoryPool::AllocDeviceMem(size_t size, DeviceMemPtr *addr) { if (size == 0) { MS_LOG(EXCEPTION) << "Failed to alloc memory pool resource, the size is zero!"; } - if (device_mem_pool_offset_ + size >= graph_dynamic_mem_offset_) { + if (device_mem_pool_offset_ - size <= graph_dynamic_mem_offset_) { MS_LOG(EXCEPTION) << "Failed to alloc memory pool memory, the current device_mem_pool_offset_ [" << device_mem_pool_offset_ << "], current graph_dynamic_mem_offset_ " << graph_dynamic_mem_offset_ << "], need memory size [" << size << "]"; } + device_mem_pool_offset_ -= size; *addr = device_mem_pool_base_ + device_mem_pool_offset_; - device_mem_pool_offset_ += size; if (*addr == nullptr) { MS_LOG(EXCEPTION) << "Alloc device memory pool address is nullptr, failed to alloc memory pool resource!"; } @@ -57,6 +57,12 @@ void AscendMemoryPool::set_device_mem_pool_base(uint8_t *device_mem_pool_base) { device_mem_pool_base_ = device_mem_pool_base; } +void AscendMemoryPool::set_device_mem_size(uint64_t device_mem_size) { device_mem_size_ = device_mem_size; } + +void AscendMemoryPool::set_device_mem_pool_offset(uint64_t device_mem_pool_offset) { + device_mem_pool_offset_ = device_mem_pool_offset; +} + void AscendMemoryPool::set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_offset) { graph_dynamic_mem_offset_ = graph_dynamic_mem_offset; } @@ -64,14 +70,14 @@ void AscendMemoryPool::set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_o uint64_t AscendMemoryPool::device_mem_pool_offset() const { return device_mem_pool_offset_; } size_t AscendMemoryPool::free_mem_size() { - if (graph_dynamic_mem_offset_ <= device_mem_pool_offset_) { + if (graph_dynamic_mem_offset_ >= device_mem_pool_offset_) { MS_LOG(EXCEPTION) << "graph dynamic mem offset [" << graph_dynamic_mem_offset_ << "] less than or equal to device mem pool offset [" << device_mem_pool_offset_ << "]!"; } - return graph_dynamic_mem_offset_ - device_mem_pool_offset_; + return device_mem_pool_offset_ - graph_dynamic_mem_offset_; } -size_t AscendMemoryPool::total_mem_size() { return graph_dynamic_mem_offset_ == 0 ? 0 : graph_dynamic_mem_offset_ - 1; } +size_t AscendMemoryPool::total_mem_size() { return device_mem_size_ - graph_dynamic_mem_offset_; } } // namespace ascend } // namespace device } // namespace mindspore diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h index 439468fd6f..46b7fd4392 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h @@ -31,7 +31,9 @@ class AscendMemoryPool : public DynamicMemPoolBestFit { size_t AllocDeviceMem(size_t size, DeviceMemPtr *addr) override; bool FreeDeviceMem(const DeviceMemPtr &addr) override; + void set_device_mem_size(uint64_t device_mem_size); void set_device_mem_pool_base(uint8_t *device_mem_pool_base); + void set_device_mem_pool_offset(uint64_t device_mem_pool_offset); void set_graph_dynamic_mem_offset(uint64_t graph_dynamic_mem_offset); uint64_t device_mem_pool_offset() const; @@ -52,6 +54,7 @@ class AscendMemoryPool : public DynamicMemPoolBestFit { private: AscendMemoryPool() = default; uint8_t *device_mem_pool_base_{nullptr}; + uint64_t device_mem_size_{0}; uint64_t device_mem_pool_offset_{0}; uint64_t graph_dynamic_mem_offset_{0}; };