src/caffeine/syncedmem.cpp

   1 #include "cuda_runtime.h"
   2
   3 #include "caffeine/common.hpp"
   4 #include "caffeine/syncedmem.hpp"
   5
   6 namespace caffeine {
   7
   8 SyncedMemory::~SyncedMemory() {
   9   if (cpu_ptr_) {
  10     CUDA_CHECK(cudaFreeHost(cpu_ptr_));
  11   }
  12
  13   if (gpu_ptr_) {
  14     CUDA_CHECK(cudaFree(gpu_ptr_));
  15   }
  16 }
  17
  18 inline void SyncedMemory::to_cpu() {
  19   switch(head_) {
  20   case UNINITIALIZED:
  21     CUDA_CHECK(cudaMallocHost(&cpu_ptr_, size_));
  22     memset(cpu_ptr_, 0, size_);
  23     head_ = HEAD_AT_CPU;
  24     break;
  25   case HEAD_AT_GPU:
  26     if (cpu_ptr_ == NULL) {
  27       CUDA_CHECK(cudaMallocHost(&cpu_ptr_, size_));
  28       CUDA_CHECK(cudaMemcpy(cpu_ptr_, gpu_ptr_, size_, cudaMemcpyDeviceToHost));
  29     }
  30     head_ = SYNCED;
  31     break;
  32   case HEAD_AT_CPU:
  33   case SYNCED:
  34     break;
  35   }
  36 }
  37
  38 inline void SyncedMemory::to_gpu() {
  39   switch(head_) {
  40   case UNINITIALIZED:
  41     CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
  42     CUDA_CHECK(cudaMemset(gpu_ptr_, 0, size_));
  43     head_ = HEAD_AT_GPU;
  44     break;
  45   case HEAD_AT_CPU:
  46     if (gpu_ptr_ == NULL) {
  47       CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
  48       CUDA_CHECK(cudaMemcpy(gpu_ptr_, cpu_ptr_, size_, cudaMemcpyHostToDevice));
  49     }
  50     head_ = SYNCED;
  51     break;
  52   case HEAD_AT_GPU:
  53   case SYNCED:
  54     break;
  55   }
  56 }
  57
  58
  59 inline const void* SyncedMemory::cpu_data() {
  60   to_cpu();
  61   return (const void*)cpu_ptr_;
  62 }
  63
  64 inline const void* SyncedMemory::gpu_data() {
  65   to_gpu();
  66   return (const void*)gpu_ptr_;
  67 }
  68
  69 inline void* SyncedMemory::mutable_cpu_data() {
  70   to_cpu();
  71   head_ = HEAD_AT_CPU;
  72   return cpu_ptr_;
  73 }
  74
  75 inline void* SyncedMemory::mutable_gpu_data() {
  76   to_gpu();
  77   head_ = HEAD_AT_GPU;
  78   return gpu_ptr_;
  79 }
  80
  81
  82 }  // namespace caffeine
  83