Merge pull request BVLC#179 from drnikolaev/caffe-0.15-multigpu-ws-fix

Hot fix preventing dead lock during shutdown
lukeyeager · Jun 25, 2016 · a7c9144 · a7c9144
2 parents dce8bbf + c068392
commit a7c9144
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/caffe/util/gpu_memory.cpp b/src/caffe/util/gpu_memory.cpp
@@ -37,7 +37,7 @@ void GPUMemory::Manager::init(const vector<int>& gpus, Mode m, bool debug) {
       // Just in case someone installed 'no cleanup' arena before
       delete cub_allocator_;
       cub_allocator_ = new cub::CachingDeviceAllocator(BIN_GROWTH, MIN_BIN,
-          MAX_BIN, MAX_CACHED_BYTES, false, debug_);
+          MAX_BIN, MAX_CACHED_BYTES, true, debug_);
     } catch (...) {
     }
     CHECK_NOTNULL(cub_allocator_);
@@ -117,7 +117,14 @@ void GPUMemory::Manager::deallocate(void* ptr, int device,
   }
   switch (mode_) {
   case CUB_ALLOCATOR:
-    CUDA_CHECK(cub_allocator_->DeviceFree(device, ptr));
+    {
+      int current_device;
+      cudaError_t status = cudaGetDevice(&current_device);
+      // Preventing dead lock while Caffe shutting down.
+      if (status != cudaErrorCudartUnloading) {
+        CUDA_CHECK(cub_allocator_->DeviceFree(device, ptr));
+      }
+    }
     break;
   default:
     CUDA_CHECK(cudaFree(ptr));