Keson's blog

Caffe解读3 -- SyncedMemory

首先推荐一个阅读Caffe代码结构的网站
http://caffe.berkeleyvision.org/doxygen

简介

这个类主要用来做内存的分配和同步,代码量较少,包含syncedmem.hppsyncedmem.cpp,这个类相对比较简单易懂~ 但是也很重要,因为牵扯到了最底层的cpu和gpu数据的同步等问题。

首先简单介绍一下Pinned Memory和Non-Pinned Memory,详细的介绍可以参见我上一篇博客Caffe解读1-Pinned-Memory-Vs-Non-Pinned-Memory/

为了在CPU和GPU之间传输内存,关于CPU的内存分配和释放方式有以下两种:

  • 通过C标准库中的malloc函数完成内存分配,free进行内存释放
  • 调用CUDA中的cudaMallocHost函数完成内存分配,cudaFreeHost进行内存释放

mallocfree的优点是分配和释放的耗时少,缺点是CPU和GPU之间的传输相比而言比较慢。cudaMallocHostcudaFreeHost正好相反,优点是CPU和GPU之间的传输快,缺点是分配和释放内存比较耗时。

在Caffe中,让在GPU模式和CUDA可用的情况下,采用Pinned Memory方式,即使用cudaMallocHostcudaFreeHost来进行CPU内存的分配和释放。这在单个GPU的时候可能效果并不明显,但是更利于并行训练,更重要的是,在多个GPU上,这种方式更加稳定。

头文件syncedmem.hpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
#ifndef CAFFE_SYNCEDMEM_HPP_
#define CAFFE_SYNCEDMEM_HPP_
#include <cstdlib>
#include "caffe/common.hpp"
namespace caffe {
// If CUDA is available and in GPU mode, host memory will be allocated pinned,
// using cudaMallocHost. It avoids dynamic pinning for transfers (DMA).
// The improvement in performance seems negligible in the single GPU case,
// but might be more significant for parallel training. Most importantly,
// it improved stability for large models on many GPUs.
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;
CHECK(*ptr) << "host allocation of size " << size << " failed";
}
inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY
if (use_cuda) {
CUDA_CHECK(cudaFreeHost(ptr));
return;
}
#endif
free(ptr);
}
/**
* @brief Manages memory allocation and synchronization between the host (CPU)
* and device (GPU).
*
* TODO(dox): more thorough description.
*/
class SyncedMemory {
public:
SyncedMemory()
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
gpu_device_(-1) {}
explicit SyncedMemory(size_t size)
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
gpu_device_(-1) {}
~SyncedMemory();
const void* cpu_data(); //获取cpu数据指针
void set_cpu_data(void* data); //设置cpu数据
const void* gpu_data(); //获得gpu数据指针
void set_gpu_data(void* data); //设置gpu数据
void* mutable_cpu_data(); //获取可以更改cpu数据的指针
void* mutable_gpu_data(); //获取可以更改gpu数据的指针
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
SyncedHead head() { return head_; }
size_t size() { return size_; }
#ifndef CPU_ONLY
void async_gpu_push(const cudaStream_t& stream);
#endif
private:
void to_cpu(); //数据状态转移函数
void to_gpu(); //数据状态转移函数
void* cpu_ptr_; //cpu数据内存指针
void* gpu_ptr_; //gpu数据内存指针
size_t size_; //size
SyncedHead head_; //用来指明状态
bool own_cpu_data_; //共享标记,是否使用的是自己的cpu数据
bool cpu_malloc_use_cuda_;
bool own_gpu_data_; //共享标记,是否使用的是自己的gpu数据
int gpu_device_;
DISABLE_COPY_AND_ASSIGN(SyncedMemory);
}; // class SyncedMemory
} // namespace caffe
#endif // CAFFE_SYNCEDMEM_HPP_

CPU内存的分配和释放

进行了最简单的封装

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
#ifndef CPU_ONLY
if (Caffe::mode() == Caffe::GPU) { //GPU模式下使用cudaMallocHost
CUDA_CHECK(cudaMallocHost(ptr, size));
*use_cuda = true;
return;
}
#endif
*ptr = malloc(size); //CPU模式下使用malloc
*use_cuda = false;
CHECK(*ptr) << "host allocation of size " << size << " failed";
}
inline void CaffeFreeHost(void* ptr, bool use_cuda) {
#ifndef CPU_ONLY //GPU模式下使用cudaFreeHost
if (use_cuda) {
CUDA_CHECK(cudaFreeHost(ptr));
return;
}
#endif
free(ptr); //CPU模式下使用free
}

