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

#pragma once

#include <cub/config.cuh>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <cub/agent/agent_reduce.cuh>
#include <cub/device/dispatch/tuning/common.cuh>
#include <cub/util_arch.cuh>

#include <cuda/__device/compute_capability.h>

CUB_NAMESPACE_BEGIN

namespace detail::rfa
{
struct reduce_policy
{
  int threads_per_block;
  int items_per_thread;
  BlockReduceAlgorithm block_algorithm;
};

struct single_tile_policy
{
  int threads_per_block;
  int items_per_thread;
  BlockReduceAlgorithm block_algorithm;
};

struct rfa_policy
{
  reduce_policy reduce;
  single_tile_policy single_tile;

#if _CCCL_HOSTED()
  friend ::std::ostream& operator<<(::std::ostream& os, const rfa_policy& policy)
  {
    return os
        << "rfa_policy { .reduce = { .threads_per_block = " << policy.reduce.threads_per_block
        << ", .items_per_thread = " << policy.reduce.items_per_thread
        << " }, .single_tile = { .threads_per_block = " << policy.single_tile.threads_per_block
        << ", .items_per_thread = " << policy.single_tile.items_per_thread << " } }";
  }
#endif // _CCCL_HOSTED()
};

struct policy_selector
{
  type_t accum_t;
  int accum_size;

  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> rfa_policy
  {
    if (cc >= ::cuda::compute_capability{9, 0})
    {
      // only tuned for float, fall through for other types
      if (accum_t == type_t::float32)
      {
        // ipt_13.tpb_224  1.107188  1.009709  1.097114  1.316820
        const auto scaled = scale_mem_bound(224, 13, accum_size);
        return {{scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING},
                {scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING}};
      }
    }

    if (cc >= ::cuda::compute_capability{8, 6})
    {
      // only tuned for float and double, fall through for other types
      if (accum_t == type_t::float32)
      {
        // ipt_6.tpb_224  1.034383  1.000000  1.032097  1.090909
        const auto scaled = scale_mem_bound(224, 6, accum_size);
        return {{scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING},
                {scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING}};
      }
      if (accum_t == type_t::float64)
      {
        // ipt_11.tpb_128 ()  1.232089  1.002124  1.245336  1.582279
        const auto scaled = scale_mem_bound(128, 11, accum_size);
        return {{scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING},
                {scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING}};
      }
    }

    if (cc >= ::cuda::compute_capability{6, 0})
    {
      const auto scaled = scale_mem_bound(256, 16, accum_size);
      return {{scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING},
              {scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING}};
    }

    const auto scaled = scale_mem_bound(256, 20, accum_size);
    return {{scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING},
            {scaled.threads_per_block, scaled.items_per_thread, BLOCK_REDUCE_RAKING}};
  }
};

// stateless version which can be passed to kernels
template <typename AccumT>
struct policy_selector_from_types
{
  [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> rfa_policy
  {
    return policy_selector{classify_type<AccumT>, int{sizeof(AccumT)}}(cc);
  }
};
} // namespace detail::rfa

CUB_NAMESPACE_END
