/*
* (C) Copyright 2013 ECMWF.
*
* This software is licensed under the terms of the Apache Licence Version 2.0
* which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
* In applying this licence, ECMWF does not waive the privileges and immunities
* granted to it by virtue of its status as an intergovernmental organisation nor
* does it submit to any jurisdiction.
*/


#include "Allocate.h"

#include "atlas/library/config.h"
#include "atlas/runtime/Exception.h"
#include "eckit/log/CodeLocation.h"

#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
#include <cuda_runtime.h>
#endif

namespace atlas {
namespace util {

//------------------------------------------------------------------------------
namespace detail {
//------------------------------------------------------------------------------

void allocate_cudamanaged(void** ptr, size_t size) {
#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
    cudaError_t err = cudaMallocManaged(ptr, size);
    if (err != cudaSuccess)
        throw_AssertionFailed("failed to allocate GPU memory", Here());
#else
    *ptr = malloc(size);
#endif
}

void deallocate_cudamanaged(void* ptr) {
#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
    cudaError_t err = cudaDeviceSynchronize();
    if (err != cudaSuccess)
        throw_AssertionFailed("failed to synchronize memory", Here());

    err = cudaFree(ptr);
    // The following throws an invalid device memory
    if (err != cudaSuccess)
        throw_AssertionFailed("failed to free GPU memory", Here());
#else
    free(ptr);
#endif
}

void allocate_cuda(void** ptr, size_t size) {
#if ATLAS_GRIDTOOLS_STORAGE_BACKEND_CUDA
    cudaError_t err = cudaMalloc(ptr, size);
    if (err != cudaSuccess)
        throw_AssertionFailed("failed to allocate GPU memory", Here());
#else
    *ptr = malloc(size);
#endif
}

void deallocate_cuda(void* ptr) {
    deallocate_cudamanaged(ptr);
}

void allocate_host(void** ptr, size_t size) {
    *ptr = malloc(size);
}

void deallocate_host(void* ptr) {
    free(ptr);
}

//------------------------------------------------------------------------------
}  // namespace detail
//------------------------------------------------------------------------------

extern "C" {
void atlas__allocate_managedmem_double(double*& a, size_t N) {
    allocate_managedmem(a, N);
}
void atlas__allocate_managedmem_float(float*& a, size_t N) {
    allocate_managedmem(a, N);
}
void atlas__allocate_managedmem_int(int*& a, size_t N) {
    allocate_managedmem(a, N);
}
void atlas__allocate_managedmem_long(long*& a, size_t N) {
    allocate_managedmem(a, N);
}
void atlas__deallocate_managedmem(void*& a) {
    delete_managedmem(a);
}
}

//------------------------------------------------------------------------------

}  // namespace util
}  // namespace atlas