SyncedMemory类的私有成员变量

1
2
3
4
5
6
7
8
9
10
11
12
13
void* cpu_ptr_; //数据在cpu的指针
void* gpu_ptr_; //数据在gpu的指针
size_t size_; //数据的大小
/*
*用来表示数据的同步状态,有4种状态,分别是未初始化的,数据在cpu中,数据在gpu中,cpu和gpu都有。
* /
enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
SyncedHead head_; //用来指明状态
bool own_cpu_data_; //是否是自己的cpu数据
bool cpu_malloc_use_cuda_; //
bool own_gpu_data_; //是否有gpu数据
int gpu_device_; //gpu卡的标志

CPU/GPU 内存的状态

enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };

  • UNINITIALIZED
    这个状态是未初始化状态,也就是SyncedMemory最早的状态,这时候内存和显存都没有被分配,当cpu或者gpu申请内存时该状态终结。

  • HEAD_AT_CPU
    这个状态表明最近一次数据修改是由cpu引起的。此时cpu和gpu的数据还没有同步,也就是cpu和gpu的数据可能不同。

  • HEAD_AT_GPU
    这个状态表明最近一次数据修改是由gpu引起的。此时cpu和gpu的数据还没有同步,也就是cpu和gpu的数据可能不同。

  • SYNCED
    同步状态,这个状态表明此时cpu和gpu的数据一致。这个状态可以使得CPU和GPU数据相同时减少不必要的复制。

具体的状态转移和实现函数

构造函数和析构函数:

1
2
3
4
5
6
7
8
9
SyncedMemory()
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
gpu_device_(-1) {}
explicit SyncedMemory(size_t size)
: cpu_ptr_(NULL), gpu_ptr_(NULL), size_(size), head_(UNINITIALIZED),
own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false),
gpu_device_(-1) {}
~SyncedMemory();
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
SyncedMemory::~SyncedMemory() {
if (cpu_ptr_ && own_cpu_data_) {
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
#ifndef CPU_ONLY
if (gpu_ptr_ && own_gpu_data_) {
int initial_device;
cudaGetDevice(&initial_device);
if (gpu_device_ != -1) {
CUDA_CHECK(cudaSetDevice(gpu_device_));
}
CUDA_CHECK(cudaFree(gpu_ptr_));
cudaSetDevice(initial_device);
}
#endif // CPU_ONLY
}

两个内敛私有函数to_cpu()to_gpu()

to_cpu()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
inline void SyncedMemory::to_cpu() {
switch (head_) {
case UNINITIALIZED: //如果是未初始化状态,则进行cpu内存分配
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
caffe_memset(size_, 0, cpu_ptr_);
head_ = HEAD_AT_CPU; //设置状态为“数据在cpu”
own_cpu_data_ = true; //cpu拥有数据置为真
break;
case HEAD_AT_GPU: //如果是gpu拥有数据
#ifndef CPU_ONLY //如果有GPU存在
if (cpu_ptr_ == NULL) { //如果cpu还未分配内存空间,则进行分配
CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);
own_cpu_data_ = true; //cpu拥有数据置为真
}
caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_); //内存从gpu拷贝到cpu
head_ = SYNCED; //设置状态为"数据在cpu和gpu都拥有"
#else
NO_GPU; //否则log报错,这是宏定义
#endif
break;
case HEAD_AT_CPU: //CPU 已经拥有数据,则直接跳过
case SYNCED:
break;
}
}

to_gpu():

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
inline void SyncedMemory::to_gpu() {
#ifndef CPU_ONLY
switch (head_) {
case UNINITIALIZED: //如果状态是未初始化
CUDA_CHECK(cudaGetDevice(&gpu_device_));
CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_)); //分配GPU显存
caffe_gpu_memset(size_, 0, gpu_ptr_);
head_ = HEAD_AT_GPU; //设置状态为在GPU
own_gpu_data_ = true;
break;
case HEAD_AT_CPU: //如果状态是在CPU
if (gpu_ptr_ == NULL) {
CUDA_CHECK(cudaGetDevice(&gpu_device_));
CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_)); //分配GPU显存
own_gpu_data_ = true;
}
caffe_gpu_memcpy(size_, cpu_ptr_, gpu_ptr_); //将数据从CPU传输到GPU
head_ = SYNCED; //状态设置为CPU和GPU同步
break;
case HEAD_AT_GPU: //如果GPU已经拥有数据,则跳过
case SYNCED:
break;
}
#else
NO_GPU; //如果没有GPU,则log报错
#endif
}

获取cpu和gpu内存地址的方法

1
2
3
4
const void* cpu_data();
const void* gpu_data();
void* mutable_cpu_data();
void* mutable_gpu_data();

实现分别如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
const void* SyncedMemory::cpu_data() {
to_cpu(); //保证cpu内存存在数据
return (const void*)cpu_ptr_; //返回数据地址指针
}
const void* SyncedMemory::gpu_data() {
#ifndef CPU_ONLY //如果是GPU模式
to_gpu(); //保证GPU内存中存在数据
return (const void*)gpu_ptr_; //返回GPU中数据地址指针
#else
NO_GPU; //GPU不存在则log报错
return NULL;
#endif
}

mutable_cpu_data()mutable_gpu_data()

与上面两个函数类似,多了状态的设定和返回可以修改数据的指针

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
void* SyncedMemory::mutable_cpu_data() {
to_cpu();
head_ = HEAD_AT_CPU; //设置装填为数据在CPU
return cpu_ptr_;
}
void* SyncedMemory::mutable_gpu_data() {
#ifndef CPU_ONLY
to_gpu();
head_ = HEAD_AT_GPU; //设置状态为数据在GPU
return gpu_ptr_;
#else
NO_GPU;
return NULL;
#endif
}

set_cpu_data(void* data)set_gpu_data(void* data)

1
2
3
4
5
6
7
8
9
10
11
//cpu的data指针指向一个新的区域由data指针传入,并且将原来申请的内存释放
void SyncedMemory::set_cpu_data(void* data) {
CHECK(data);
if (own_cpu_data_) { //如果cpu内存有数据,则释放
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
cpu_ptr_ = data; //设置
head_ = HEAD_AT_CPU; //设置状态为数据在CPU
own_cpu_data_ = false; //表明当前使用的宿主的数据
}

我们可以看到set_cpu_data释放了当前的cpu内存,把指针指向data所指的内存中,own_cpu_data_设置为了false;表明当前使用的是宿主(data)的内存。

我们对own_cpu_data_进行标记是有必要的,因为当使用的是宿主的内存的时候,当这个类被释放而调用析构函数时,需要检查共享标记,不能释放宿主的内存。这样可以保证自己申请的内存只能由自己释放。

析构函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
SyncedMemory::~SyncedMemory() {
if (cpu_ptr_ && own_cpu_data_) {
CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
}
#ifndef CPU_ONLY
if (gpu_ptr_ && own_gpu_data_) {
int initial_device;
cudaGetDevice(&initial_device);
if (gpu_device_ != -1) {
CUDA_CHECK(cudaSetDevice(gpu_device_));
}
CUDA_CHECK(cudaFree(gpu_ptr_));
cudaSetDevice(initial_device);
}
#endif // CPU_ONLY
}

set_gpu_data(void *data)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
//
void SyncedMemory::set_gpu_data(void* data) {
#ifndef CPU_ONLY
CHECK(data);
if (own_gpu_data_) {
int initial_device;
cudaGetDevice(&initial_device);
if (gpu_device_ != -1) {
CUDA_CHECK(cudaSetDevice(gpu_device_));
}
CUDA_CHECK(cudaFree(gpu_ptr_));
cudaSetDevice(initial_device);
}
gpu_ptr_ = data;
head_ = HEAD_AT_GPU;
own_gpu_data_ = false;
#else
NO_GPU;
#endif
}

async_gpu_push(const cudaStream_t& stream)

最后还有一个异步传输的函数,cuda拷贝的异步传输,从数据从cpu拷贝到gpu,实现如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#ifndef CPU_ONLY
void SyncedMemory::async_gpu_push(const cudaStream_t& stream) {
CHECK(head_ == HEAD_AT_CPU);
if (gpu_ptr_ == NULL) {
CUDA_CHECK(cudaGetDevice(&gpu_device_));
CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
own_gpu_data_ = true;
}
const cudaMemcpyKind put = cudaMemcpyHostToDevice;
CUDA_CHECK(cudaMemcpyAsync(gpu_ptr_, cpu_ptr_, size_, put, stream));
// Assume caller will synchronize on the stream before use
head_ = SYNCED;
}
#endif