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

#ifndef KOKKOS_SIMD_HPP
#define KOKKOS_SIMD_HPP

#include <Kokkos_SIMD_Common.hpp>
#include <Kokkos_SIMD_Base.hpp>
#include <Kokkos_SIMD_Scalar.hpp>
#include <Kokkos_Macros.hpp>

#include <climits>
#include <cstdint>

#if defined(KOKKOS_ARCH_AVX) && !defined(__AVX__)
#error "__AVX__ must be defined for KOKKOS_ARCH_AVX"
#endif

#if defined(KOKKOS_ARCH_AVX2)
#if !defined(__AVX2__)
#error "__AVX2__ must be defined for KOKKOS_ARCH_AVX2"
#endif
#include <Kokkos_SIMD_AVX2.hpp>
#endif

#if defined(KOKKOS_ARCH_AVX512XEON)
#if !defined(__AVX512F__)
#error "__AVX512F__ must be defined for KOKKOS_ARCH_AVX512XEON"
#endif
#include <Kokkos_SIMD_AVX512.hpp>
#endif

#if defined(KOKKOS_ARCH_ARM_SVE)
#if !defined(__ARM_FEATURE_SVE_BITS) || !defined(__ARM_NEON)
#error \
    "Both __ARM_FEATURE_SVE_BITS and __ARM_NEON must be definded for KOKKOS_ARCH_ARM_SVE"
#endif
#include <Kokkos_SIMD_SVE.hpp>
#endif

#if defined(KOKKOS_ARCH_ARM_NEON)
#if !defined(__ARM_NEON)
#error "__ARM_NEON must be defined for KOKKOS_ARCH_ARM_NEON"
#endif
#include <Kokkos_SIMD_NEON.hpp>
#endif

#include <Kokkos_SIMD_Common_Math.hpp>

namespace Kokkos {
namespace Experimental {

namespace simd_abi {

namespace Impl {

#if defined(KOKKOS_ARCH_AVX512XEON)
template <class T>
using host_fixed_native = avx512_fixed_size<8>;
template <typename T, Experimental::Impl::simd_size_t N>
using host_native_abi =
    std::conditional_t<N == 0, avx512_fixed_size<512 / (CHAR_BIT * sizeof(T))>,
                       avx512_fixed_size<N>>;

#elif defined(KOKKOS_ARCH_AVX2)
template <class T>
using host_fixed_native = avx2_fixed_size<4>;
template <typename T, Experimental::Impl::simd_size_t N>
using host_native_abi =
    std::conditional_t<N == 0, avx2_fixed_size<256 / (CHAR_BIT * sizeof(T))>,
                       avx2_fixed_size<N>>;

#elif defined(KOKKOS_ARCH_ARM_SVE)
template <class T>
using host_fixed_native =
    sve_fixed_size<(__ARM_FEATURE_SVE_BITS / (CHAR_BIT * sizeof(T)))>;
template <typename T, Experimental::Impl::simd_size_t N>
using host_native_abi =
    std::conditional_t<N == 0, host_fixed_native<T>, sve_fixed_size<N>>;

#elif defined(KOKKOS_ARCH_ARM_NEON)
template <class T>
using host_fixed_native = neon_fixed_size<2>;
template <typename T, Experimental::Impl::simd_size_t N>
using host_native_abi =
    std::conditional_t<N == 0, neon_fixed_size<128 / (CHAR_BIT * sizeof(T))>,
                       neon_fixed_size<N>>;

#else
template <class T>
using host_fixed_native = scalar;
template <typename T, Experimental::Impl::simd_size_t N>
using host_native_abi = scalar;
#endif

template <class S>
struct ForSpace;

#ifdef KOKKOS_ENABLE_SERIAL
template <>
struct ForSpace<Kokkos::Serial> {
  template <class T>
  using type = host_fixed_native<T>;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = host_native_abi<T, N>;
};
#endif

#ifdef KOKKOS_ENABLE_CUDA
template <>
struct ForSpace<Kokkos::Cuda> {
  template <class T>
  using type = scalar;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = scalar;
};
#endif

#ifdef KOKKOS_ENABLE_THREADS
template <>
struct ForSpace<Kokkos::Threads> {
  template <class T>
  using type = host_fixed_native<T>;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = host_native_abi<T, N>;
};
#endif

#ifdef KOKKOS_ENABLE_HPX
template <>
struct ForSpace<Kokkos::Experimental::HPX> {
  template <class T>
  using type = scalar;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = scalar;
};
#endif

#ifdef KOKKOS_ENABLE_OPENMP
template <>
struct ForSpace<Kokkos::OpenMP> {
  template <class T>
  using type = host_fixed_native<T>;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = host_native_abi<T, N>;
};
#endif

#ifdef KOKKOS_ENABLE_OPENACC
template <>
struct ForSpace<Kokkos::Experimental::OpenACC> {
  template <class T>
  using type = scalar;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = scalar;
};
#endif

#ifdef KOKKOS_ENABLE_HIP
template <>
struct ForSpace<Kokkos::HIP> {
  template <class T>
  using type = scalar;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = scalar;
};
#endif

#ifdef KOKKOS_ENABLE_SYCL
template <>
struct ForSpace<Kokkos::SYCL> {
  template <class T>
  using type = scalar;

