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

#ifndef KOKKOS_SYCL_INSTANCE_HPP_
#define KOKKOS_SYCL_INSTANCE_HPP_

#include <optional>
#include <sycl/sycl.hpp>

#include <impl/Kokkos_Error.hpp>
#include <impl/Kokkos_Profiling.hpp>
namespace Kokkos {
namespace Impl {

class SYCLInternal {
 public:
  using size_type = unsigned int;

  SYCLInternal() = default;
  ~SYCLInternal();

  SYCLInternal(const SYCLInternal&)            = delete;
  SYCLInternal& operator=(const SYCLInternal&) = delete;
  SYCLInternal& operator=(SYCLInternal&&)      = delete;
  SYCLInternal(SYCLInternal&&)                 = delete;

  Kokkos::Impl::sycl_device_ptr<void> scratch_space(const std::size_t size);
  Kokkos::Impl::sycl_device_ptr<void> scratch_flags(const std::size_t size);
  Kokkos::Impl::sycl_host_ptr<void> scratch_host(const std::size_t size);
  int acquire_team_scratch_space();
  Kokkos::Impl::sycl_device_ptr<void> resize_team_scratch_space(
      int scratch_pool_id, std::int64_t bytes, bool force_shrink = false);
  void register_team_scratch_event(int scratch_pool_id, sycl::event event);

  uint32_t impl_get_instance_id() const;
  static int m_syclDev;

  size_t m_maxWorkgroupSize   = 0;
  uint32_t m_maxConcurrency   = 0;
  uint64_t m_maxShmemPerBlock = 0;

  std::size_t m_scratchSpaceCount                         = 0;
  Kokkos::Impl::sycl_device_ptr<size_type> m_scratchSpace = nullptr;
  std::size_t m_scratchHostCount                          = 0;
  Kokkos::Impl::sycl_host_ptr<size_type> m_scratchHost    = nullptr;
  std::size_t m_scratchFlagsCount                         = 0;
  Kokkos::Impl::sycl_device_ptr<size_type> m_scratchFlags = nullptr;
  // mutex to access shared memory
  mutable std::mutex m_mutexScratchSpace;

  // Team Scratch Level 1 Space
  static constexpr int m_n_team_scratch                         = 10;
  mutable int64_t m_team_scratch_current_size[m_n_team_scratch] = {};
  mutable Kokkos::Impl::sycl_device_ptr<void>
      m_team_scratch_ptr[m_n_team_scratch]                   = {};
  mutable int m_current_team_scratch                         = 0;
  mutable sycl::event m_team_scratch_event[m_n_team_scratch] = {};
  mutable std::mutex m_team_scratch_mutex;

  uint32_t m_instance_id =
      Kokkos::Tools::Experimental::Impl::idForInstance<Kokkos::SYCL>(
          reinterpret_cast<uintptr_t>(this));
  std::optional<sycl::queue> m_queue;

  // Using std::vector<std::optional<sycl::queue>> reveals a compiler bug when
  // compiling for the CUDA backend. Storing pointers instead works around this.
  static std::vector<std::optional<sycl::queue>*> all_queues;
  // We need a mutex for thread safety when modifying all_queues.
  static std::mutex mutex;

  // USMObjectMem is a reusable buffer for a single object
  // in USM memory
  template <sycl::usm::alloc Kind>
  class USMObjectMem {
   public:
    void reset();

    void reset(sycl::queue q, uint32_t instance_id) {
      m_instance_id = instance_id;
      reset();
      m_q.emplace(std::move(q));
    }
    USMObjectMem() = default;
    explicit USMObjectMem(sycl::queue q, uint32_t instance_id) noexcept
        : m_q(std::move(q)), m_instance_id(instance_id) {}

    USMObjectMem(USMObjectMem const&)            = delete;
    USMObjectMem(USMObjectMem&&)                 = delete;
    USMObjectMem& operator=(USMObjectMem&&)      = delete;
    USMObjectMem& operator=(USMObjectMem const&) = delete;

    ~USMObjectMem() { reset(); }

    void* data() noexcept { return m_data; }
    const void* data() const noexcept { return m_data; }

    size_t capacity() const noexcept { return m_capacity; }

    // reserve() allocates space for at least n bytes
    // returns the new capacity
    size_t reserve(size_t n);

   private:
    using AllocationSpace =
        std::conditional_t<Kind == sycl::usm::alloc::device,
                           Kokkos::SYCLDeviceUSMSpace,
                           std::conditional_t<Kind == sycl::usm::alloc::shared,
                                              Kokkos::SYCLSharedUSMSpace,
                                              Kokkos::SYCLHostUSMSpace>>;

   public:
    // Performs either sycl::memcpy (for USM device memory) or std::memcpy
    // (otherwise) and returns a reference to the copied object.
    template <typename T>
    T& copy_from(const T& t) {
      m_mutex.lock();
      fence();
      reserve(sizeof(T));
      if constexpr (sycl::usm::alloc::device == Kind) {
        std::memcpy(static_cast<void*>(m_staging.get()),
                    static_cast<const void*>(std::addressof(t)), sizeof(T));
        m_copy_event = m_q->memcpy(m_data, m_staging.get(), sizeof(T));
      } else
        std::memcpy(static_cast<void*>(m_data),
                    static_cast<const void*>(std::addressof(t)), sizeof(T));
      return *reinterpret_cast<T*>(m_data);
    }

    void fence() {
      SYCLInternal::fence(m_last_event,
                          "Kokkos::SYCLInternal::USMObject fence to wait for "
                          "last event to finish",
                          m_instance_id);
    }

    void register_event(sycl::event event) {
      KOKKOS_ASSERT(
          m_last_event
              .get_info<sycl::info::event::command_execution_status>() ==
          sycl::info::event_command_status::complete);
      m_last_event = event;
      m_mutex.unlock();
    }

