malloc.cpp

Go to the documentation of this file.

#include <cstdlib>
#include <cstdio>
#include <string>
#include <map>
#include <unistd.h> // for getpagesize()
#include <execinfo.h> // for backtrace
#include <quda_internal.h>

#ifdef USE_QDPJIT
#include "qdp_quda.h"
#endif

namespace quda {

  enum AllocType {
    DEVICE,
    HOST,
    PINNED,
    MAPPED,
    N_ALLOC_TYPE
  };

  class MemAlloc {

  public:
    std::string func;
    std::string file;
    int line;
    size_t size;
    size_t base_size;

    MemAlloc()
      : line(-1), size(0), base_size(0) { }

    MemAlloc(std::string func, std::string file, int line)
      : func(func), file(file), line(line), size(0), base_size(0) { }

    MemAlloc& operator=(const MemAlloc &a) {
      if (&a != this) {
  func = a.func;
  file = a.file;
  line = a.line;
  size = a.size;
  base_size = a.base_size;
      }
      return *this;
    }
  };


  static std::map<void *, MemAlloc> alloc[N_ALLOC_TYPE];
  static long total_bytes[N_ALLOC_TYPE] = {0};
  static long max_total_bytes[N_ALLOC_TYPE] = {0};
  static long total_host_bytes, max_total_host_bytes;
  static long total_pinned_bytes, max_total_pinned_bytes;

  long device_allocated_peak() { return max_total_bytes[DEVICE]; }

  long pinned_allocated_peak() { return max_total_bytes[PINNED]; }

  long mapped_allocated_peak() { return max_total_bytes[MAPPED]; }

  long host_allocated_peak() { return max_total_bytes[HOST]; }

  static void print_trace (void) {
    void *array[10];
    size_t size;
    char **strings;
    size = backtrace (array, 10);
    strings = backtrace_symbols (array, size);
    printfQuda("Obtained %zd stack frames.\n", size);
    for (size_t i=0; i<size; i++) printfQuda("%s\n", strings[i]);
    free(strings);
  }

  static void print_alloc_header()
  {
    printfQuda("Type    Pointer          Size             Location\n");
    printfQuda("----------------------------------------------------------\n");
  }


  static void print_alloc(AllocType type)
  {
    const char *type_str[] = {"Device", "Host  ", "Pinned", "Mapped"};
    std::map<void *, MemAlloc>::iterator entry;

    for (entry = alloc[type].begin(); entry != alloc[type].end(); entry++) {
      void *ptr = entry->first;
      MemAlloc a = entry->second;
      printfQuda("%s  %15p  %15lu  %s(), %s:%d\n", type_str[type], ptr, (unsigned long) a.base_size,
     a.func.c_str(), a.file.c_str(), a.line);
    }
  }


  static void track_malloc(const AllocType &type, const MemAlloc &a, void *ptr)
  {
    total_bytes[type] += a.base_size;
    if (total_bytes[type] > max_total_bytes[type]) {
      max_total_bytes[type] = total_bytes[type];
    }
    if (type != DEVICE) {
      total_host_bytes += a.base_size;
      if (total_host_bytes > max_total_host_bytes) {
  max_total_host_bytes = total_host_bytes;
      }
    }
    if (type == PINNED || type == MAPPED) {
      total_pinned_bytes += a.base_size;
      if (total_pinned_bytes > max_total_pinned_bytes) {
  max_total_pinned_bytes = total_pinned_bytes;
      }
    }
    alloc[type][ptr] = a;
  }


  static void track_free(const AllocType &type, void *ptr)
  {
    size_t size = alloc[type][ptr].base_size;
    total_bytes[type] -= size;
    if (type != DEVICE) {
      total_host_bytes -= size;
    }
    if (type == PINNED || type == MAPPED) {
      total_pinned_bytes -= size;
    }
    alloc[type].erase(ptr);
  }


