[cuda][caffe]统一内存管理

统一内存管理简介

最近和一个朋友聊到了统一内存管理的话题，统一内存是cuda中的一个很重要的概念，通过统一内存管理，用户可以直接使用内存，而不用在意数据在内存中位置，做到透明管理。

统—内存编程模型由CUDA6引入，从开普勒架构开始就可用，但开普勒架构和麦克斯韦架构的GPU提供的统一内存编程功能相对较弱。从帕斯卡架构到现在的伏特架构和图灵架构统一内存的功能加强了很多,主要是因为这些架构的GPU具有了精细的缺页异常处理（page faulthandling）能力。

统一内存带来的优势

1. 代码更简洁，编程更简单。在没有统一内存管理之前，需要开辟两个指针，分别是内存指针和设备指针，并分别开辟空间，手动管理两者之间的内存同步，会写更多代码。
2. 自动将数据搬运到合适的位置。在某些训练或者推理场景下，需要将显存offload或者将内存上的数据搬运到设备端，通过统一内存管理，有望实现自动数据搬运。
3. 统一内存管理可以超额分配显存。当显存大小不够时，仍然可以分配统一内存，超出的部分分配到内存上。

统一内存可能存在的问题

1. 不熟悉的新手可能会错误使用。内存中数据搬运是自动的，在某些情况下，新手可能会误用其特性导致内存效率降低。
2. 内存抖动问题，当频繁在在不同侧的内存进行写的时候存在内存抖动，数据会来回移动。

统一内存管理编程实现

动态申请

在该程序中,我们使用cudaMallocManaged函数对x,y,z进行内存分配,并在Host侧赋予了每一个位置一个初始值,接着将所有值放入到device端进行运算,并将运算完成的结果在Host端打印。
由于动态申请，在内存侧的数据是放在堆上的，而在设备侧是动态申请的。
注意统一内存仅可在主机侧代码中申请。

__global__ void addKernel(float *x, float *y, float *z, const int N) {
  int tid = threadIdx.x + blockDim.x * blockIdx.x;
 
  if (tid < N) {
    z[tid] = x[tid] + y[tid];
  }
}
 
int main() {
  float *x, *y, *z;
  const int N = 10000;
  int size = sizeof(float) * 10000;
  cudaMallocManaged((void **)&x, size);
  cudaMallocManaged((void **)&y, size);
  cudaMallocManaged((void **)&z, size);
 
  for (int i = 0; i < N; i++) {
    x[i] = 0.1 * i;
    y[i] = 0.2 * i;
    z[i] = 0.f;
  }
 
  const int block_size = 128;
  const int grid_s1ze = N / block_size;
  addKernel<<<grid_s1ze, block_size>>>(x, y, z, N);
 
  cudaDeviceSynchronize();
 
  for (int i = 0; i < N; i++) {
    printf("%f ", z[i]);
  }
  printf("\n");
 
  cudaFree(x);
  cudaFree(y);
  cudaFree(z);
 
  return 0;
}

静态申请

相比于动态申请，静态申请则是直接在设备上申请一块内存空间。
需要注意的是，这块内存空间在源文件所有可视范围内均可使用。

__device__ __managed__ int ret[1000];
__global__ void plusKernel(int a, int b) {
  int tid = threadIdx.x + blockDim.x * blockIdx.x;
  ret[tid] = a + b + threadIdx.x;
}
 
int main_plus() {
  int a = 10, b = 100;
  plusKernel<<<1, 1000>>>(a, b);
  cudaDeviceSynchronize();
 
  for (int i = 0; i < 1000; i++) {
    printf("%d ", ret[i]);
  }
 
  printf("\n");
 
  return 0;
}

超额内存分配

采用命令nvcc --compiler-bindir /usr/bin/g++-10 -DUNIFIED UMM_malloc.cu进行编译，在使用统一内存时，可以将内存开辟到60 GB，当使用普通的内存分配时，仅可分配6GB。

const int N = 60;
 
// #define UNIFIED
 
int main_malloc() {
 
  uint64_t *x;
  for (int n = 1; n <= N; n++) {
    size_t sz = size_t(n) * 1024 * 1024 * 1024;
 
#ifdef UNIFIED
    CHECK(cudaMallocManaged(&x, sz));
    CHECK(cudaFree(x));
    printf("cudaMallocManaged %d GB data\n", n);
#else
    CHECK(cudaMalloc(&x, sz));
    CHECK(cudaFree(x));
    printf("cudaMalloc %d GB data\n", n);
 
#endif
  }
  return 0;
}

GPU和CPU初始化

当分别使用GPU和CPU分别对内存中的参数进行初始化时，理论上当使用GPU初始化时能够使用GPU显存+CPU内存，当使用CPU初始化时，仅可使用CPU主存。

__global__ void gpu_torch(uint64_t *x, size_t SZ) {
  size_t tid = threadIdx.x + blockDim.x + blockIdx.x;
  if (tid < SZ) {
    x[tid] = 0;
  }
}
 
void cpu_touch(uint64_t *x, size_t SZ) {
  for (int i = 0; i < SZ / sizeof(uint64_t); i++) {
    x[i] = 0;
  }
}
int main_touch() {
  uint64_t *x;
  for (int n = 1; n <= N; n++) {
    size_t sz = size_t(n) * 1024 * 1024 * 1024;
 
    CHECK(cudaMallocManaged(&x, sz));
 
    size_t block_size = 1024;
    size_t grid_size = sz / sizeof(uint64_t) / 1024;
    // gpu_torch<<<grid_size, block_size>>>(x, sz);
    cpu_touch(x, sz);
    // CHECK(cudaGetLastError());
    // CHECK(cudaDeviceSynchronize());
    CHECK(cudaFree(x));
    printf("cudaMallocManaged %d GB data and init it!\n", n);
  }
}

统一内存管理实现

由于cuda并未开源统一内存的具体实现，但是我们在caffe源码发现了类似的实现。
在caffe中基础类为blob，而blob的内存管理是SyncedMemory，SyncedMemory便是通过封装cpu、gpu实现，屏蔽内存管理以及数据同步细节，惰性内存分配和同步，提高效率和节省内存。

数据结构

先来看一看数据结构，其中包含了CPU指针、GPU指针和数据大小，同步头以确定当前最新数据的位置，通过cpu、gpu所有权标志位来设置外部数据，通过是否使用标志位来标记是否使用pinned memory。

class SyncedMemory {
 public:
  SyncedMemory();
  explicit SyncedMemory(size_t size);
  ~SyncedMemory();
  const void* cpu_data();  // 同步gpu到cpu，只读数据
  void set_cpu_data(void* data); // 从外部设置cpu数据，不拥有所有权
  const void* gpu_data();  // 同步cpu数据到gpu，只读数据
  void set_gpu_data(void* data); // 从外部设置gpu数据，不拥有所有权
  void* mutable_cpu_data();  // 同步gpu到cpu，读写数据
  void* mutable_gpu_data();  // 同步cpu到gpu，读写数据
  enum SyncedHead { UNINITIALIZED, HEAD_AT_CPU, HEAD_AT_GPU, SYNCED };
  SyncedHead head() const { return head_; }  // 获取状态头
  size_t size() const { return size_; } // 获取数据大小
 
 private:
  void check_device();
 
  void to_cpu();
  void to_gpu();
  void* cpu_ptr_;  // cpu数据指针
  void* gpu_ptr_;  // gpu数据指针
  size_t size_;    // 数据大小
  SyncedHead head_;  // 数据的同步头
  bool own_cpu_data_;  // 是否具有cpu所有权
  bool cpu_malloc_use_cuda_; // Pinned memory, 否则使用的普通的malloc, 会慢一些
  bool own_gpu_data_; // 是否具有gpu所有权
  int device_;  // device id
};  // class SyncedMemory

函数及实现

分为无参构造函数、有参构造函数、析构函数、数据转移函数、设置外部数据几个函数。

构造函数：构造函数是初始化数据，这里head_被初始化为UNINITIALIZED，其他默认值均为NULL或者False，有参和无参函数的区别在于size_会不会被设置。
析构函数：只有在cpu侧有数据且拥有所有权时才会释放cpu内存，只有在gpu侧有数据且拥有所有权时才会释放数据。
to_cpu函数：将数据同步搬运至cpu上，如果是非同步状态，则需要开辟空间，并设置状态头在CPU侧；如果在GPU侧，则需要将数据GPU数据转移到CPU，设置状态为已同步状态；在CPU侧不做任何处理。
to_gpu函数：与to_cpu 类似，只不过是当数据是在gpu上而不是cpu侧。
cpu_data函数：会调用to_cpu函数将数据从gpu侧搬运到cpu侧，返回指针，这里是惰性搬运，也就是只有当数据需要使用时才搬运；需要mutable_cpu_data只比只读多了一个标志，声明状态头在CPU侧。
set_cpu_data函数：使用外部的数据设置cpu数据，发生的是浅拷贝，但由于是外部数据，所以不拥有所有权，设置所有权为false，但需要注意的是，再搬运到gpu侧时，gpu是有可能拥有权限的。
gpu函数与以上类似。

SyncedMemory::SyncedMemory()
  : cpu_ptr_(NULL), gpu_ptr_(NULL), size_(0), head_(UNINITIALIZED),
    own_cpu_data_(false), cpu_malloc_use_cuda_(false), own_gpu_data_(false) {
#ifndef CPU_ONLY
#ifdef DEBUG
  CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}
 
SyncedMemory::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) {
#ifndef CPU_ONLY
#ifdef DEBUG
  CUDA_CHECK(cudaGetDevice(&device_));
#endif
#endif
}
 
SyncedMemory::~SyncedMemory() {
  check_device();
  if (cpu_ptr_ && own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }
 
#ifndef CPU_ONLY
  if (gpu_ptr_ && own_gpu_data_) {
    CUDA_CHECK(cudaFree(gpu_ptr_));
  }
#endif  // CPU_ONLY
}
 
inline void SyncedMemory::to_cpu() {
  check_device();
  switch (head_) {
  case UNINITIALIZED:  // 非初始化状态
    CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);  // 分配cpu侧内存，可以为pinned memory
    caffe_memset(size_, 0, cpu_ptr_);
    head_ = HEAD_AT_CPU;   // 设置状态头在cpu侧
    own_cpu_data_ = true;  // 拥有所有权
    break;
  case HEAD_AT_GPU:  // 数据在gpu侧
#ifndef CPU_ONLY
    if (cpu_ptr_ == NULL) {
      CaffeMallocHost(&cpu_ptr_, size_, &cpu_malloc_use_cuda_);  // 先在cpu
      own_cpu_data_ = true;
    }
    caffe_gpu_memcpy(size_, gpu_ptr_, cpu_ptr_);
    head_ = SYNCED;
#else
    NO_GPU;
#endif
    break;
  case HEAD_AT_CPU:
  case SYNCED:
    break;
  }
}
 
