//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
// Copyright (c) 2016-21, Lawrence Livermore National Security, LLC
// and RAJA project contributors. See the RAJA/COPYRIGHT file for details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

#pragma once

#include "RAJA/RAJA.hpp"

#if defined(RAJA_ENABLE_CUDA)
  #include "RAJA/policy/cuda/raja_cudaerrchk.hpp"
#endif

#if defined(RAJA_ENABLE_HIP)
  #include "RAJA/policy/hip/raja_hiperrchk.hpp"
#endif

/*
  As RAJA does not manage memory we include a general purpose memory
  manager which may be used to perform c++ style allocation/deallocation
  or allocate/deallocate CUDA unified memory. The type of memory allocated
  is dependent on how RAJA was configured.
*/
namespace memoryManager
{

#if defined(RAJA_ENABLE_SYCL)
static camp::resources::Resource *sycl_res;
#endif

template <typename T>
T *allocate(RAJA::Index_type size)
{
  T *ptr;
#if defined(RAJA_ENABLE_CUDA)
  cudaErrchk(cudaMallocManaged((void **)&ptr, sizeof(T) * size, cudaMemAttachGlobal));
#elif defined(RAJA_ENABLE_HIP)
  hipErrchk(hipMalloc((void **)&ptr, sizeof(T) * size));
#elif defined(RAJA_ENABLE_SYCL)
  ptr = sycl_res->allocate<T>(size);
#else
  ptr = new T[size];
#endif
  return ptr;
}

template <typename T>
void deallocate(T *&ptr)
{
  if (ptr) {
#if defined(RAJA_ENABLE_CUDA)
    cudaErrchk(cudaFree(ptr));
#elif defined(RAJA_ENABLE_HIP)
    hipErrchk(hipFree(ptr));
#elif defined(RAJA_ENABLE_SYCL)
    sycl_res->deallocate(ptr);
#else
    delete[] ptr;
#endif
    ptr = nullptr;
  }
}

#if defined(RAJA_ENABLE_CUDA) || defined(RAJA_ENABLE_HIP) || defined(RAJA_ENABLE_SYCL)
template <typename T>
T *allocate_gpu(RAJA::Index_type size)
{
  T *ptr;
  #if defined(RAJA_ENABLE_CUDA)
  cudaErrchk(cudaMalloc((void **)&ptr, sizeof(T) * size));
  #elif defined(RAJA_ENABLE_HIP)
  hipErrchk(hipMalloc((void **)&ptr, sizeof(T) * size));
  #elif defined(RAJA_ENABLE_SYCL)
  auto qu = sycl_res->get<camp::resources::Sycl>().get_queue();
  ptr     = cl::sycl::malloc_device<T>(size, *qu);
  #endif
  return ptr;
}

template <typename T>
void deallocate_gpu(T *&ptr)
{
  if (ptr) {
  #if defined(RAJA_ENABLE_CUDA)
    cudaErrchk(cudaFree(ptr));
  #elif defined(RAJA_ENABLE_HIP)
    hipErrchk(hipFree(ptr));
  #elif defined(RAJA_ENABLE_SYCL)
    sycl_res->deallocate(ptr);
  #endif
    ptr = nullptr;
  }
}
#endif

}; // namespace memoryManager
