Caffe源码分析-syncedmem内存

2020-04-20 本文已影响0人陈泽_06aa

Caffe为了简化CPU主机和GPU内存管理，采用了一个syncedmem类来封装了这两部分的内存管理，其中它主要有以下功能：

支持CPU模式下的内存管理，基于malloc和free接口
支持GPU模式下的CPU内存和GPU内存管理，基于CUDA相关api接口
支持CPU与GPU之间的数据复制

在详细了解这些之前我们先介绍几个cuda内存管理函数，完成链接参考CUDA Memory Management

__host__cudaError_t cudaMallocHost ( void** ptr, size_t size )
分配一个host cpu内存。
__host__cudaError_t cudaFreeHost ( void* ptr )
释放由cudaMallocHost分配的Host主机内存
__host____device__cudaError_t cudaMalloc ( void** devPtr, size_t size )
在GPU设备上分配一个GPU设备内存
__host____device__cudaError_t cudaFree ( void* devPtr )
释放由cudaMalloc分配的GPU设备内存
__host__cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind )
用于Host和设备内存之间复制。

其中有一个很重要点就是，如果启用了CUDA，Host主机上的CPU内存是采用cudaMallocHost/cudaFreeHost进行申请和释放，而不是采用传统的malloc/free，在CUDA产品文档是这样描述：

The cuda driver tracks the virtual memory ranges allocated with cudaMallocHost and automatically accelerates calls to functions such as cudaMemcpy*(). Since the memory can be accessed directly by the device, it can be read or written with much higher bandwidth than pageable memory obtained with functions such as malloc().

因此就有了下面两个基础函数的实现

inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
  if (Caffe::mode() == Caffe::GPU) {
    CUDA_CHECK(cudaMallocHost(ptr, size));
    *use_cuda = true;
    return;
  }
#endif
  *ptr = malloc(size);
  *use_cuda = false;
}

inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
  if (use_cuda) {
    CUDA_CHECK(cudaFreeHost(ptr));
    return;
  }
#endif
  free(ptr);
}

有了以上基础知识以后，我们看看syncedmem如何实现。首先看看syncedmem的主类

class SyncedMemory {
 public:
  SyncedMemory();
  ~SyncedMemory();
  const void* cpu_data();
  void set_cpu_data(void* data);
  const void* gpu_data();
  void set_gpu_data(void* data);
  void* mutable_cpu_data();
  void* mutable_gpu_data();
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  SyncedHead head() const { return head_; }
  size_t size() const { return size_; }
 private:
  void to_cpu();
  void to_gpu();
  void* cpu_ptr_;
  void* gpu_ptr_;
  size_t size_;
  SyncedHead head_;
  bool own_cpu_data_;
  bool cpu_malloc_use_cuda_;
  bool own_gpu_data_;
  int device_;
};

cpu_ptr_和gpu_ptr_为host内存和GPU内存指针，支持通过cpu_data()/mutable_cpu_data()和gpu_data()/mutable_gpu_data()进行指针的获取。
to_cpu()/to_gpu()用于内存分配和内存数据转移。syncedmem内存是延迟分配，即我们需要指定位置的内存时候，才需要进行分配，其中size_指定内存大小。
SyncedHead当前内存状态，是UNINITIALIZED，还是仅仅在CPU端：HEAD_AT_CPU，还是仅仅在GPU端：HEAD_AT_GPU，还是GPU和CPU已经完成同步SYNCED。 set_cpu_data()/set_gpu_data()和mutable_cpu_data()/mutable_gpu_data()等操作都会将内存设置为非同步状态。
set_cpu_data()/set_gpu_data()支持设置一个外部指针，此时syncedmem对象对该内存没有own权限，即own_cpu_data_和own_gpu_data字段为false。在syncedmem析构时候，是不允许对该内存进行释放操作。但是如果将外部cpu指针进行to_gpu()，或者将外部gpu支持进行to_cpu()操作以后，会产生自己新的数据复制，此时对应own_cpu_data_和own_gpu_data字段会重新设置为true。

整体上就这些关键点，END

Caffe源码分析-syncedmem内存

猜你喜欢

热点阅读