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

#ifndef KOKKOS_OPENACC_PARALLEL_SCAN_RANGE_HPP
#define KOKKOS_OPENACC_PARALLEL_SCAN_RANGE_HPP

#include <OpenACC/Kokkos_OpenACC.hpp>
#include <OpenACC/Kokkos_OpenACC_FunctorAdapter.hpp>
#include <OpenACC/Kokkos_OpenACC_Macros.hpp>
#include <Kokkos_Parallel.hpp>
#include <Kokkos_BitManipulation.hpp>

#include <type_traits>

// Clacc uses an alternative implementation to work around not-yet-implemented
// OpenACC features: Clacc does not fully support private clauses for
// gang-private variables, and the alternative implementation allocates
// the gang-private arrays on GPU global memory using array expansion,
// instead of using the private clause.
/* clang-format off */
#ifdef KOKKOS_COMPILER_CLANG
#define KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(THREADID) \
  element_values[team_id * 2 * chunk_size + THREADID]
#define KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(THREADID) \
  red_elem_vals[team_id * chunk_size + THREADID]
#define KOKKOS_IMPL_ACC_ELEMENT_VALUES_CLAUSE create(element_values [0:num_elements])
#define KOKKOS_IMPL_ACC_RED_ELEM_VALS_CLAUSE create(red_elem_vals [0:num_red_elems])
#else
#define KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(THREADID) element_values[THREADID]
#define KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(THREADID) red_elem_vals[THREADID]
#define KOKKOS_IMPL_ACC_ELEMENT_VALUES_CLAUSE private(element_values [0:num_elements])
#define KOKKOS_IMPL_ACC_RED_ELEM_VALS_CLAUSE private(red_elem_vals [0:num_red_elems])
#endif
/* clang-format on */

namespace Kokkos::Impl {

template <class Functor, class GivenValueType, class... Traits>
class ParallelScanOpenACCBase {
 protected:
  using Policy = Kokkos::RangePolicy<Traits...>;
  using Analysis =
      Kokkos::Impl::FunctorAnalysis<Kokkos::Impl::FunctorPatternInterface::SCAN,
                                    Policy, Functor, GivenValueType>;
  using PointerType = typename Analysis::pointer_type;
  using ValueType   = typename Analysis::value_type;
  using MemberType  = typename Policy::member_type;
  using IndexType   = typename Policy::index_type;
  Functor m_functor;
  Policy m_policy;
  ValueType* m_result_ptr;
  bool m_result_ptr_device_accessible;
  static constexpr MemberType default_scan_chunk_size = 128;

 public:
  ParallelScanOpenACCBase(Functor const& arg_functor, Policy const& arg_policy,
                          ValueType* arg_result_ptr,
                          bool arg_result_ptr_device_accessible)
      : m_functor(arg_functor),
        m_policy(arg_policy),
        m_result_ptr(arg_result_ptr),
        m_result_ptr_device_accessible(arg_result_ptr_device_accessible) {}

