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 72765f4ecf..9ad49fbab6 100644 --- a/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h +++ b/mindspore/ccsrc/backend/optimizer/mem_reuse/mem_dynamic_allocator.h @@ -87,6 +87,7 @@ class DynamicMemPoolBestFit { void ReleaseDeviceRes(); // Display the information of memory block and memory buf. void DumpDynamicMemPoolInfo(); + SizeMapMemBuf GetIdleMemBufMap() { return global_idle_mem_buf_map_; } // Get the related memory statistics information. size_t total_mem_statistics() const { return total_mem_statistics_; } diff --git a/mindspore/ccsrc/backend/session/ascend_session.h b/mindspore/ccsrc/backend/session/ascend_session.h index a42377bbaa..e8b624163c 100755 --- a/mindspore/ccsrc/backend/session/ascend_session.h +++ b/mindspore/ccsrc/backend/session/ascend_session.h @@ -29,6 +29,7 @@ #include "backend/kernel_compiler/kernel.h" #include "backend/session/session_factory.h" #include "backend/session/ascend_control_parser.h" +#include "runtime/device/ascend/ascend_memory_pool.h" namespace mindspore { namespace session { @@ -37,7 +38,7 @@ enum GraphType : int { COMMON_GRAPH = 0, CONDITION_GRAPH = 1, BRANCH_START = 2, class AscendSession : public SessionBasic { public: AscendSession() { final_graph_id_ = kInvalidGraphId; } - ~AscendSession() override = default; + ~AscendSession() override { mindspore::device::ascend::AscendMemoryPool::GetInstance().ResetIdleMemBuf(); } void Init(uint32_t device_id) override { SessionBasic::Init(device_id); context_ = std::make_shared(kAscendDevice, device_id); diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc index 6d711dacd9..d70f120370 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.cc @@ -43,6 +43,13 @@ bool AscendMemoryPool::FreeDeviceMem(const DeviceMemPtr &addr) { return true; } +void AscendMemoryPool::ResetIdleMemBuf() { + auto idle_mem_buf_map = DynamicMemPoolBestFit::GetIdleMemBufMap(); + for (auto &it : idle_mem_buf_map) { + rtMemset(it.second->device_addr_, it.first, 0, it.first); + } +} + size_t AscendMemoryPool::AlignMemorySize(size_t size) const { if (size == 0) { MS_LOG(EXCEPTION) << "The align memory size is a zero !"; diff --git a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h index 46b7fd4392..c9007d3f5e 100644 --- a/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h +++ b/mindspore/ccsrc/runtime/device/ascend/ascend_memory_pool.h @@ -31,6 +31,7 @@ class AscendMemoryPool : public DynamicMemPoolBestFit { size_t AllocDeviceMem(size_t size, DeviceMemPtr *addr) override; bool FreeDeviceMem(const DeviceMemPtr &addr) override; + void ResetIdleMemBuf(); 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);