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

#ifndef KOKKOS_TEST_SIMD_SHIFT_OPS_HPP
#define KOKKOS_TEST_SIMD_SHIFT_OPS_HPP

#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.simd;
import kokkos.simd_impl;
#else
#include <Kokkos_SIMD.hpp>
#endif
#include <SIMDTesting_Utilities.hpp>

template <typename Abi, typename Loader, typename ShiftOp, typename DataType>
inline void host_check_shift_on_one_loader(ShiftOp shift_op,
                                           DataType test_vals[],
                                           DataType shift_by[], std::size_t n) {
  using simd_type             = Kokkos::Experimental::basic_simd<DataType, Abi>;
  constexpr std::size_t width = simd_type::size();
  Loader loader;

  for (std::size_t i = 0; i < n; ++i) {
    simd_type simd_vals;
    bool const loaded_arg = loader.host_load(test_vals, width, simd_vals);
    if (!loaded_arg) {
      continue;
    }

    DataType expected_val[width];
    for (std::size_t lane = 0; lane < width; ++lane) {
      expected_val[lane] =
          shift_op.on_host(simd_vals[lane], static_cast<int>(shift_by[i]));
    }

    simd_type expected_result =
        Kokkos::Experimental::simd_unchecked_load<simd_type>(
            expected_val, Kokkos::Experimental::simd_flag_default);
    simd_type const computed_result =
        shift_op.on_host(simd_vals, static_cast<int>(shift_by[i]));

    host_check_equality(expected_result, computed_result, width);
  }
}

template <typename Abi, typename Loader, typename ShiftOp, typename DataType>
inline void host_check_shift_by_lanes_on_one_loader(
    ShiftOp shift_op, DataType test_vals[],
    Kokkos::Experimental::basic_simd<DataType, Abi>& shift_by) {
  using simd_type             = Kokkos::Experimental::basic_simd<DataType, Abi>;
  constexpr std::size_t width = simd_type::size();
  Loader loader;

  simd_type simd_vals;
  bool const loaded_arg = loader.host_load(test_vals, width, simd_vals);
  ASSERT_TRUE(loaded_arg);

  DataType expected_val[width];
  for (std::size_t lane = 0; lane < width; ++lane) {
    expected_val[lane] =
        shift_op.on_host(simd_vals[lane], static_cast<int>(shift_by[lane]));
  }

  simd_type expected_result =
      Kokkos::Experimental::simd_unchecked_load<simd_type>(
          expected_val, Kokkos::Experimental::simd_flag_default);
  simd_type const computed_result = shift_op.on_host(simd_vals, shift_by);

  host_check_equality(expected_result, computed_result, width);
}

template <typename Abi, typename ShiftOp, typename DataType>
inline void host_check_shift_op_all_loaders(ShiftOp shift_op,
                                            DataType test_vals[],
                                            DataType shift_by[],
                                            std::size_t n) {
  host_check_shift_on_one_loader<Abi, load_element_aligned>(shift_op, test_vals,
                                                            shift_by, n);
  host_check_shift_on_one_loader<Abi, load_masked>(shift_op, test_vals,
                                                   shift_by, n);
  host_check_shift_on_one_loader<Abi, load_as_scalars>(shift_op, test_vals,
                                                       shift_by, n);
  host_check_shift_on_one_loader<Abi, load_vector_aligned>(shift_op, test_vals,
                                                           shift_by, n);

  using simd_type = Kokkos::Experimental::basic_simd<DataType, Abi>;
  simd_type shift_by_lanes =
      Kokkos::Experimental::simd_unchecked_load<simd_type>(
          shift_by, Kokkos::Experimental::simd_flag_default);

  host_check_shift_by_lanes_on_one_loader<Abi, load_element_aligned>(
      shift_op, test_vals, shift_by_lanes);
  host_check_shift_by_lanes_on_one_loader<Abi, load_masked>(shift_op, test_vals,
                                                            shift_by_lanes);
  host_check_shift_by_lanes_on_one_loader<Abi, load_as_scalars>(
      shift_op, test_vals, shift_by_lanes);
  host_check_shift_by_lanes_on_one_loader<Abi, load_vector_aligned>(
      shift_op, test_vals, shift_by_lanes);
}

