forked from mindspore-Ecosystem/mindspore
!33205 [master][Bugfix][PyNative] Device memory leak
Merge pull request !33205 from caifubi/master-pynative-memory-leak
This commit is contained in:
commit
3e69715397
|
@ -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) {
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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()) {
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue