統一內存管理簡介

最近和一個朋友聊到了統一內存管理的話題，統一內存是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;
}