template <typename Abi, typename DataType>
inline void host_check_shift_ops() {
  if constexpr (is_simd_avail_v<DataType, Abi>) {
    if constexpr (std::is_integral_v<DataType>) {
      using simd_type = Kokkos::Experimental::basic_simd<DataType, Abi>;
      constexpr std::size_t width     = simd_type::size();
      constexpr std::size_t num_cases = 16;
      constexpr size_t alignment =
          Kokkos::Experimental::basic_simd<DataType, Abi>::size() *
          sizeof(DataType);

      DataType max = std::numeric_limits<DataType>::max();

      alignas(alignment) DataType shift_by[num_cases] = {
          0, 1, 3, width / 2, width / 2 + 1, width - 1, width, width + 1,
          0, 1, 3, width / 2, width / 2 + 1, width - 1, width, width + 1};
      alignas(alignment) DataType test_vals[width];
      for (std::size_t i = 0; i < width; ++i) {
        DataType inc = max / width;
        test_vals[i] = i * inc + 1;
      }

      host_check_shift_op_all_loaders<Abi>(shift_right(), test_vals, shift_by,
                                           num_cases);
      host_check_shift_op_all_loaders<Abi>(shift_right_eq(), test_vals,
                                           shift_by, num_cases);
      host_check_shift_op_all_loaders<Abi>(shift_left(), test_vals, shift_by,
                                           num_cases);
      host_check_shift_op_all_loaders<Abi>(shift_left_eq(), test_vals, shift_by,
                                           num_cases);

      if constexpr (std::is_signed_v<DataType>) {
        for (std::size_t i = 0; i < width; ++i) test_vals[i] *= -1;
        host_check_shift_op_all_loaders<Abi>(shift_right(), test_vals, shift_by,
                                             num_cases);
        host_check_shift_op_all_loaders<Abi>(shift_right_eq(), test_vals,
                                             shift_by, num_cases);
      }
    }
  }
}

template <typename Abi, typename... DataTypes>
inline void host_check_shift_ops_all_types(
    Kokkos::Experimental::Impl::data_types<DataTypes...>) {
  (host_check_shift_ops<Abi, DataTypes>(), ...);
}

template <typename... Abis>
inline void host_check_shift_ops_all_abis(
    Kokkos::Experimental::Impl::abi_set<Abis...>) {
  using DataTypes = Kokkos::Experimental::Impl::data_type_set;
  (host_check_shift_ops_all_types<Abis>(DataTypes()), ...);
}

template <typename Abi, typename Loader, typename ShiftOp, typename DataType>
KOKKOS_INLINE_FUNCTION void device_check_shift_on_one_loader(
    ShiftOp shift_op, DataType test_vals[], DataType shift_by[],
    std::size_t n) {
  using simd_type             = Kokkos::Experimental::basic_simd<DataType, Abi>;
  constexpr std::size_t width = simd_type::size();
  Loader loader;

  for (std::size_t i = 0; i < n; ++i) {
    simd_type simd_vals;
    bool const loaded_arg = loader.device_load(test_vals, width, simd_vals);
    if (!loaded_arg) {
      continue;
    }

    simd_type expected_result(KOKKOS_LAMBDA(std::size_t lane) {
      return shift_op.on_device(simd_vals[lane], static_cast<int>(shift_by[i]));
    });

    simd_type const computed_result =
        shift_op.on_device(simd_vals, static_cast<int>(shift_by[i]));
    device_check_equality(expected_result, computed_result, width);
  }
}

template <typename Abi, typename Loader, typename ShiftOp, typename DataType>
KOKKOS_INLINE_FUNCTION void device_check_shift_by_lanes_on_one_loader(
    ShiftOp shift_op, DataType test_vals[],
    Kokkos::Experimental::basic_simd<DataType, Abi>& shift_by) {
  using simd_type             = Kokkos::Experimental::basic_simd<DataType, Abi>;
  constexpr std::size_t width = simd_type::size();
  Loader loader;
  simd_type simd_vals;
  loader.device_load(test_vals, width, simd_vals);

  simd_type expected_result(KOKKOS_LAMBDA(std::size_t lane) {
    return shift_op.on_device(simd_vals[lane],
                              static_cast<int>(shift_by[lane]));
  });

  simd_type const computed_result = shift_op.on_device(simd_vals, shift_by);
  device_check_equality(expected_result, computed_result, width);
}

