//==---------------- atomic.hpp - SYCL atomics -----------------------------==//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#pragma once

#include <CL/__spirv/spirv_types.hpp>         // for Scope, MemorySemanticsMask
#include <sycl/access/access.hpp>             // for address_space, decorated
#include <sycl/detail/defines_elementary.hpp> // for __SYCL2020_DEPRECATED
#include <sycl/detail/helpers.hpp>            // for getSPIRVMemorySemanticsMask
#include <sycl/memory_enums.hpp> // for memory_order, getStdMemoryO...
#include <sycl/multi_ptr.hpp>    // for multi_ptr

#include <type_traits> // for is_same

#ifndef __SYCL_DEVICE_ONLY__
#include <atomic> // for atomic, memory_order
#endif

#define __SYCL_STATIC_ASSERT_NOT_FLOAT(T)                                      \
  static_assert(!std::is_same<T, float>::value,                                \
                "SYCL atomic function not available for float type")

namespace sycl {
inline namespace _V1 {

// Forward declaration
template <typename pointerT, access::address_space addressSpace,
          access::decorated isDecorated>
class multi_ptr;

namespace detail {

using memory_order = sycl::memory_order;

template <typename T> struct IsValidAtomicType {
  static constexpr bool value =
      (std::is_same_v<T, int> || std::is_same_v<T, unsigned int> ||
       std::is_same_v<T, long> || std::is_same_v<T, unsigned long> ||
       std::is_same_v<T, long long> || std::is_same_v<T, unsigned long long> ||
       std::is_same_v<T, float>);
};

template <sycl::access::address_space AS> struct IsValidAtomicAddressSpace {
  static constexpr bool value =
      (AS == access::address_space::global_space ||
       AS == access::address_space::local_space ||
       AS == access::address_space::ext_intel_global_device_space);
};

// Type trait to translate a sycl::access::address_space to
// a SPIR-V memory scope
template <access::address_space AS> struct GetSpirvMemoryScope {};
template <> struct GetSpirvMemoryScope<access::address_space::global_space> {
  static constexpr auto scope = __spv::Scope::Device;
};
template <>
struct GetSpirvMemoryScope<
    access::address_space::ext_intel_global_device_space> {
  static constexpr auto scope = __spv::Scope::Device;
};
template <> struct GetSpirvMemoryScope<access::address_space::local_space> {
  static constexpr auto scope = __spv::Scope::Workgroup;
};

} // namespace detail
} // namespace _V1
} // namespace sycl

#ifndef __SYCL_DEVICE_ONLY__
// host implementation of SYCL atomics
namespace sycl {
inline namespace _V1 {
namespace detail {
// Translate sycl::memory_order or __spv::MemorySemanticsMask::Flag
// into std::memory_order
// Only relaxed memory semantics are supported currently
inline std::memory_order getStdMemoryOrder(__spv::MemorySemanticsMask::Flag) {
  return std::memory_order_relaxed;
}
} // namespace detail
} // namespace _V1
} // namespace sycl

// std::atomic version of atomic SPIR-V builtins

