caffe SyncedMemory 学习

#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_
#include <cstdlib>
#include "caffe/common.hpp"
#include "caffe/util/math_functions.hpp"
namespace caffe {
//分配内存空间的内联函数，大小为size字节
//typedef unsigned long long size_t
inline void CaffeMallocHost(void** ptr, size_t size) {
*ptr = malloc(size);
CHECK(*ptr) << "host allocation of size " << size << " failed";
}
//释放内存空间的内联函数
inline void CaffeFreeHost(void* ptr) {
free(ptr);
}
class SyncedMemory {
public:
SyncedMemory()
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
own_cpu_data_(false) {}
explicit SyncedMemory(size_t size)
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
own_cpu_data_(false) {}
~SyncedMemory();
const void* cpu_data();
void set_cpu_data(void* data);
const void* gpu_data();
void* mutable_cpu_data();
void* mutable_gpu_data();
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
SyncedHead head() { return head_; }
size_t size() { return size_; }

private:
void to_cpu();
void to_gpu();
void* cpu_ptr_;
void* gpu_ptr_;
size_t size_;
SyncedHead head_;
bool own_cpu_data_;

DISABLE_COPY_AND_ASSIGN(SyncedMemory);//使用宏定义，禁止复制和引用
/*
宏定义
#define DISABLE_COPY_AND_ASSIGN(classname) \
private:\
classname(const classname&);\
classname& operator=(const classname&)
写成代码的形式
private:
SyncedMemory(const SyncedMemory&);
SyncedMemory& operator=(const SyncedMemory&)
*/
};  // class SyncedMemory

}  // namespace caffe

#endif  // CAFFE_SYNCEDMEM_HPP_

#include <cstring>
#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

SyncedMemory::~SyncedMemory() {//析构函数，释放内存
if (cpu_ptr_ && own_cpu_data_) {
CaffeFreeHost(cpu_ptr_);
}

#ifndef CPU_ONLY
if (gpu_ptr_) {
CUDA_CHECK(cudaFree(gpu_ptr_));
}
#endif  // CPU_ONLY
}

inline void SyncedMemory::to_cpu() {
switch (head_) {
case UNINITIALIZED://未初始化
CaffeMallocHost(&cpu_ptr_, size_);
caffe_memset(size_, 0, cpu_ptr_);
head_ = HEAD_AT_CPU;
own_cpu_data_ = true;
break;
case HEAD_AT_GPU:
#ifndef CPU_ONLY
if (cpu_ptr_ == NULL) {
CaffeMallocHost(&cpu_ptr_, size_);
own_cpu_data_ = true;
}
caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);
/*
void caffe_gpu_memcpy(const size_t N, const void* X, void* Y) {
if (X != Y) {//内存数据的复制
CUDA_CHECK(cudaMemcpy(Y, X, N, cudaMemcpyDefault));  // NOLINT(caffe/alt_fn)
}
}
*/
head_ = SYNCED;//表示同步
#else
NO_GPU;
#endif
break;
case HEAD_AT_CPU:
case SYNCED:
break;
}
}

inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
switch (head_) {
case UNINITIALIZED:
CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
caffe_gpu_memset(size_, 0, gpu_ptr_);
head_ = HEAD_AT_GPU;
break;
case HEAD_AT_CPU:
if (gpu_ptr_ == NULL) {
CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
}
caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_);
head_ = SYNCED;
break;
case HEAD_AT_GPU:
case SYNCED:
break;
}
#else
NO_GPU;
#endif
}

const void* SyncedMemory::cpu_data() {//只读
to_cpu();
return (const void*)cpu_ptr_;
}

void SyncedMemory::set_cpu_data(void* data) {
CHECK(data);
if (own_cpu_data_) {
CaffeFreeHost(cpu_ptr_);
}
cpu_ptr_ = data;
head_ = HEAD_AT_CPU;
own_cpu_data_ = false;// 并非自己申请的数据内存，不需要自己释放。
}

const void* SyncedMemory::gpu_data() {//只读
#ifndef CPU_ONLY
to_gpu();
return (const void*)gpu_ptr_;
#else
NO_GPU;
#endif
}

void* SyncedMemory::mutable_cpu_data() {//可以修改
to_cpu();
head_ = HEAD_AT_CPU;
return cpu_ptr_;
}

void* SyncedMemory::mutable_gpu_data() {//可以修改的
#ifndef CPU_ONLY
to_gpu();
head_ = HEAD_AT_GPU;
return gpu_ptr_;
#else
NO_GPU;
#endif
}
}  // namespace caffe