/*
* (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.
*/

#pragma once

#include <cstddef>

#include "atlas/library/config.h"

namespace atlas {
namespace util {

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

namespace detail {
void allocate_cudamanaged(void** ptr, size_t size);
void deallocate_cudamanaged(void* ptr);

void allocate_cuda(void** ptr, size_t size);
void deallocate_cuda(void* ptr);

void allocate_host(void** ptr, size_t size);
void deallocate_host(void* ptr);

}  // namespace detail

template <typename T>
void allocate_managedmem(T*& data, size_t N) {
    if (N != 0) {
        detail::allocate_cudamanaged(reinterpret_cast<void**>(&data), N * sizeof(T));
    }
}

template <typename T>
void delete_managedmem(T*& data) {
    if (data) {
        detail::deallocate_cudamanaged(data);
        data = nullptr;
    }
}

template <typename T>
void allocate_devicemem(T*& data, size_t N) {
    if (N != 0) {
        detail::allocate_cuda(reinterpret_cast<void**>(&data), N * sizeof(T));
    }
}

template <typename T>
void delete_devicemem(T*& data) {
    if (data) {
        detail::deallocate_cuda(data);
        data = nullptr;
    }
}

template <typename T>
void allocate_hostmem(T*& data, size_t N) {
    if (N != 0) {
        detail::allocate_host(reinterpret_cast<void**>(&data), N * sizeof(T));
    }
}

template <typename T>
void delete_hostmem(T*& data) {
    if (data) {
        detail::deallocate_host(data);
        data = nullptr;
    }
}


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

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

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

}  // namespace util
}  // namespace atlas