inline void SyncedMemory::to_gpu() {
  check_device();
#ifndef CPU_ONLY
  switch (head_) {
  case UNINITIALIZED:
    CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
    caffe_gpu_memset(size_, 0, gpu_ptr_);
    head_ = HEAD_AT_GPU;
    own_gpu_data_ = true;
    break;
  case HEAD_AT_CPU:
    if (gpu_ptr_ == NULL) {
      CUDA_CHECK(cudaMalloc(&gpu_ptr_, size_));
      own_gpu_data_ = true;
    }
    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() {
  check_device();
  to_cpu();
  return (const void*)cpu_ptr_;
}
 
void SyncedMemory::set_cpu_data(void* data) {
  check_device();
  CHECK(data);
  if (own_cpu_data_) {
    CaffeFreeHost(cpu_ptr_, cpu_malloc_use_cuda_);
  }
  cpu_ptr_ = data;
  head_ = HEAD_AT_CPU;
  own_cpu_data_ = false;
}
 
const void* SyncedMemory::gpu_data() {
  check_device();
#ifndef CPU_ONLY
  to_gpu();
  return (const void*)gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}
 
void SyncedMemory::set_gpu_data(void* data) {
  check_device();
#ifndef CPU_ONLY
  CHECK(data);
  if (own_gpu_data_) {
    CUDA_CHECK(cudaFree(gpu_ptr_));
  }
  gpu_ptr_ = data;
  head_ = HEAD_AT_GPU;
  own_gpu_data_ = false;
#else
  NO_GPU;
#endif
}
 
void* SyncedMemory::mutable_cpu_data() {
  check_device();
  to_cpu();
  head_ = HEAD_AT_CPU;
  return cpu_ptr_;
}
 
void* SyncedMemory::mutable_gpu_data() {
  check_device();
#ifndef CPU_ONLY
  to_gpu();
  head_ = HEAD_AT_GPU;
  return gpu_ptr_;
#else
  NO_GPU;
  return NULL;
#endif
}

场景分析

统一内存访问的惰性管理会减少内存搬运（需要时才搬运），提高程序效率。

1. 只读场景：在初始状态下，访问to_cpu或to_gpu函数会将数据从一侧搬运至另外一侧，最多发生一次内存搬运，最终达到synced状态，此时无论怎么读都是最新数据。
2. 设置外部数据只读场景：set_cpu_data函数从外部拿到的数据需要外部进行指针管理，没有所有权；但是当数据移动到另外一侧时，发生了内存搬运，便有了所有权，存在一侧有所有权一侧无所有权的情况，此时如果外部修改指针内存，可能会存在cpu数据和gpu数据不一致的问题。
3. 读写场景：mutable_gpu_data相比于gpu_data多了一个head的操作，会将状态头设置为对应侧，如果是不同设备的读写则和场景2类似；如果数据在不同设备是只写的，head将不会出现SYNCED状态，那么频繁地从cpu侧移动到gpu侧或相反，这就造成了内存抖动，数据会频繁在两种设备间来回移动。

统一内存管理程序的优化

上述我们谈到了内存抖动的问题，即频繁在两侧数据进行内存搬运，一个解决方法是cudaMemPrefetchAsync，该函数的作用是在CUDA流stream中将统一内存缓冲区devPtr内的count字节的内存迁移到设备dstDevice(主机的设备号用cudaCpuDeviceID表示)中的内存区域，从而方式(或减少)缺页异常，并提高数据的局部性。

int main_add() {
  float *x, *y, *z;
  const int N = 10000;
  int size = sizeof(float) * 10000;
  cudaMallocManaged((void **)&x, size);
  cudaMallocManaged((void **)&y, size);
  cudaMallocManaged((void **)&z, size);
 
  for (int i = 0; i < N; i++) {
    x[i] = 0.1 * i;
    y[i] = 0.2 * i;
    z[i] = 0.f;
  }
 
  const int block_size = 128;
  const int grid_s1ze = N / block_size;
 
 
  int device_id=0;
  cudaGetDevice(&device_id);
 
  cudaMemPrefetchAsync(x,size,device_id,NULL);
  cudaMemPrefetchAsync(y,size,device_id,NULL);
  cudaMemPrefetchAsync(z,size,device_id,NULL);
 
  addKernel<<<grid_s1ze, block_size>>>(x, y, z, N);
 
  cudaDeviceSynchronize();
 
  for (int i = 0; i < N; i++) {
    printf("%f ", z[i]);
  }
  printf("\n");
 
  cudaFree(x);
  cudaFree(y);
  cudaFree(z);
 
  return 0;
}