YAKL_allocators.h

Go to the documentation of this file.

 
#pragma once
// Included by YAKL.h
 
__YAKL_NAMESPACE_WRAPPER_BEGIN__
namespace yakl {
 
  inline bool use_pool() { return get_yakl_instance().pool_enabled; }
 
 
  inline Gator & get_pool() { return get_yakl_instance().pool; }
 
 
  // Set the allocation and deallocation functions for YAKL
  inline void set_device_alloc_free(std::function<void *( size_t )> &alloc , std::function<void ( void * )> &dealloc) {
    #if   defined(YAKL_ARCH_CUDA)
      #if defined (YAKL_MANAGED_MEMORY)
        alloc   = [] ( size_t bytes ) -> void* {
          if (bytes == 0) return nullptr;
          void *ptr;
          cudaMallocManaged(&ptr,bytes);      // Allocate managed memory
          if (ptr == nullptr) yakl_throw("ERROR: cudaMallocManaged returned nullptr. You have likely run out of memory");
          #ifdef _OPENMP45
            // if using OMP target offload, make sure OMP runtime knows to leave this memory alone
            omp_target_associate_ptr(ptr,ptr,bytes,0,0);
          #endif
          #ifdef _OPENACC
            // if using OpenACC, make sure OpenACC runtime knows to leave this memory alone
            acc_map_data(ptr,ptr,bytes);
          #endif
          check_last_error();
          return ptr;
        };
        dealloc = [] ( void *ptr    ) {
          cudaFree(ptr);
          check_last_error();
        };
      #else
        alloc   = [] ( size_t bytes ) -> void* {
          if (bytes == 0) return nullptr;
          void *ptr;
          cudaMalloc(&ptr,bytes);
          if (ptr == nullptr) yakl_throw("ERROR: cudaMalloc returned nullptr. You have likely run out of memory");
          check_last_error();
          return ptr;
        };
        dealloc = [] ( void *ptr    ) {
          cudaFree(ptr);
          check_last_error();
        };
      #endif
    #elif defined(YAKL_ARCH_HIP)
      #if defined (YAKL_MANAGED_MEMORY)
        alloc = [] ( size_t bytes ) -> void* {
          if (bytes == 0) return nullptr;
          void *ptr;
          hipMallocManaged(&ptr,bytes);  // This is the current standin for managed memory for HIP
          if (ptr == nullptr) yakl_throw("ERROR: hipMallocManaged returned nullptr. You have likely run out of memory");
          #ifdef _OPENMP45
            // if using OMP target offload, make sure OMP runtime knows to leave this memory alone
            omp_target_associate_ptr(ptr,ptr,bytes,0,0);
          #endif
          #ifdef _OPENACC
            // if using OpenACC, make sure OpenACC runtime knows to leave this memory alone
            acc_map_data(ptr,ptr,bytes);
          #endif
          check_last_error();
          return ptr;
        };
        dealloc = [] ( void *ptr    ) {
          hipFree(ptr);
          check_last_error();
        };
      #else
        alloc = [] ( size_t bytes ) -> void* {
          if (bytes == 0) return nullptr;
          void *ptr;
          hipMalloc(&ptr,bytes);
          if (ptr == nullptr) yakl_throw("ERROR: hipMalloc returned nullptr. You have likely run out of memory");
          check_last_error();
          return ptr;
        };
        dealloc = [] ( void *ptr ) {
          hipFree(ptr);
          check_last_error();
        };
      #endif
    #elif defined (YAKL_ARCH_SYCL)
      #if defined (YAKL_MANAGED_MEMORY)
        alloc = [] ( size_t bytes ) -> void* {
          if (bytes == 0) return nullptr;
          // Allocate unified shared memory
          void *ptr = sycl::malloc_shared(bytes,sycl_default_stream());
          if (ptr == nullptr) yakl_throw("ERROR: sycl::malloc_shared returned nullptr. You have likely run out of memory");
          #ifdef _OPENMP45
            // if using OMP target offload, make sure OMP runtime knows to leave this memory alone
            omp_target_associate_ptr(ptr,ptr,bytes,0,0);
          #endif
          #ifdef _OPENACC
            // if using OpenACC, make sure OpenACC runtime knows to leave this memory alone
            acc_map_data(ptr,ptr,bytes);
          #endif
          return ptr;
        };
        dealloc = [] ( void *ptr ) {
          sycl::free(ptr, sycl_default_stream());
          check_last_error();
        };
      #else
        alloc = [] ( size_t bytes ) -> void* {
          if (bytes == 0) return nullptr;
          void *ptr = sycl::malloc_device(bytes,sycl_default_stream());
          if (ptr == nullptr) yakl_throw("ERROR: sycl::malloc_device returned nullptr. You have likely run out of memory");
          check_last_error();
          return ptr;
        };
        dealloc = [] ( void *ptr ) {
          sycl::free(ptr, sycl_default_stream());
          check_last_error();
        };
      #endif
    #else
      alloc   = [] ( size_t bytes ) -> void* {
        if (bytes == 0) return nullptr;
        void *ptr = ::malloc(bytes);
        if (ptr == nullptr) yakl_throw("ERROR: malloc returned nullptr. You have likely run out of memory");
        return ptr;
      };
      dealloc = [] ( void *ptr ) {
        ::free(ptr);
      };
    #endif
  }
 
 
  inline void set_yakl_allocators_to_default() {
    fence();
    if (use_pool()) {
      get_yakl_instance().alloc_device_func = [] (size_t bytes , char const *label) -> void * {
        #ifdef YAKL_MEMORY_DEBUG
          if (yakl_mainproc()) std::cout << "MEMORY_DEBUG: Allocating label \"" << label << "\" of size " << bytes << " bytes" << std::endl;
        #endif
        void * ptr = get_yakl_instance().pool.allocate( bytes , label );
        #ifdef YAKL_MEMORY_DEBUG
          if (yakl_mainproc()) std::cout << "MEMORY_DEBUG: Successfully allocated label \"" << label
                                         << " with pointer address " << ptr << std::endl;
        #endif
        return ptr;
      };
      get_yakl_instance().free_device_func  = [] (void *ptr , char const *label)              {
        #ifdef YAKL_MEMORY_DEBUG
          if (yakl_mainproc()) std::cout << "MEMORY_DEBUG: Freeing label \"" << label << "\" with pointer address " << ptr << std::endl;
        #endif
        get_yakl_instance().pool.free( ptr , label );
      };
    } else {
      std::function<void *( size_t)> alloc;
      std::function<void ( void *)>  dealloc;
      set_device_alloc_free(alloc , dealloc);
      get_yakl_instance().alloc_device_func = [=] (size_t bytes , char const *label) -> void * {
        #ifdef YAKL_MEMORY_DEBUG
          if (yakl_mainproc()) std::cout << "MEMORY_DEBUG: Allocating label \"" << label << "\" of size " << bytes << " bytes" << std::endl;
        #endif
        void * ptr = alloc(bytes);
        #ifdef YAKL_MEMORY_DEBUG
          if (yakl_mainproc()) std::cout << "MEMORY_DEBUG: Successfully allocated label \"" << label
                                         << " with pointer address " << ptr << std::endl;
        #endif
        return ptr;
      };
      get_yakl_instance().free_device_func  = [=] (void *ptr , char const *label)              {
        #ifdef YAKL_MEMORY_DEBUG
          if (yakl_mainproc()) std::cout << "MEMORY_DEBUG: Freeing label \"" << label << "\" with pointer address " << ptr << std::endl;
        #endif
        dealloc(ptr);
      };
    }
 
    get_yakl_instance().device_allocators_are_default = true;
  }
 
 
  inline void set_device_allocator  ( std::function<void *(size_t)> func ) {
    fence();   get_yakl_instance().alloc_device_func = [=] (size_t bytes , char const *label) -> void * { return func(bytes); };
    get_yakl_instance().device_allocators_are_default = false;
  }
 
 
  inline void set_device_deallocator( std::function<void (void *)>  func ) {
    fence();   get_yakl_instance().free_device_func  = [=] (void *ptr , char const *label) { func(ptr); };
    get_yakl_instance().device_allocators_are_default = false;
  }
 
 
  inline void set_device_allocator  ( std::function<void *( size_t , char const *)> func ) { fence();  get_yakl_instance().alloc_device_func = func; }
 
 
  inline void set_device_deallocator( std::function<void ( void * , char const *)>  func ) { fence();  get_yakl_instance().free_device_func  = func; }
 
  inline void * alloc_device( size_t bytes, char const *label) { return get_yakl_instance().alloc_device_func(bytes,label); }
 
  inline void   free_device ( void * ptr  , char const *label) {        get_yakl_instance().free_device_func (ptr  ,label); }
 
}
__YAKL_NAMESPACE_WRAPPER_END__