Do not free nullptr and if a nullptr is freed, handle it correctly.

nouiz · nouiz · commit c7a3b3a9dac2 · 2021-07-27T07:35:24.000-07:00
diff --git a/tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc b/tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc
@@ -314,6 +314,8 @@ void* GpuCudaMallocAsyncAllocator::AllocateRaw(size_t alignment,
 }
 void GpuCudaMallocAsyncAllocator::DeallocateRaw(void* ptr) {
 #if TF_CUDA_MALLOC_ASYNC_SUPPORTED
+  if (ptr == nullptr)
+    return;
   if (auto result = cuMemFreeAsync(reinterpret_cast<const CUdeviceptr&>(ptr),
                                    cuda_stream_)) {
     if (result == CUDA_ERROR_DEINITIALIZED) {
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc
@@ -397,7 +397,8 @@ BaseGPUDevice::BaseGPUDevice(const SessionOptions& options, const string& name,
 
 BaseGPUDevice::~BaseGPUDevice() {
   delete gpu_device_info_;
-  gpu_allocator_->DeallocateRaw(scratch_);
+  if (scratch_)
+    gpu_allocator_->DeallocateRaw(scratch_);
   device_context_->Unref();
 }
 

Original file line number	Diff line number	Diff line change
`@@ -314,6 +314,8 @@ void* GpuCudaMallocAsyncAllocator::AllocateRaw(size_t alignment,`
`314`	`314`	`}`
`315`	`315`	`void GpuCudaMallocAsyncAllocator::DeallocateRaw(void* ptr) {`
`316`	`316`	`#if TF_CUDA_MALLOC_ASYNC_SUPPORTED`
	`317`	`+ if (ptr == nullptr)`
	`318`	`+ return;`
`317`	`319`	`if (auto result = cuMemFreeAsync(reinterpret_cast<const CUdeviceptr&>(ptr),`
`318`	`320`	`cuda_stream_)) {`
`319`	`321`	`if (result == CUDA_ERROR_DEINITIALIZED) {`
Original file line number	Diff line number	Diff line change
`@@ -397,7 +397,8 @@ BaseGPUDevice::BaseGPUDevice(const SessionOptions& options, const string& name,`
`397`	`397`
`398`	`398`	`BaseGPUDevice::~BaseGPUDevice() {`
`399`	`399`	`delete gpu_device_info_;`
`400`		`- gpu_allocator_->DeallocateRaw(scratch_);`
	`400`	`+ if (scratch_)`
	`401`	`+ gpu_allocator_->DeallocateRaw(scratch_);`
`401`	`402`	`device_context_->Unref();`
`402`	`403`	`}`
`403`	`404`