// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
#include <Kokkos_Macros.hpp>
static_assert(false,
              "Including non-public Kokkos header files is not allowed.");
#endif
#ifndef KOKKOS_HOSTSPACE_HPP
#define KOKKOS_HOSTSPACE_HPP

#include <cstring>
#include <string>
#include <iosfwd>
#include <typeinfo>

#include <Kokkos_Core_fwd.hpp>
#include <Kokkos_Concepts.hpp>
#include <Kokkos_MemoryTraits.hpp>

#include <impl/Kokkos_Traits.hpp>
#include <impl/Kokkos_Error.hpp>
#include <impl/Kokkos_SharedAlloc.hpp>
#include <impl/Kokkos_Tools.hpp>

#include "impl/Kokkos_HostSpace_deepcopy.hpp"

/*--------------------------------------------------------------------------*/

namespace Kokkos {
/// \class HostSpace
/// \brief Memory management for host memory.
///
/// HostSpace is a memory space that governs host memory.  "Host"
/// memory means the usual CPU-accessible memory.
class HostSpace {
 public:
  //! Tag this class as a kokkos memory space
  using memory_space = HostSpace;
  using size_type    = size_t;
  using index_type   = std::make_signed_t<size_type>;

  /// \typedef execution_space
  /// \brief Default execution space for this memory space.
  ///
  /// Every memory space has a default execution space.  This is
  /// useful for things like initializing a View (which happens in
  /// parallel using the View's default execution space).
  using execution_space = DefaultHostExecutionSpace;

  //! This memory space preferred device_type
  using device_type = Kokkos::Device<execution_space, memory_space>;

  HostSpace() = default;

  /**\brief  Allocate untracked memory in the space */
  template <typename ExecutionSpace>
  void* allocate(const ExecutionSpace&, const size_t arg_alloc_size) const {
    return allocate(arg_alloc_size);
  }
  template <typename ExecutionSpace>
  void* allocate(const ExecutionSpace&, const char* arg_label,
                 const size_t arg_alloc_size) const {
    return allocate(arg_label, arg_alloc_size);
  }
  template <typename ExecutionSpace>
  void* allocate(const ExecutionSpace&, const char* arg_label,
                 const size_t arg_alloc_size,
                 const size_t arg_logical_size) const {
    return allocate(arg_label, arg_alloc_size, arg_logical_size);
  }
  void* allocate(const size_t arg_alloc_size) const;
  void* allocate(const char* arg_label, const size_t arg_alloc_size) const;
  void* allocate(const char* arg_label, const size_t arg_alloc_size,
                 const size_t arg_logical_size) const;

  /**\brief  Deallocate untracked memory in the space */
  void deallocate(void* const arg_alloc_ptr, const size_t arg_alloc_size) const;
  void deallocate(const char* arg_label, void* const arg_alloc_ptr,
                  const size_t arg_alloc_size) const;
  void deallocate(const char* arg_label, void* const arg_alloc_ptr,
                  const size_t arg_alloc_size,
                  const size_t arg_logical_size) const;

  void* impl_allocate(const char* arg_label, const size_t arg_alloc_size,
                      const Kokkos::Tools::SpaceHandle =
                          Kokkos::Tools::make_space_handle(name())) const;
  void* impl_allocate(const char* arg_label, const size_t arg_alloc_size,
                      const size_t arg_reported_size,
                      const Kokkos::Tools::SpaceHandle =
                          Kokkos::Tools::make_space_handle(name())) const;
  void impl_deallocate(const char* arg_label, void* const arg_alloc_ptr,
                       const size_t arg_alloc_size,
                       const Kokkos::Tools::SpaceHandle =
                           Kokkos::Tools::make_space_handle(name())) const;
  void impl_deallocate(const char* arg_label, void* const arg_alloc_ptr,
                       const size_t arg_alloc_size,
                       const size_t arg_reported_size,
                       const Kokkos::Tools::SpaceHandle =
                           Kokkos::Tools::make_space_handle(name())) const;

  /**\brief Return Name of the MemorySpace */
  static constexpr const char* name() { return m_name; }

 private:
  static constexpr const char* m_name = "Host";
};

}  // namespace Kokkos

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

namespace Kokkos {

namespace Impl {

static_assert(Kokkos::Impl::MemorySpaceAccess<Kokkos::HostSpace,
                                              Kokkos::HostSpace>::assignable);

template <typename MemSpace>
struct HostMirror {
 private:
  static_assert(is_memory_space_v<MemSpace>);

  // If input execution space can access HostSpace then keep it.
  // Example: Kokkos::OpenMP can access, Kokkos::Cuda cannot
  enum {
    keep_exe = Kokkos::SpaceAccessibility<typename MemSpace::execution_space,
                                          Kokkos::HostSpace>::accessible
  };
  // If HostSpace can access memory space then keep it.
  // Example: Cannot access Kokkos::CudaSpace, can access Kokkos::CudaUVMSpace
  enum {
    keep_mem =
        Kokkos::Impl::MemorySpaceAccess<Kokkos::HostSpace, MemSpace>::accessible
  };

 public:
  // Construct a device mirror type
  // Decision logic: First check if HostSpace can access the memory space.
  // If yes, keep it and check execution space compatibility.
  // If no, fall back to HostSpace::device_type.

  // keep_exe | keep_mem | Result
  // ---------|----------|-------
  //    T     |    T     | MemSpace::device_type
  //    F     |    T     | Device<HostSpace::execution_space, MemSpace>
  //    T     |    F     | HostSpace::device_type
  //    F     |    F     | HostSpace::device_type

  using device_type = std::conditional_t<
      keep_mem,
      std::conditional_t<
          keep_exe, typename MemSpace::device_type,
          Kokkos::Device<Kokkos::HostSpace::execution_space, MemSpace>>,
      Kokkos::HostSpace::device_type>;

  using execution_space = typename device_type::execution_space;
  using memory_space    = typename device_type::memory_space;

  // FIXME: should be deprecated eventually
  using Space = memory_space;
};

}  // namespace Impl

}  // namespace Kokkos

namespace Kokkos {

namespace Impl {

template <class ExecutionSpace>
struct DeepCopy<HostSpace, HostSpace, ExecutionSpace> {
  DeepCopy(void* dst, const void* src, size_t n) {
    hostspace_parallel_deepcopy(dst, src, n);
  }

  DeepCopy(const ExecutionSpace& exec, void* dst, const void* src, size_t n) {
    if constexpr (!Kokkos::SpaceAccessibility<ExecutionSpace,
                                              Kokkos::HostSpace>::accessible) {
      exec.fence(
          "Kokkos::Impl::DeepCopy<HostSpace, HostSpace, "
          "ExecutionSpace>::DeepCopy: fence before copy");
      hostspace_parallel_deepcopy_async(dst, src, n);
    } else {
      hostspace_parallel_deepcopy_async(exec, dst, src, n);
    }
  }
};

}  // namespace Impl

}  // namespace Kokkos

#endif  // #define KOKKOS_HOSTSPACE_HPP