!33205 [master][Bugfix][PyNative] Device memory leak

Merge pull request !33205 from caifubi/master-pynative-memory-leak
This commit is contained in:
i-robot 2022-04-29 06:19:13 +00:00 committed by Gitee
commit 3e69715397
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
4 changed files with 25 additions and 3 deletions

View File

@ -543,6 +543,16 @@ void AscendDeviceContext::LoadModel(const NotNull<KernelGraphPtr> &root_graph) c
bool AscendDeviceContext::AllocateMemory(DeviceAddress *const &address, size_t size) const {
MS_EXCEPTION_IF_NULL(address);
MS_EXCEPTION_IF_NULL(runtime_instance_);
if (address->DeviceType() != DeviceAddressType::kAscend) {
MS_LOG(ERROR) << "The device address type is wrong: " << address->DeviceType();
return false;
}
if (address->GetPtr() != nullptr) {
MS_LOG(ERROR) << "Memory leak detected!";
return false;
}
runtime_instance_->SetContext();
auto device_ptr = mem_manager_->MallocMemFromMemPool(size, address->from_persistent_mem_);
if (!device_ptr) {

View File

@ -84,7 +84,13 @@ bool CPUDeviceContext::AllocateMemory(DeviceAddress *const &address, size_t size
MS_EXCEPTION_IF_NULL(address);
MS_EXCEPTION_IF_NULL(mem_manager_);
if (address->DeviceType() != DeviceAddressType::kCPU) {
MS_LOG(EXCEPTION) << "The device address type is wrong: " << address->DeviceType();
MS_LOG(ERROR) << "The device address type is wrong: " << address->DeviceType();
return false;
}
if (address->GetPtr() != nullptr) {
MS_LOG(ERROR) << "Memory leak detected!";
return false;
}
auto device_ptr = mem_manager_->MallocMemFromMemPool(size, 0);

View File

@ -186,7 +186,13 @@ void GPUDeviceContext::Destroy() {
bool GPUDeviceContext::AllocateMemory(DeviceAddress *const &address, size_t size) const {
MS_EXCEPTION_IF_NULL(address);
if (address->DeviceType() != DeviceAddressType::kGPU) {
MS_LOG(EXCEPTION) << "The device address type is wrong: " << address->DeviceType();
MS_LOG(ERROR) << "The device address type is wrong: " << address->DeviceType();
return false;
}
if (address->GetPtr() != nullptr) {
MS_LOG(ERROR) << "Memory leak detected!";
return false;
}
if (!BindDeviceToCurrentThread()) {

View File

@ -256,7 +256,7 @@ void OutputActor::UpdateOutputDeviceAddress() {
tensor_device_address->set_from_persistent_mem(device_tensor->from_persistent_mem());
tensor_device_address->set_host_shape(device_tensor->host_shape());
// The outputs may have the same output node, so need skip when the node has been done.
if (device_tensor->GetPtr() == nullptr) {
if (tensor_device_address->GetPtr() != nullptr) {
continue;
}