  // This function implements the parallel scan alogithm based on the parallel
  // prefix sum algorithm proposed by Hillis and Steele (doi:10.1145/7902.7903),
  // which offers a shorter span and more parallelism but may not be
  // work-efficient.
  // It takes a Reduce-Scan-Scan approach to parallelize across chunks.
  void OpenACCParallelScanRangePolicy(const IndexType begin,
                                      const IndexType end, IndexType chunk_size,
                                      const int async_arg) const {
    if (chunk_size > 1) {
      if (!Kokkos::has_single_bit<std::make_unsigned_t<IndexType>>(chunk_size))
        Kokkos::abort(
            "RangePolicy blocking granularity must be power of two to be used "
            "with OpenACC parallel_scan()");
    } else {
      chunk_size = default_scan_chunk_size;
    }
    const Kokkos::Experimental::Impl::FunctorAdapter<
        Functor, Policy, Kokkos::Experimental::Impl::RoutineClause::seq>
        functor(m_functor);
    const IndexType N        = end - begin;
    const IndexType n_chunks = (N + chunk_size - 1) / chunk_size;
#ifdef KOKKOS_COMPILER_CLANG
    int const num_elements  = n_chunks * 2 * chunk_size;
    int const num_red_elems = n_chunks * chunk_size;
#else
    int const num_elements  = 2 * chunk_size;
    int const num_red_elems = chunk_size;
#endif
    Kokkos::View<ValueType*, Kokkos::Experimental::OpenACCSpace> chunk_values(
        "Kokkos::OpenACCParallelScan::chunk_values", n_chunks);
    Kokkos::View<ValueType*, Kokkos::Experimental::OpenACCSpace> offset_values(
        "Kokkos::OpenACCParallelScan::offset_values", n_chunks);
    Kokkos::View<ValueType, Kokkos::Experimental::OpenACCSpace> m_result_total(
        "Kokkos::OpenACCParallelScan::m_result_total");
    std::unique_ptr<ValueType[]> element_values_owner(
        new ValueType[num_elements]);
    ValueType* element_values = element_values_owner.get();
    std::unique_ptr<ValueType[]> red_elem_vals_owner(
        new ValueType[num_red_elems]);
    ValueType* red_elem_vals = red_elem_vals_owner.get();
    typename Analysis::Reducer final_reducer(m_functor);

#pragma acc enter data copyin(functor, final_reducer) \
    copyin(chunk_values, offset_values) async(async_arg)

    /* clang-format off */
KOKKOS_IMPL_ACC_PRAGMA(parallel loop gang vector_length(chunk_size) KOKKOS_IMPL_ACC_RED_ELEM_VALS_CLAUSE present(functor, chunk_values, final_reducer) async(async_arg))
    /* clang-format on */
    for (IndexType team_id = 0; team_id < n_chunks; ++team_id) {
#pragma acc loop vector
      for (IndexType thread_id = 0; thread_id < chunk_size; ++thread_id) {
        const IndexType local_offset = team_id * chunk_size + begin;
        const IndexType idx          = local_offset + thread_id;
        ValueType update;
        final_reducer.init(&update);
        if (idx < end) functor(idx, update, false);
        KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(thread_id) = update;
      }
      // chunk-local reduction
      IndexType t_size = chunk_size;
      for (IndexType step_size = (chunk_size >> 1); step_size > 0;
           step_size >>= 1) {
#pragma acc loop vector
        for (IndexType thread_id = 0; thread_id < chunk_size; ++thread_id) {
          if (thread_id < step_size) {
            final_reducer.join(
                &KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(thread_id),
                &KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(thread_id + step_size));
          }
          IndexType check_odd = t_size & 1;
          if (check_odd == 1) {
            if (thread_id == 0) {
              final_reducer.join(&KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(thread_id),
                                 &KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(t_size - 1));
            }
          }
        }
        t_size = step_size;
      }
      // save the reduction result of each chunk
      chunk_values(team_id) = KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS(0);
    }

    // serial exclusive scan of the reduction results of all the chunks
    ValueType tempValue;
#pragma acc parallel loop seq num_gangs(1) num_workers(1) vector_length(1) \
    present(chunk_values, offset_values, final_reducer) async(async_arg)
    for (IndexType team_id = 0; team_id < n_chunks; ++team_id) {
      if (team_id == 0) {
        final_reducer.init(&offset_values(0));
        final_reducer.init(&tempValue);
      } else {
        final_reducer.join(&tempValue, &chunk_values(team_id - 1));
        offset_values(team_id) = tempValue;
      }
    }

    /* clang-format off */
KOKKOS_IMPL_ACC_PRAGMA(parallel loop gang vector_length(chunk_size) KOKKOS_IMPL_ACC_ELEMENT_VALUES_CLAUSE present(functor, offset_values, final_reducer) copyin(m_result_total) async(async_arg))
    /* clang-format on */
    for (IndexType team_id = 0; team_id < n_chunks; ++team_id) {
#pragma acc loop vector
      for (IndexType thread_id = 0; thread_id < chunk_size; ++thread_id) {
        const IndexType local_offset = team_id * chunk_size + begin;
        const IndexType idx          = local_offset + thread_id;
        ValueType update;
        final_reducer.init(&update);
        if (idx < end) functor(idx, update, false);
        KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(thread_id) = update;
      }
      IndexType current_step = 0;
      IndexType next_step    = 1;
      IndexType temp;
      for (IndexType step_size = 1; step_size < chunk_size; step_size *= 2) {
#pragma acc loop vector
        for (IndexType thread_id = 0; thread_id < chunk_size; ++thread_id) {
          if (thread_id < step_size) {
            KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(next_step * chunk_size +
                                            thread_id) =
                KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(current_step * chunk_size +
                                                thread_id);
          } else {
            ValueType localValue = KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(
                current_step * chunk_size + thread_id);
            final_reducer.join(&localValue, &KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(
                                                current_step * chunk_size +
                                                thread_id - step_size));
            KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(next_step * chunk_size +
                                            thread_id) = localValue;
          }
        }
        temp         = current_step;
        current_step = next_step;
        next_step    = temp;
      }
#pragma acc loop vector
      for (IndexType thread_id = 0; thread_id < chunk_size; ++thread_id) {
        const IndexType local_offset = team_id * chunk_size + begin;
        const IndexType idx          = local_offset + thread_id;

        ValueType update;
        final_reducer.init(&update);
        final_reducer.join(&update, &offset_values(team_id));
        // Reconstruct the exclusive prefix from the chunk-local inclusive scan
        // to avoid rereading the previous chunk tail.
        if (thread_id > 0) {
          final_reducer.join(&update,
                             &KOKKOS_IMPL_ACC_ACCESS_ELEMENTS(
                                 current_step * chunk_size + thread_id - 1));
        }

        if (idx < end) functor(idx, update, true);
        if (idx == end - 1) {
          if (m_result_ptr_device_accessible) {
            *m_result_ptr = update;
          } else {
            m_result_total() = update;
          }
        }
      }
    }
    if (!m_result_ptr_device_accessible && m_result_ptr != nullptr) {
      DeepCopy<HostSpace, Kokkos::Experimental::OpenACCSpace,
               Kokkos::Experimental::OpenACC>(m_policy.space(), m_result_ptr,
                                              m_result_total.data(),
                                              sizeof(ValueType));
    }

#pragma acc exit data delete (functor, chunk_values, offset_values, \
                                  final_reducer)async(async_arg)
    acc_wait(async_arg);
  }

