// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause

#pragma once

// keep checks at the top so compilation of discarded variants fails really fast
#include <cub/device/dispatch/dispatch_transform.cuh>
#if !TUNE_BASE
#  if _CCCL_PP_COUNT(__CUDA_ARCH_LIST__) != 1
#    error "When tuning, this benchmark does not support being compiled for multiple architectures"
#  endif
#  if TUNE_ALGORITHM == 3
#    if (__CUDA_ARCH_LIST__) < 900
#      error "Cannot compile algorithm 3 (ublkcp) below sm90"
#    endif
#  endif // TUNE_ALGORITHM == 3
#endif // !TUNE_BASE

#include <cub/util_namespace.cuh>

#include <cuda/__numeric/narrow.h>
#include <cuda/std/cstdint>
#include <cuda/std/type_traits>

#include <stdexcept>

#include <nvbench_helper.cuh>

#if !TUNE_BASE
struct policy_selector
{
  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability cc) const -> cub::TransformPolicy
  {
    const int min_bytes_in_flight = cub::detail::transform::cc_to_min_bytes_in_flight(cc) + TUNE_BIF_BIAS;
#  if TUNE_ALGORITHM == 0 || TUNE_ALGORITHM == 1
    // setup prefetch, since it's either used directly or the fallback to vectorized
    auto algorithm                = cub::TransformAlgorithm::prefetch;
    auto pref_policy              = cub::TransformPrefetchPolicy{};
    pref_policy.threads_per_block = TUNE_THREADS;
    pref_policy.unroll_factor     = TUNE_UNROLL_FACTOR;
#    ifdef TUNE_PREFETCH_MULT
    pref_policy.prefetch_byte_stride = 32 * TUNE_PREFETCH_MULT;
#    endif //  TUNE_PREFETCH_MULT
#    ifdef TUNE_ITEMS_PER_THREAD_NO_INPUT
    pref_policy.items_per_thread_no_input = TUNE_ITEMS_PER_THREAD_NO_INPUT;
#    endif // TUNE_ITEMS_PER_THREAD_NO_INPUT

    // setup vectorized if requested
    auto vec_policy = cub::TransformVectorizedPolicy{};
#    if TUNE_ALGORITHM == 1
    algorithm                    = cub::TransformAlgorithm::vectorized;
    vec_policy.threads_per_block = TUNE_THREADS;
    vec_policy.vec_size          = (1 << TUNE_VEC_SIZE_POW2);
    vec_policy.items_per_thread  = vec_policy.vec_size * TUNE_UNROLL_FACTOR;
#    endif
    return {min_bytes_in_flight, algorithm, pref_policy, vec_policy, {}};
#  elif TUNE_ALGORITHM == 2
    constexpr auto algorithm = cub::TransformAlgorithm::ldgsts;
    auto policy              = cub::TransformAsyncCopyPolicy{};
    policy.threads_per_block = TUNE_THREADS;
    policy.unroll_factor     = TUNE_UNROLL_FACTOR;
    return {min_bytes_in_flight, algorithm, {}, {}, policy};
#  elif TUNE_ALGORITHM == 3
    constexpr auto algorithm = cub::TransformAlgorithm::ublkcp;
    auto policy              = cub::TransformAsyncCopyPolicy{};
    policy.threads_per_block = TUNE_THREADS;
    policy.unroll_factor     = TUNE_UNROLL_FACTOR;
    return {min_bytes_in_flight, algorithm, {}, {}, policy};
#  else // TUNE_ALGORITHM
#    error Policy hub does not yet implement the specified value for algorithm
#  endif // TUNE_ALGORITHM
  }
};
#endif // !TUNE_BASE

template <typename... RandomAccessIteratorsIn, typename RandomAccessIteratorOut, typename TransformOp>
void bench_transform(nvbench::state& state,
                     cuda::std::tuple<RandomAccessIteratorsIn...> inputs,
                     RandomAccessIteratorOut output,
                     ::cuda::std::int64_t num_items,
                     TransformOp transform_op)
{
  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](const nvbench::launch& launch) {
    cub::DeviceTransform::Transform(
      inputs,
      output,
      num_items,
      transform_op,
      cuda::std::execution::env{::cuda::stream_ref{launch.get_stream().get_stream()}
#if !TUNE_BASE
                                ,
                                cuda::execution::tune(policy_selector{})
#endif // !TUNE_BASE
      });
  });
}