Merge pull request #101 from borisfom/caffe-0.14

gpu_memory cleanup, fixes possible init issue

3rdparty/cub/cub/util_allocator.cuh

@@ -408,10 +408,10 @@ struct CachingDeviceAllocator
  && (block_itr->device == device)
  && (block_itr->bin == search_key.bin)) {
 
-  // use special rule for the last ("exact size") bin: set max memory overuse to 1/8th
+ // use special rule for the last ("exact size") bin: set max memory overuse to 1/8th
  if (search_key.bin == (unsigned int) -1 && (block_itr->bytes - search_key.bytes)*8UL > search_key.bytes)
  break;
- 
+
  cudaStream_t prev_stream = block_itr->associated_stream;
  if ((active_stream == prev_stream)
  || (cudaEventQuery(block_itr->ready_event) != cudaErrorNotReady)) {
@@ -454,22 +454,12 @@ struct CachingDeviceAllocator
  if (error != cudaSuccess) {
  if (debug) CubLog("\tdevice %d failed to allocate %lld bytes for stream %lld",
  device, (long long) search_key.bytes, (long long) search_key.associated_stream);
-
- // if (search_key.bytes < cached_bytes[device]) {
- // free all cached memory (for all devices), synchrionize and retry once
- cudaDeviceSynchronize();
- cudaThreadSynchronize();
- FreeAllCached();
- cudaDeviceSynchronize();
- cudaThreadSynchronize();
- error = cudaMalloc(&search_key.d_ptr, search_key.bytes);
- // }
  }
- if (CubDebug(error)) 
+ if (CubDebug(error))
  return error;
- if (CubDebug(error = cudaEventCreateWithFlags(&search_key.ready_event, cudaEventDisableTiming))) 
+ if (CubDebug(error = cudaEventCreateWithFlags(&search_key.ready_event, cudaEventDisableTiming)))
  return error;
- 
+
  // Insert into live blocks
  Lock(&spin_lock);
  live_blocks.insert(search_key);
@@ -548,30 +538,30 @@ struct CachingDeviceAllocator
 
  // Lock
  Lock(&spin_lock);
- 
+
  // Find corresponding block descriptor
  BusyBlocks::iterator block_itr = live_blocks.find(search_key);
  if (block_itr != live_blocks.end()) {
  // Remove from live blocks
  search_key = *block_itr;
- live_blocks.erase(block_itr); 
+ live_blocks.erase(block_itr);
  cached_bytes[device].busy -= search_key.bytes;
- 
+
  // Check if we should keep the returned allocation
  if (cached_bytes[device].free + search_key.bytes <= max_cached_bytes)
- { 
+ {
  // Insert returned allocation into free blocks
  cached_blocks.insert(search_key);
  cached_bytes[device].free += search_key.bytes;
  recached = true;
- if (debug) { 
+ if (debug) {
  CubLog("\tdevice %d returned %lld bytes from associated stream %lld.\n\t\t %lld available blocks cached (%lld bytes), %lld live blocks outstanding. (%lld bytes)\n",
- device, (long long) search_key.bytes, (long long) search_key.associated_stream, (long long) cached_blocks.size(), 
+ device, (long long) search_key.bytes, (long long) search_key.associated_stream, (long long) cached_blocks.size(),
  (long long) cached_bytes[device].free, (long long) live_blocks.size(), (long long) cached_bytes[device].busy);
  }
  }
  }
- 
+
  Unlock(&spin_lock);
 
  if (recached) {
@@ -585,7 +575,7 @@ struct CachingDeviceAllocator
  return error;
  if (CubDebug(error = cudaSetDevice(device))) return error;
  }
- 
+
  // Actually free device memory
  if (CubDebug(error = cudaFree(d_ptr))) return error;
  if (CubDebug(error = cudaEventDestroy(search_key.ready_event))) return error;

Makefile

@@ -571,7 +571,7 @@ $(DYNAMIC_NAME): $(OBJS)| $(LIB_BUILD_DIR)
  @ cd $(BUILD_DIR)/lib; rm -f $(DYNAMIC_SONAME_SHORT); ln -s $(DYNAMIC_VERSIONED_NAME_SHORT) $(DYNAMIC_SONAME_SHORT)
  @ cd $(BUILD_DIR)/lib; rm -f $(DYNAMIC_NAME_SHORT); ln -s $(DYNAMIC_SONAME_SHORT) $(DYNAMIC_NAME_SHORT)
 
-$(STATIC_NAME): $(OBJS) | $(LIB_BUILD_DIR) 
+$(STATIC_NAME): $(OBJS) | $(LIB_BUILD_DIR)
  @ echo AR -o $@
  $(Q)ar rcs $@ $(OBJS)
 

include/caffe/util/gpu_memory.hpp

@@ -77,6 +77,8 @@ class gpu_memory {
  cudaStream_t stream_;
  size_t size_;
  };
+ static void update_dev_info(int device);
+
 # endif
 
  private:
@@ -89,8 +91,12 @@ class gpu_memory {
 
 #ifndef CPU_ONLY
  struct MemInfo {
- size_t free;
- size_t total;
+ MemInfo() {
+ free = total = flush_count = 0;
+ }
+ size_t free;
+ size_t total;
+ unsigned flush_count;
  };
 
  static vector<MemInfo> dev_info_;
@@ -102,7 +108,6 @@ class gpu_memory {
  cudaStream_t stream = cudaStreamDefault);
  static void deallocate(pointer ptr, cudaStream_t = cudaStreamDefault);
 
- static void registerStream(cudaStream_t stream);
  static void getInfo(size_t *free_mem, size_t *used_mem);
 
  private:

src/caffe/util/gpu_memory.cpp

@@ -26,9 +26,12 @@ namespace caffe {
  return "No GPU: CPU Only Memory";
  }
 #else
+
  void gpu_memory::init(const std::vector<int>& gpus,
  PoolMode m, bool debug) {
- debug_ = debug;
+ bool debug_env = (getenv("DEBUG_GPU_MEM") != 0);
+ debug_ = debug || debug_env;
+
  if (gpus.size() <= 0) {
  // should we report an error here ?
  m = gpu_memory::NoPool;
@@ -63,7 +66,29 @@ namespace caffe {
  CHECK((ptr) != NULL);
  switch (mode_) {
  case CubPool:
- CUDA_CHECK(cubAlloc->DeviceAllocate(ptr, size, stream));
+ if (cubAlloc->DeviceAllocate(ptr, size, stream) != cudaSuccess) {
+ int cur_device;
+ CUDA_CHECK(cudaGetDevice(&cur_device));
+ // free all cached memory (for all devices), synchrionize
+ cudaDeviceSynchronize();
+ cudaThreadSynchronize();
+ cubAlloc->FreeAllCached();
+ cudaDeviceSynchronize();
+ cudaThreadSynchronize();
+
+ // Refresh per-device saved values.
+ for (int i = 0; i < dev_info_.size(); i++) {
+ // only query devices that were initialized
+ if (dev_info_[i].total) {
+ update_dev_info(i);
+ // record which device caused cache flush
+ if (i == cur_device)
+ dev_info_[i].flush_count++;
+ }
+ }
+ // retry once
+ CUDA_CHECK(cubAlloc->DeviceAllocate(ptr, size, stream));
+ }
  break;
  default:
  CUDA_CHECK(cudaMalloc(ptr, size));
@@ -85,47 +110,42 @@ namespace caffe {
  }
  }
 
- void gpu_memory::registerStream(cudaStream_t stream) {
- switch (mode_) {
- case CubPool:
- default:
- break;
- }
- }
 
- void gpu_memory::initMEM(const std::vector<int>& gpus, PoolMode m) {
- mode_ = m;
+ void gpu_memory::update_dev_info(int device) {
  int initial_device;
-
  CUDA_CHECK(cudaGetDevice(&initial_device));
 
- for (int i = 0; i < gpus.size(); i++) {
- int cur_device = gpus[i];
- if (cur_device+1 > dev_info_.size())
- dev_info_.resize(cur_device+1);
-
- CUDA_CHECK(cudaSetDevice(gpus[i]));
- cudaDeviceProp props;
- CUDA_CHECK(cudaGetDeviceProperties(&props, cur_device));
- CUDA_CHECK(cudaMemGetInfo(&dev_info_[cur_device].free,
- &dev_info_[cur_device].total));
-
- if (debug_) {
- std::cout << "cudaGetDeviceProperties: Mem = "
- << props.totalGlobalMem <<std::endl;
- std::cout << "cudaMemGetInfo_[" << cur_device
- <<": Free= " << dev_info_[cur_device].free
- << " Total= " << dev_info_[cur_device].total << std::endl;
- }
-
- // make sure we don't ask for more that total device memory
- dev_info_[i].free = std::min(dev_info_[cur_device].total,
- dev_info_[cur_device].free);
- dev_info_[i].free = std::min(props.totalGlobalMem,
- dev_info_[cur_device].free);
+ if (device+1 > dev_info_.size())
+ dev_info_.resize(device+1);
+
+ CUDA_CHECK(cudaSetDevice(device));
+ cudaDeviceProp props;
+ CUDA_CHECK(cudaGetDeviceProperties(&props, device));
+ CUDA_CHECK(cudaMemGetInfo(&dev_info_[device].free,
+ &dev_info_[device].total));
+
+ if (debug_) {
+ std::cout << "cudaGetDeviceProperties: Mem = "
+ << props.totalGlobalMem <<std::endl;
+ std::cout << "cudaMemGetInfo_[" << device
+ <<": Free= " << dev_info_[device].free
+ << " Total= " << dev_info_[device].total << std::endl;
  }
 
+ // make sure we don't have more that total device memory
+ dev_info_[device].total = std::min(props.totalGlobalMem,
+ dev_info_[device].total);
+
+ // here we are adding existing 'busy' allocations to CUDA free memory
+ dev_info_[device].free =
+ std::min(dev_info_[device].total,
+ dev_info_[device].free
+ + cubAlloc->cached_bytes[device].busy);
+ CUDA_CHECK(cudaSetDevice(initial_device));
+ }
 
+ void gpu_memory::initMEM(const std::vector<int>& gpus, PoolMode m) {
+ mode_ = m;
  switch ( mode_ ) {
  case CubPool:
  try {
@@ -141,12 +161,13 @@ namespace caffe {
  }
  catch (...) {}
  CHECK(cubAlloc);
+ for (int i = 0; i < gpus.size(); i++) {
+ update_dev_info(gpus[i]);
+ }
  break;
  default:
  break;
  }
-
- CUDA_CHECK(cudaSetDevice(initial_device));
  }
 
  const char* gpu_memory::getPoolName() {
@@ -167,7 +188,7 @@ namespace caffe {
  // Free memory is initial free memory minus outstanding allocations.
  // Assuming we only allocate via gpu_memory since its constructon.
  *free_mem = dev_info_[cur_device].free -
-  cubAlloc->cached_bytes[cur_device].busy;
+ cubAlloc->cached_bytes[cur_device].busy;
  break;
  default:
  CUDA_CHECK(cudaMemGetInfo(free_mem, total_mem));