  void execute() const {
    const IndexType begin = m_policy.begin();
    const IndexType end   = m_policy.end();
    IndexType chunk_size  = m_policy.chunk_size();

    if (end <= begin) {
      if (!m_result_ptr_device_accessible && m_result_ptr != nullptr) {
        *m_result_ptr = 0;
      }
      return;
    }

    int const async_arg = m_policy.space().acc_async_queue();

    OpenACCParallelScanRangePolicy(begin, end, chunk_size, async_arg);
  }
};

}  // namespace Kokkos::Impl

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

template <class Functor, class... Traits>
class Kokkos::Impl::ParallelScan<Functor, Kokkos::RangePolicy<Traits...>,
                                 Kokkos::Experimental::OpenACC>
    : public ParallelScanOpenACCBase<Functor, void, Traits...> {
  using base_t    = ParallelScanOpenACCBase<Functor, void, Traits...>;
  using IndexType = typename base_t::IndexType;

 public:
  void execute() const {
    const IndexType begin = base_t::m_policy.begin();
    const IndexType end   = base_t::m_policy.end();
    IndexType chunk_size  = base_t::m_policy.chunk_size();

    int const async_arg = base_t::m_policy.space().acc_async_queue();

    base_t::OpenACCParallelScanRangePolicy(begin, end, chunk_size, async_arg);
  }

  ParallelScan(const Functor& arg_functor,
               const typename base_t::Policy& arg_policy)
      : base_t(arg_functor, arg_policy, nullptr, false) {}
};

template <class FunctorType, class ReturnType, class... Traits>
class Kokkos::Impl::ParallelScanWithTotal<
    FunctorType, Kokkos::RangePolicy<Traits...>, ReturnType,
    Kokkos::Experimental::OpenACC>
    : public ParallelScanOpenACCBase<FunctorType, ReturnType, Traits...> {
  using base_t    = ParallelScanOpenACCBase<FunctorType, ReturnType, Traits...>;
  using IndexType = typename base_t::IndexType;

 public:
  void execute() const {
    const IndexType begin = base_t::m_policy.begin();
    const IndexType end   = base_t::m_policy.end();
    IndexType chunk_size  = base_t::m_policy.chunk_size();

    if (end <= begin) {
      if (!base_t::m_result_ptr_device_accessible &&
          base_t::m_result_ptr != nullptr) {
        *base_t::m_result_ptr = 0;
      }
      return;
    }

    int const async_arg = base_t::m_policy.space().acc_async_queue();

    base_t::OpenACCParallelScanRangePolicy(begin, end, chunk_size, async_arg);
  }

  template <class ViewType>
  ParallelScanWithTotal(const FunctorType& arg_functor,
                        const typename base_t::Policy& arg_policy,
                        const ViewType& arg_result_view)
      : base_t(arg_functor, arg_policy, arg_result_view.data(),
               MemorySpaceAccess<Kokkos::Experimental::OpenACCSpace,
                                 typename ViewType::memory_space>::accessible) {
  }
};

#undef KOKKOS_IMPL_ACC_ACCESS_ELEMENTS
#undef KOKKOS_IMPL_ACC_ELEMENT_VALUES_CLAUSE
#undef KOKKOS_IMPL_ACC_ACCESS_RED_ELEMS
#undef KOKKOS_IMPL_ACC_RED_ELEM_VALS_CLAUSE

#endif