template <typename Abi, typename ShiftOp, typename DataType>
KOKKOS_INLINE_FUNCTION void device_check_shift_op_all_loaders(
    ShiftOp shift_op, DataType test_vals[], DataType shift_by[],
    std::size_t n) {
  device_check_shift_on_one_loader<Abi, load_element_aligned>(
      shift_op, test_vals, shift_by, n);
  device_check_shift_on_one_loader<Abi, load_masked>(shift_op, test_vals,
                                                     shift_by, n);
  device_check_shift_on_one_loader<Abi, load_as_scalars>(shift_op, test_vals,
                                                         shift_by, n);
  device_check_shift_on_one_loader<Abi, load_vector_aligned>(
      shift_op, test_vals, shift_by, n);

  using simd_type = Kokkos::Experimental::basic_simd<DataType, Abi>;
  simd_type shift_by_lanes =
      Kokkos::Experimental::simd_unchecked_load<simd_type>(
          shift_by, Kokkos::Experimental::simd_flag_default);

  device_check_shift_by_lanes_on_one_loader<Abi, load_element_aligned>(
      shift_op, test_vals, shift_by_lanes);
  device_check_shift_by_lanes_on_one_loader<Abi, load_masked>(
      shift_op, test_vals, shift_by_lanes);
  device_check_shift_by_lanes_on_one_loader<Abi, load_as_scalars>(
      shift_op, test_vals, shift_by_lanes);
  device_check_shift_by_lanes_on_one_loader<Abi, load_vector_aligned>(
      shift_op, test_vals, shift_by_lanes);
}

template <typename Abi, typename DataType>
KOKKOS_INLINE_FUNCTION void device_check_shift_ops() {
  if constexpr (is_type_v<Kokkos::Experimental::basic_simd<DataType, Abi>>) {
    if constexpr (std::is_integral_v<DataType>) {
      using simd_type = Kokkos::Experimental::basic_simd<DataType, Abi>;
      constexpr std::size_t width     = simd_type::size();
      constexpr std::size_t num_cases = 16;

      DataType max = Kokkos::reduction_identity<DataType>::max();

      DataType shift_by[num_cases] = {
          0, 1, 3, width / 2, width / 2 + 1, width - 1, width, width + 1,
          0, 1, 3, width / 2, width / 2 + 1, width - 1, width, width + 1};
      DataType test_vals[width];

      for (std::size_t i = 0; i < width; ++i) {
        DataType inc = max / width;
        test_vals[i] = i * inc + 1;
      }

      device_check_shift_op_all_loaders<Abi>(shift_right(), test_vals, shift_by,
                                             num_cases);
      device_check_shift_op_all_loaders<Abi>(shift_right_eq(), test_vals,
                                             shift_by, num_cases);
      device_check_shift_op_all_loaders<Abi>(shift_left(), test_vals, shift_by,
                                             num_cases);
      device_check_shift_op_all_loaders<Abi>(shift_left_eq(), test_vals,
                                             shift_by, num_cases);

      if constexpr (std::is_signed_v<DataType>) {
        for (std::size_t i = 0; i < width; ++i) test_vals[i] *= -1;
        device_check_shift_op_all_loaders<Abi>(shift_right(), test_vals,
                                               shift_by, num_cases);
        device_check_shift_op_all_loaders<Abi>(shift_right_eq(), test_vals,
                                               shift_by, num_cases);
      }
    }
  }
}

template <typename Abi, typename... DataTypes>
KOKKOS_INLINE_FUNCTION void device_check_shift_ops_all_types(
    Kokkos::Experimental::Impl::data_types<DataTypes...>) {
  (device_check_shift_ops<Abi, DataTypes>(), ...);
}

template <typename... Abis>
KOKKOS_INLINE_FUNCTION void device_check_shift_ops_all_abis(
    Kokkos::Experimental::Impl::abi_set<Abis...>) {
  using DataTypes = Kokkos::Experimental::Impl::data_type_set;
  (device_check_shift_ops_all_types<Abis>(DataTypes()), ...);
}

class simd_device_shift_ops_functor {
 public:
  KOKKOS_INLINE_FUNCTION void operator()(int) const {
    device_check_shift_ops_all_abis(
        Kokkos::Experimental::Impl::device_abi_set());
  }
};

TEST(simd, host_shift_ops) {
  host_check_shift_ops_all_abis(Kokkos::Experimental::Impl::host_abi_set());
}

TEST(simd, device_shift_ops) {
  Kokkos::parallel_for(Kokkos::RangePolicy<Kokkos::IndexType<int>>(0, 1),
                       simd_device_shift_ops_functor());
}

#endif