  static void *aligned_malloc(MemAlloc &a, size_t size)
  {
    void *ptr = nullptr;

    a.size = size;

#if (CUDA_VERSION > 4000) && 0 // we need to manually align to page boundaries to allow us to bind a texture to mapped memory
    a.base_size = size;
    ptr = malloc(size);
    if (!ptr ) {
#else
    static int page_size = 2*getpagesize();
    a.base_size = ((size + page_size - 1) / page_size) * page_size; // round up to the nearest multiple of page_size
    int align = posix_memalign(&ptr, page_size, a.base_size);
    if (!ptr || align != 0) {
#endif
      printfQuda("ERROR: Failed to allocate aligned host memory of size %zu (%s:%d in %s())\n", size, a.file.c_str(), a.line, a.func.c_str());
      errorQuda("Aborting");
    }
    return ptr;
  }


  void *device_malloc_(const char *func, const char *file, int line, size_t size)
  {
    MemAlloc a(func, file, line);
    void *ptr;

    a.size = a.base_size = size;

    cudaError_t err = cudaMalloc(&ptr, size);
    if (err != cudaSuccess) {
      printfQuda("ERROR: Failed to allocate device memory of size %zu (%s:%d in %s())\n", size, file, line, func);
      errorQuda("Aborting");
    }
    track_malloc(DEVICE, a, ptr);
#ifdef HOST_DEBUG
    cudaMemset(ptr, 0xff, size);
#endif
    return ptr;
  }


  void *device_pinned_malloc_(const char *func, const char *file, int line, size_t size)
  {
    MemAlloc a(func, file, line);
    void *ptr;

    a.size = a.base_size = size;

    CUresult err = cuMemAlloc((CUdeviceptr*)&ptr, size);
    if (err != CUDA_SUCCESS) {
      printfQuda("ERROR: Failed to allocate device memory of size %zu (%s:%d in %s())\n", size, file, line, func);
      errorQuda("Aborting");
    }
    track_malloc(DEVICE, a, ptr);
#ifdef HOST_DEBUG
    cudaMemset(ptr, 0xff, size);
#endif
    return ptr;
  }


  void *safe_malloc_(const char *func, const char *file, int line, size_t size)
  {
    MemAlloc a(func, file, line);
    a.size = a.base_size = size;

    void *ptr = malloc(size);
    if (!ptr) {
      printfQuda("ERROR: Failed to allocate host memory of size %zu (%s:%d in %s())\n", size, file, line, func);
      errorQuda("Aborting");
    }
    track_malloc(HOST, a, ptr);
#ifdef HOST_DEBUG
    memset(ptr, 0xff, size);
#endif
    return ptr;
  }


  void *pinned_malloc_(const char *func, const char *file, int line, size_t size)
  {
    MemAlloc a(func, file, line);
    void *ptr = aligned_malloc(a, size);
    
    cudaError_t err = cudaHostRegister(ptr, a.base_size, cudaHostRegisterDefault);
    if (err != cudaSuccess) {
      printfQuda("ERROR: Failed to register pinned memory of size %zu (%s:%d in %s())\n", size, file, line, func);
      errorQuda("Aborting");
    }
    track_malloc(PINNED, a, ptr);
#ifdef HOST_DEBUG
    memset(ptr, 0xff, a.base_size);
#endif
    return ptr;
  }


  void *mapped_malloc_(const char *func, const char *file, int line, size_t size)
  {
    MemAlloc a(func, file, line);
    void *ptr = aligned_malloc(a, size);
    
    cudaError_t err = cudaHostRegister(ptr, a.base_size, cudaHostRegisterMapped);
    if (err != cudaSuccess) {
      printfQuda("ERROR: Failed to register host-mapped memory of size %zu (%s:%d in %s())\n", size, file, line, func);
      errorQuda("Aborting");
    }
    track_malloc(MAPPED, a, ptr);
#ifdef HOST_DEBUG
    memset(ptr, 0xff, a.base_size);
#endif
    return ptr;
  }  


  void device_free_(const char *func, const char *file, int line, void *ptr)
  {
    if (!ptr) {
      printfQuda("ERROR: Attempt to free NULL device pointer (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    if (!alloc[DEVICE].count(ptr)) {
      printfQuda("ERROR: Attempt to free invalid device pointer (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    cudaError_t err = cudaFree(ptr);
    if (err != cudaSuccess) {
      printfQuda("ERROR: Failed to free device memory (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    track_free(DEVICE, ptr);
  }


  void device_pinned_free_(const char *func, const char *file, int line, void *ptr)
  {
    if (!ptr) {
      printfQuda("ERROR: Attempt to free NULL device pointer (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    if (!alloc[DEVICE].count(ptr)) {
      printfQuda("ERROR: Attempt to free invalid device pointer (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    CUresult err = cuMemFree((CUdeviceptr)ptr);
    if (err != CUDA_SUCCESS) {
      printfQuda("ERROR: Failed to free device memory (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    track_free(DEVICE, ptr);
  }


  void host_free_(const char *func, const char *file, int line, void *ptr)
  {
    if (!ptr) {
      printfQuda("ERROR: Attempt to free NULL host pointer (%s:%d in %s())\n", file, line, func);
      errorQuda("Aborting");
    }
    if (alloc[HOST].count(ptr)) {
      track_free(HOST, ptr);
    } else if (alloc[PINNED].count(ptr)) {
      cudaError_t err = cudaHostUnregister(ptr);
      if (err != cudaSuccess) {
  printfQuda("ERROR: Failed to unregister pinned memory (%s:%d in %s())\n", file, line, func);
  errorQuda("Aborting");
      }
      track_free(PINNED, ptr);
    } else if (alloc[MAPPED].count(ptr)) {
      cudaError_t err = cudaHostUnregister(ptr);
      if (err != cudaSuccess) {
  printfQuda("ERROR: Failed to unregister host-mapped memory (%s:%d in %s())\n", file, line, func);
  errorQuda("Aborting");
      }
      track_free(MAPPED, ptr);
    } else {
      printfQuda("ERROR: Attempt to free invalid host pointer (%s:%d in %s())\n", file, line, func);
      print_trace();
      errorQuda("Aborting");
    }
    free(ptr);
  }


  void printPeakMemUsage()
  {
    printfQuda("Device memory used = %.1f MB\n", max_total_bytes[DEVICE] / (double)(1<<20));
    printfQuda("Page-locked host memory used = %.1f MB\n", max_total_pinned_bytes / (double)(1<<20));
    printfQuda("Total host memory used >= %.1f MB\n", max_total_host_bytes / (double)(1<<20));
  }


  void assertAllMemFree()
  {
    if (!alloc[DEVICE].empty() || !alloc[HOST].empty() || !alloc[PINNED].empty() || !alloc[MAPPED].empty()) {
      warningQuda("The following internal memory allocations were not freed.");
      printfQuda("\n");
      print_alloc_header();
      print_alloc(DEVICE);
      print_alloc(HOST);
      print_alloc(PINNED);
      print_alloc(MAPPED);
      printfQuda("\n");
    }
  }


  namespace pool {

    static std::multimap<size_t, void *> pinnedCache;

    static std::map<void *, size_t> pinnedSize;

    static std::multimap<size_t, void *> deviceCache;

    static std::map<void *, size_t> deviceSize;

    static bool pool_init = false;

    static bool device_memory_pool = true;

    static bool pinned_memory_pool = true;

    void init() {
      if (!pool_init) {
  // device memory pool
  char *enable_device_pool = getenv("QUDA_ENABLE_DEVICE_MEMORY_POOL");
  if (!enable_device_pool || strcmp(enable_device_pool,"0")!=0) {
    warningQuda("Using device memory pool allocator");
    device_memory_pool = true;
  } else {
    warningQuda("Not using device memory pool allocator");
    device_memory_pool = false;
  }

  // pinned memory pool
  char *enable_pinned_pool = getenv("QUDA_ENABLE_PINNED_MEMORY_POOL");
  if (!enable_pinned_pool || strcmp(enable_pinned_pool,"0")!=0) {
    warningQuda("Using pinned memory pool allocator");
    pinned_memory_pool = true;
  } else {
    warningQuda("Not using pinned memory pool allocator");
    pinned_memory_pool = false;
  }
  pool_init = true;
      }
    }

    void* pinned_malloc_(const char *func, const char *file, int line, size_t nbytes)
    {
      void *ptr = nullptr;
      if (pinned_memory_pool) {
  std::multimap<size_t, void *>::iterator it;

  if (pinnedCache.empty()) {
    ptr = quda::pinned_malloc_(func, file, line, nbytes);
  } else {
    it = pinnedCache.lower_bound(nbytes);
    if (it != pinnedCache.end()) { // sufficiently large allocation found
      nbytes = it->first;
      ptr = it->second;
      pinnedCache.erase(it);
    } else { // sacrifice the smallest cached allocation
      it = pinnedCache.begin();
      ptr = it->second;
      pinnedCache.erase(it);
      host_free(ptr);
      ptr = quda::pinned_malloc_(func, file, line, nbytes);
    }
  }
  pinnedSize[ptr] = nbytes;
      } else {
  ptr = quda::pinned_malloc_(func, file, line, nbytes);
      }
      return ptr;
    }

    void pinned_free_(const char *func, const char *file, int line, void *ptr)
    {
      if (pinned_memory_pool) {
  if (!pinnedSize.count(ptr)) {
    errorQuda("Attempt to free invalid pointer");
  }
  pinnedCache.insert(std::make_pair(pinnedSize[ptr], ptr));
  pinnedSize.erase(ptr);
      } else {
  quda::host_free_(func, file, line, ptr);
      }
    }

    void* device_malloc_(const char *func, const char *file, int line, size_t nbytes)
    {
      void *ptr = nullptr;
      if (device_memory_pool) {
  std::multimap<size_t, void *>::iterator it;

  if (deviceCache.empty()) {
    ptr = quda::device_malloc_(func, file, line, nbytes);
  } else {
    it = deviceCache.lower_bound(nbytes);
    if (it != deviceCache.end()) { // sufficiently large allocation found
      nbytes = it->first;
      ptr = it->second;
      deviceCache.erase(it);
    } else { // sacrifice the smallest cached allocation
      it = deviceCache.begin();
      ptr = it->second;
      deviceCache.erase(it);
      quda::device_free_(func, file, line, ptr);
      ptr = quda::device_malloc_(func, file, line, nbytes);
    }
  }
  deviceSize[ptr] = nbytes;
      } else {
  ptr = quda::device_malloc_(func, file, line, nbytes);
      }
      return ptr;
    }

    void device_free_(const char *func, const char *file, int line, void *ptr)
    {
      if (device_memory_pool) {
  if (!deviceSize.count(ptr)) {
    errorQuda("Attempt to free invalid pointer");
  }
  deviceCache.insert(std::make_pair(deviceSize[ptr], ptr));
  deviceSize.erase(ptr);
      } else {
  quda::device_free_(func, file, line, ptr);
      }
    }

    void flush_pinned()
    {
      if (pinned_memory_pool) {
  std::multimap<size_t, void *>::iterator it;
  for (it = pinnedCache.begin(); it != pinnedCache.end(); it++) {
    void *ptr = it->second;
    host_free(ptr);
  }
  pinnedCache.clear();
      }
    }

    void flush_device()
    {
      if (device_memory_pool) {
  std::multimap<size_t, void *>::iterator it;
  for (it = deviceCache.begin(); it != deviceCache.end(); it++) {
    void *ptr = it->second;
    device_free(ptr);
  }
  deviceCache.clear();
      }
    }

  } // namespace pool

} // namespace quda