    sycl::event get_copy_event() const { return m_copy_event; }

   private:
    // USMObjectMem class invariants
    // All four expressions below must evaluate to true:
    //
    //  !m_data == (m_capacity == 0)
    //      m_q || !m_data
    //
    //  The above invariants mean that:
    //  if m_data != nullptr then m_capacity != 0 && m_q != nullopt
    //  if m_data == nullptr then m_capacity == 0

    sycl::event m_copy_event;

    std::optional<sycl::queue> m_q;
    void* m_data = nullptr;
    std::unique_ptr<char[]> m_staging;

    size_t m_capacity = 0;
    sycl::event m_last_event;

    uint32_t m_instance_id;

    // mutex to access the underlying memory
    mutable std::mutex m_mutex;
  };

  // An indirect kernel is one where the functor to be executed is explicitly
  // copied to USM memory before being executed, to get around the
  // trivially copyable limitation of SYCL.
  using IndirectKernelMem = USMObjectMem<sycl::usm::alloc::host>;
  IndirectKernelMem& get_indirect_kernel_mem();

  bool was_finalized = false;

  static SYCLInternal& singleton();

  int verify_is_initialized(const char* const label) const;

  void initialize(const sycl::device& d);

  void initialize(const sycl::queue& q);

  int is_initialized() const { return m_queue.has_value(); }

  void finalize();

 private:
  // fence(...) takes any type with a .wait_and_throw() method
  // (sycl::event and sycl::queue)
  template <typename WAT>
  static void fence_helper(WAT& wat, const std::string& name,
                           uint32_t instance_id);

  const static size_t m_usm_pool_size = 4;
  std::vector<IndirectKernelMem> m_indirectKernelMem{m_usm_pool_size};

  size_t m_pool_next{0};

 public:
  static void fence(sycl::queue& q, const std::string& name,
                    uint32_t instance_id) {
    fence_helper(q, name, instance_id);
  }
  static void fence(sycl::event& e, const std::string& name,
                    uint32_t instance_id) {
    fence_helper(e, name, instance_id);
  }
};

// FIXME_SYCL the limit is 2048 bytes for all arguments handed to a kernel,
// assume for now that the rest doesn't need more than 248 bytes.
#if defined(SYCL_DEVICE_COPYABLE) && defined(KOKKOS_ARCH_INTEL_GPU)
template <typename Functor, typename Storage,
          bool ManualCopy = (sizeof(Functor) >= 1800)>
class SYCLFunctionWrapper;
#else
template <typename Functor, typename Storage,
          bool ManualCopy = (sizeof(Functor) >= 1800 ||
                             !std::is_trivially_copyable_v<Functor>)>
class SYCLFunctionWrapper;
#endif

#if defined(SYCL_DEVICE_COPYABLE) && defined(KOKKOS_ARCH_INTEL_GPU)
template <typename Functor, typename Storage>
class SYCLFunctionWrapper<Functor, Storage, false> {
  // We need a union here so that we can avoid calling a constructor for m_f
  // and can controll all the special member functions.
  union TrivialWrapper {
    TrivialWrapper() {}

    TrivialWrapper(const Functor& f) {
      std::memcpy(static_cast<void*>(&m_f), static_cast<const void*>(&f),
                  sizeof(m_f));
    }

    TrivialWrapper(const TrivialWrapper& other) {
      std::memcpy(static_cast<void*>(&m_f),
                  static_cast<const void*>(&other.m_f), sizeof(m_f));
    }
    TrivialWrapper& operator=(const TrivialWrapper& other) {
      std::memcpy(static_cast<void*>(&m_f),
                  static_cast<const void*>(&other.m_f), sizeof(m_f));
      return *this;
    }
    ~TrivialWrapper() {}

    Functor m_f;
  } m_functor;

 public:
  SYCLFunctionWrapper(const Functor& functor, Storage&) : m_functor(functor) {}

  const Functor& get_functor() const { return m_functor.m_f; }

  sycl::event get_copy_event() const { return {}; }

  static void register_event(sycl::event) {}
};
#else
template <typename Functor, typename Storage>
class SYCLFunctionWrapper<Functor, Storage, false> {
  const Functor m_functor;

 public:
  SYCLFunctionWrapper(const Functor& functor, Storage&) : m_functor(functor) {}

  const Functor& get_functor() const { return m_functor; }

  sycl::event get_copy_event() const { return {}; }

  static void register_event(sycl::event) {}
};
#endif

template <typename Functor, typename Storage>
class SYCLFunctionWrapper<Functor, Storage, true> {
  std::reference_wrapper<const Functor> m_kernelFunctor;
  std::reference_wrapper<Storage> m_storage;

 public:
  SYCLFunctionWrapper(const Functor& functor, Storage& storage)
      : m_kernelFunctor(storage.copy_from(functor)), m_storage(storage) {}

  std::reference_wrapper<const Functor> get_functor() const {
    return m_kernelFunctor;
  }

  sycl::event get_copy_event() const {
    return m_storage.get().get_copy_event();
  }

  void register_event(sycl::event event) {
    m_storage.get().register_event(event);
  }
};

template <typename Functor, typename Storage>
auto make_sycl_function_wrapper(const Functor& functor, Storage& storage) {
  return SYCLFunctionWrapper<Functor, Storage>(functor, storage);
}
}  // namespace Impl
}  // namespace Kokkos

#if defined(SYCL_DEVICE_COPYABLE) && defined(KOKKOS_ARCH_INTEL_GPU)
template <typename Functor, typename Storage>
struct sycl::is_device_copyable<
    Kokkos::Impl::SYCLFunctionWrapper<Functor, Storage, false>>
    : std::true_type {};

#endif
#endif