template <typename T>
void __spirv_AtomicStore(std::atomic<T> *Ptr, __spv::Scope::Flag,
                         __spv::MemorySemanticsMask::Flag MS, T V) {
  Ptr->store(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
T __spirv_AtomicLoad(const std::atomic<T> *Ptr, __spv::Scope::Flag,
                     __spv::MemorySemanticsMask::Flag MS) {
  return Ptr->load(::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
T __spirv_AtomicExchange(std::atomic<T> *Ptr, __spv::Scope::Flag,
                         __spv::MemorySemanticsMask::Flag MS, T V) {
  return Ptr->exchange(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
extern T __spirv_AtomicIAdd(std::atomic<T> *Ptr, __spv::Scope::Flag,
                            __spv::MemorySemanticsMask::Flag MS, T V) {
  return Ptr->fetch_add(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
extern T __spirv_AtomicISub(std::atomic<T> *Ptr, __spv::Scope::Flag,
                            __spv::MemorySemanticsMask::Flag MS, T V) {
  return Ptr->fetch_sub(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
extern T __spirv_AtomicAnd(std::atomic<T> *Ptr, __spv::Scope::Flag,
                           __spv::MemorySemanticsMask::Flag MS, T V) {
  return Ptr->fetch_and(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
extern T __spirv_AtomicOr(std::atomic<T> *Ptr, __spv::Scope::Flag,
                          __spv::MemorySemanticsMask::Flag MS, T V) {
  return Ptr->fetch_or(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
extern T __spirv_AtomicXor(std::atomic<T> *Ptr, __spv::Scope::Flag,
                           __spv::MemorySemanticsMask::Flag MS, T V) {
  return Ptr->fetch_xor(V, ::sycl::detail::getStdMemoryOrder(MS));
}

template <typename T>
extern T __spirv_AtomicMin(std::atomic<T> *Ptr, __spv::Scope::Flag,
                           __spv::MemorySemanticsMask::Flag MS, T V) {
  std::memory_order MemoryOrder = ::sycl::detail::getStdMemoryOrder(MS);
  T Val = Ptr->load(MemoryOrder);
  while (V < Val) {
    if (Ptr->compare_exchange_strong(Val, V, MemoryOrder, MemoryOrder))
      break;
    Val = Ptr->load(MemoryOrder);
  }
  return Val;
}

template <typename T>
extern T __spirv_AtomicMax(std::atomic<T> *Ptr, __spv::Scope::Flag,
                           __spv::MemorySemanticsMask::Flag MS, T V) {
  std::memory_order MemoryOrder = ::sycl::detail::getStdMemoryOrder(MS);
  T Val = Ptr->load(MemoryOrder);
  while (V > Val) {
    if (Ptr->compare_exchange_strong(Val, V, MemoryOrder, MemoryOrder))
      break;
    Val = Ptr->load(MemoryOrder);
  }
  return Val;
}

#endif // !defined(__SYCL_DEVICE_ONLY__)

namespace sycl {
inline namespace _V1 {

template <typename T, access::address_space addressSpace =
                          access::address_space::global_space>
class __SYCL2020_DEPRECATED(
    "sycl::atomic is deprecated since SYCL 2020") atomic {
  friend class atomic<T, access::address_space::global_space>;
  static_assert(detail::IsValidAtomicType<T>::value,
                "Invalid SYCL atomic type. Valid types are: int, "
                "unsigned int, long, unsigned long, long long, unsigned "
                "long long, float");
  static_assert(detail::IsValidAtomicAddressSpace<addressSpace>::value,
                "Invalid SYCL atomic address_space. Valid address spaces are: "
                "global_space, local_space, ext_intel_global_device_space");
  static constexpr auto SpirvScope =
      detail::GetSpirvMemoryScope<addressSpace>::scope;

  template <typename pointerT, access::decorated IsDecorated>
  static auto
  GetDecoratedPtr(multi_ptr<pointerT, addressSpace, IsDecorated> ptr) {
    if constexpr (IsDecorated == access::decorated::legacy)
      return ptr.get();
    else
      return ptr.get_decorated();
  }

public:
  template <typename pointerT, access::decorated IsDecorated>
#ifdef __SYCL_DEVICE_ONLY__
  atomic(multi_ptr<pointerT, addressSpace, IsDecorated> ptr)
      : Ptr(GetDecoratedPtr(ptr))
#else
  atomic(multi_ptr<pointerT, addressSpace, IsDecorated> ptr)
      : Ptr(reinterpret_cast<std::atomic<T> *>(ptr.get()))
#endif
  {
    static_assert(sizeof(T) == sizeof(pointerT),
                  "T and pointerT must be same size");
  }

#ifdef __ENABLE_USM_ADDR_SPACE__
  // Create atomic in global_space with one from ext_intel_global_device_space
  template <access::address_space _Space = addressSpace,
            typename = typename std::enable_if_t<
                _Space == addressSpace &&
                addressSpace == access::address_space::global_space>>
  atomic(const atomic<T, access::address_space::ext_intel_global_device_space>
             &RHS) {
    Ptr = RHS.Ptr;
  }

  template <access::address_space _Space = addressSpace,
            typename = typename std::enable_if_t<
                _Space == addressSpace &&
                addressSpace == access::address_space::global_space>>
  atomic(
      atomic<T, access::address_space::ext_intel_global_device_space> &&RHS) {
    Ptr = RHS.Ptr;
  }
#endif // __ENABLE_USM_ADDR_SPACE__

  void store(T Operand, memory_order Order = memory_order::relaxed) {
    __spirv_AtomicStore(Ptr, SpirvScope,
                        detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

#ifdef __SYCL_DEVICE_ONLY__
  template <typename T2 = T>
  std::enable_if_t<!std::is_same<cl_float, T2>::value, T>
  load(memory_order Order = memory_order::relaxed) const {
    return __spirv_AtomicLoad(Ptr, SpirvScope,
                              detail::getSPIRVMemorySemanticsMask(Order));
  }
  template <typename T2 = T>
  std::enable_if_t<std::is_same<cl_float, T2>::value, T>
  load(memory_order Order = memory_order::relaxed) const {
    auto *TmpPtr = reinterpret_cast<typename multi_ptr<
        cl_int, addressSpace, access::decorated::yes>::pointer>(Ptr);
    cl_int TmpVal = __spirv_AtomicLoad(
        TmpPtr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order));
    cl_float ResVal = sycl::bit_cast<cl_float>(TmpVal);
    return ResVal;
  }
#else
  T load(memory_order Order = memory_order::relaxed) const {
    return __spirv_AtomicLoad(Ptr, SpirvScope,
                              detail::getSPIRVMemorySemanticsMask(Order));
  }
#endif

  T exchange(T Operand, memory_order Order = memory_order::relaxed) {
    return __spirv_AtomicExchange(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  bool
  compare_exchange_strong(T &Expected, T Desired,
                          memory_order SuccessOrder = memory_order::relaxed,
                          memory_order FailOrder = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
#ifdef __SYCL_DEVICE_ONLY__
    T Value = __spirv_AtomicCompareExchange(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(SuccessOrder),
        detail::getSPIRVMemorySemanticsMask(FailOrder), Desired, Expected);

    if (Value == Expected)
      return true;

    Expected = Value;
    return false;
#else
    return Ptr->compare_exchange_strong(Expected, Desired,
                                        detail::getStdMemoryOrder(SuccessOrder),
                                        detail::getStdMemoryOrder(FailOrder));
#endif
  }

  T fetch_add(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicIAdd(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  T fetch_sub(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicISub(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  T fetch_and(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicAnd(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  T fetch_or(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicOr(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  T fetch_xor(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicXor(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  T fetch_min(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicMin(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

  T fetch_max(T Operand, memory_order Order = memory_order::relaxed) {
    __SYCL_STATIC_ASSERT_NOT_FLOAT(T);
    return __spirv_AtomicMax(
        Ptr, SpirvScope, detail::getSPIRVMemorySemanticsMask(Order), Operand);
  }

private:
#ifdef __SYCL_DEVICE_ONLY__
  typename detail::DecoratedType<T, addressSpace>::type *Ptr;
#else
  std::atomic<T> *Ptr;
#endif
};

template <typename T, access::address_space addressSpace>
void atomic_store(atomic<T, addressSpace> Object, T Operand,
                  memory_order MemoryOrder = memory_order::relaxed) {
  Object.store(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_load(atomic<T, addressSpace> Object,
              memory_order MemoryOrder = memory_order::relaxed) {
  return Object.load(MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_exchange(atomic<T, addressSpace> Object, T Operand,
                  memory_order MemoryOrder = memory_order::relaxed) {
  return Object.exchange(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
bool atomic_compare_exchange_strong(
    atomic<T, addressSpace> Object, T &Expected, T Desired,
    memory_order SuccessOrder = memory_order::relaxed,
    memory_order FailOrder = memory_order::relaxed) {
  return Object.compare_exchange_strong(Expected, Desired, SuccessOrder,
                                        FailOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_add(atomic<T, addressSpace> Object, T Operand,
                   memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_add(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_sub(atomic<T, addressSpace> Object, T Operand,
                   memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_sub(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_and(atomic<T, addressSpace> Object, T Operand,
                   memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_and(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_or(atomic<T, addressSpace> Object, T Operand,
                  memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_or(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_xor(atomic<T, addressSpace> Object, T Operand,
                   memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_xor(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_min(atomic<T, addressSpace> Object, T Operand,
                   memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_min(Operand, MemoryOrder);
}

template <typename T, access::address_space addressSpace>
T atomic_fetch_max(atomic<T, addressSpace> Object, T Operand,
                   memory_order MemoryOrder = memory_order::relaxed) {
  return Object.fetch_max(Operand, MemoryOrder);
}

} // namespace _V1
} // namespace sycl

#undef __SYCL_STATIC_ASSERT_NOT_FLOAT