  template <typename T, Experimental::Impl::simd_size_t N>
  using simd_abi = scalar;
};
#endif

template <class T, class Space = Kokkos::DefaultExecutionSpace>
using native_fixed_abi = typename ForSpace<Space>::template type<T>;

template <typename T, Experimental::Impl::simd_size_t N,
          class Space = Kokkos::DefaultExecutionSpace>
using native_abi = typename ForSpace<Space>::template simd_abi<T, N>;

}  // namespace Impl

}  // namespace simd_abi

template <class T, Impl::simd_size_t N = 0>
using simd = basic_simd<T, simd_abi::Impl::native_abi<T, N>>;

template <class T, Impl::simd_size_t N = 0>
using simd_mask = basic_simd_mask<T, simd_abi::Impl::native_abi<T, N>>;

template <typename T, typename... Flags>
  requires Impl::NonScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION
    basic_simd<T, simd_abi::Impl::host_fixed_native<T>>
    simd_unchecked_load(const T* ptr,
                        simd_flags<Flags...> flag = simd_flag_default) {
  return simd_unchecked_load<
      basic_simd<T, simd_abi::Impl::host_fixed_native<T>>>(ptr, flag);
}

template <typename T, typename... Flags>
  requires Impl::ScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_FORCEINLINE_FUNCTION constexpr basic_simd<T, simd_abi::scalar>
simd_unchecked_load(const T* ptr,
                    simd_flags<Flags...> flag = simd_flag_default) {
  return simd_unchecked_load<basic_simd<T, simd_abi::scalar>>(ptr, flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::NonScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION auto unchecked_gather_from(
    R&& in, const I& indices, simd_flags<Flags...> flag = simd_flag_default) {
  return unchecked_gather_from<
      basic_simd<T, simd_abi::Impl::host_fixed_native<T>>>(in, indices, flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::ScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_FORCEINLINE_FUNCTION auto unchecked_gather_from(
    R&& in, const I& indices, simd_flags<Flags...> flag = simd_flag_default) {
  return unchecked_gather_from<basic_simd<T, simd_abi::scalar>>(in, indices,
                                                                flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::NonScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION auto unchecked_gather_from(
    R&& in, const typename I::mask_type& mask, const I& indices,
    simd_flags<Flags...> flag = simd_flag_default) {
  return unchecked_gather_from<
      basic_simd<T, simd_abi::Impl::host_fixed_native<T>>>(in, mask, indices,
                                                           flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::ScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_FORCEINLINE_FUNCTION auto unchecked_gather_from(
    R&& in, const typename I::mask_type& mask, const I& indices,
    simd_flags<Flags...> flag = simd_flag_default) {
  return unchecked_gather_from<basic_simd<T, simd_abi::scalar>>(in, mask,
                                                                indices, flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::NonScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION auto partial_gather_from(
    R&& in, const I& indices, simd_flags<Flags...> flag = simd_flag_default) {
  return partial_gather_from<
      basic_simd<T, simd_abi::Impl::host_fixed_native<T>>>(in, indices, flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::ScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_FORCEINLINE_FUNCTION auto partial_gather_from(
    R&& in, const I& indices, simd_flags<Flags...> flag = simd_flag_default) {
  return partial_gather_from<basic_simd<T, simd_abi::scalar>>(in, indices,
                                                              flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::NonScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_IMPL_HOST_FORCEINLINE_FUNCTION auto partial_gather_from(
    R&& in, const typename I::mask_type& mask, const I& indices,
    simd_flags<Flags...> flag = simd_flag_default) {
  return partial_gather_from<
      basic_simd<T, simd_abi::Impl::host_fixed_native<T>>>(in, mask, indices,
                                                           flag);
}

template <Impl::Ranges::contiguous_range R, Impl::SimdIntegral I,
          typename... Flags, typename T = Impl::Ranges::range_value_t<R>>
  requires Impl::Ranges::sized_range<R> &&
           Impl::ScalarAbi<simd_abi::Impl::host_fixed_native<T>>
KOKKOS_FORCEINLINE_FUNCTION auto partial_gather_from(
    R&& in, const typename I::mask_type& mask, const I& indices,
    simd_flags<Flags...> flag = simd_flag_default) {
  return partial_gather_from<basic_simd<T, simd_abi::scalar>>(in, mask, indices,
                                                              flag);
}

namespace Impl {

template <class... Abis>
class abi_set {};

template <typename... Ts>
class data_types {};

#if defined(KOKKOS_ARCH_AVX512XEON)
using host_abi_set  = abi_set<simd_abi::scalar, simd_abi::avx512_fixed_size<8>,
                             simd_abi::avx512_fixed_size<16>>;
using data_type_set = data_types<std::int32_t, std::uint32_t, std::int64_t,
                                 std::uint64_t, double, float>;
#elif defined(KOKKOS_ARCH_AVX2)
using host_abi_set = abi_set<simd_abi::scalar, simd_abi::avx2_fixed_size<4>,
                             simd_abi::avx2_fixed_size<8>>;
using data_type_set =
    data_types<std::int32_t, std::int64_t, std::uint64_t, double, float>;
#elif defined(KOKKOS_ARCH_ARM_SVE)
using host_abi_set =
    abi_set<simd_abi::scalar, simd_abi::sve_fixed_size<2>,
            simd_abi::sve_fixed_size<4>, simd_abi::sve_fixed_size<8>,
            simd_abi::sve_fixed_size<16>>;
using data_type_set = data_types<std::int32_t, std::uint32_t, std::int64_t,
                                 std::uint64_t, double, float>;
#elif defined(KOKKOS_ARCH_ARM_NEON)
using host_abi_set  = abi_set<simd_abi::scalar, simd_abi::neon_fixed_size<2>,
                             simd_abi::neon_fixed_size<4>>;
using data_type_set = data_types<std::int32_t, std::uint32_t, std::int64_t,
                                 std::uint64_t, double, float>;
#else
using host_abi_set    = abi_set<simd_abi::scalar>;
using data_type_set   = data_types<std::int32_t, std::uint32_t, std::int64_t,
                                 std::uint64_t, double, float>;
#endif

using device_abi_set = abi_set<simd_abi::scalar>;

}  // namespace Impl

}  // namespace Experimental
}  // namespace Kokkos

#endif