// SPDX-FileCopyrightText: Copyright (c) 2011-2021, 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_adjacent_difference.cuh>
#include <cub/detail/launcher/cuda_runtime.cuh>
#include <cub/detail/type_traits.cuh>
#include <cub/device/dispatch/dispatch_common.cuh>
#include <cub/device/dispatch/tuning/tuning_adjacent_difference.cuh>
#include <cub/util_arch.cuh>
#include <cub/util_debug.cuh>
#include <cub/util_device.cuh>
#include <cub/util_math.cuh>
#include <cub/util_namespace.cuh>

#include <thrust/system/cuda/detail/core/triple_chevron_launch.h>

#include <cuda/__cmath/ceil_div.h>
#include <cuda/__device/compute_capability.h>
#include <cuda/std/__functional/invoke.h>
#include <cuda/std/__host_stdlib/sstream>
#include <cuda/std/__type_traits/is_empty.h>

CUB_NAMESPACE_BEGIN

namespace detail::adjacent_difference
{
template <typename AgentDifferenceInitT, typename InputIteratorT, typename InputT, typename OffsetT>
_CCCL_KERNEL_ATTRIBUTES void DeviceAdjacentDifferenceInitKernel(
  const InputIteratorT first, InputT* const result, const OffsetT num_tiles, const int items_per_tile)
{
  const int tile_idx = static_cast<int>(blockIdx.x * blockDim.x + threadIdx.x);
  AgentDifferenceInitT::Process(tile_idx, first, result, num_tiles, items_per_tile);
}

template <typename PolicySelector,
          typename InputIteratorT,
          typename OutputIteratorT,
          typename DifferenceOpT,
          typename OffsetT,
          typename InputT,
          bool MayAlias,
          bool ReadLeft>
_CCCL_KERNEL_ATTRIBUTES void DeviceAdjacentDifferenceDifferenceKernel(
  const InputIteratorT input,
  InputT* const first_tile_previous,
  const OutputIteratorT result,
  DifferenceOpT difference_op,
  const OffsetT num_items)
{
  static_assert(::cuda::std::is_empty_v<PolicySelector>);
  static constexpr adjacent_difference_policy policy = current_policy<PolicySelector>();
  using AdjacentDifferencePolicyT =
    AgentAdjacentDifferencePolicy<policy.threads_per_block,
                                  policy.items_per_thread,
                                  policy.load_algorithm,
                                  policy.load_modifier,
                                  policy.store_algorithm>;

  // It is OK to introspect the return type or parameter types of the
  // `operator()` function of `__device__` extended lambda within device code.
  using OutputT = ::cuda::std::invoke_result_t<DifferenceOpT, InputT, InputT>;

  using Agent =
    AgentDifference<AdjacentDifferencePolicyT,
                    InputIteratorT,
                    OutputIteratorT,
                    DifferenceOpT,
                    OffsetT,
                    InputT,
                    OutputT,
                    MayAlias,
                    ReadLeft>;

  __shared__ typename Agent::TempStorage storage;

  Agent agent(storage, input, first_tile_previous, result, difference_op, num_items);

  int tile_idx      = static_cast<int>(blockIdx.x);
  OffsetT tile_base = static_cast<OffsetT>(tile_idx) * AdjacentDifferencePolicyT::ITEMS_PER_TILE;

  agent.Process(tile_idx, tile_base);
}

template <typename PolicyHub>
struct policy_selector_from_hub
{
  // this is only called in device code, so we can ignore the cc parameter
  _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> adjacent_difference_policy
  {
    using p = typename PolicyHub::MaxPolicy::ActivePolicy::AdjacentDifferencePolicy;
    return adjacent_difference_policy{
      p::BLOCK_THREADS, p::ITEMS_PER_THREAD, p::LOAD_ALGORITHM, p::LOAD_MODIFIER, p::STORE_ALGORITHM};
  }
};
} // namespace detail::adjacent_difference

enum class ReadOption
{
  Left,
  Right
};

// TODO(bgruber): remove in CCL 4.0
template <typename InputIteratorT,
          typename OutputIteratorT,
          typename DifferenceOpT,
          typename OffsetT,
          MayAlias AliasOpt,
          ReadOption ReadOpt,
          typename PolicyHub = detail::adjacent_difference::policy_hub<InputIteratorT, AliasOpt == MayAlias::Yes>>
struct DispatchAdjacentDifference
{
  using InputT = detail::it_value_t<InputIteratorT>;

  void* d_temp_storage;
  size_t& temp_storage_bytes;
  InputIteratorT d_input;
  OutputIteratorT d_output;
  OffsetT num_items;
  DifferenceOpT difference_op;
  cudaStream_t stream;

  CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchAdjacentDifference(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    InputIteratorT d_input,
    OutputIteratorT d_output,
    OffsetT num_items,
    DifferenceOpT difference_op,
    cudaStream_t stream)
      : d_temp_storage(d_temp_storage)
      , temp_storage_bytes(temp_storage_bytes)
      , d_input(d_input)
      , d_output(d_output)
      , num_items(num_items)
      , difference_op(difference_op)
      , stream(stream)
  {}

  /// Invocation
  template <typename ActivePolicyT>
  CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke()
  {
    using AdjacentDifferencePolicyT = typename ActivePolicyT::AdjacentDifferencePolicy;

    cudaError error = cudaSuccess;

    do
    {
      constexpr int tile_size = AdjacentDifferencePolicyT::ITEMS_PER_TILE;
      const int num_tiles     = static_cast<int>(::cuda::ceil_div(num_items, tile_size));

      size_t first_tile_previous_size = (AliasOpt == MayAlias::Yes) * num_tiles * sizeof(InputT);

      void* allocations[1]       = {nullptr};
      size_t allocation_sizes[1] = {(AliasOpt == MayAlias::Yes) * first_tile_previous_size};

      error = CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes));

      if (cudaSuccess != error)
      {
        break;
      }

      if (d_temp_storage == nullptr)
      {
        // Return if the caller is simply requesting the size of the storage
        // allocation

        if (temp_storage_bytes == 0)
        {
          temp_storage_bytes = 1;
        }

        break;
      }

      if (num_items == OffsetT{})
      {
        break;
      }

      auto first_tile_previous = reinterpret_cast<InputT*>(allocations[0]);

      if constexpr (AliasOpt == MayAlias::Yes)
      {
        using AgentDifferenceInitT =
          detail::adjacent_difference::AgentDifferenceInit<InputIteratorT, InputT, OffsetT, ReadOpt == ReadOption::Left>;

        constexpr int init_block_size = AgentDifferenceInitT::BLOCK_THREADS;
        const int init_grid_size      = ::cuda::ceil_div(num_tiles, init_block_size);

#ifdef CUB_DEBUG_LOG
        _CubLog("Invoking DeviceAdjacentDifferenceInitKernel"
                "<<<%d, %d, 0, %lld>>>()\n",
                init_grid_size,
                init_block_size,
                reinterpret_cast<long long>(stream));
#endif // CUB_DEBUG_LOG

        error = CubDebug(
          THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_block_size, 0, stream)
            .doit(detail::adjacent_difference::
                    DeviceAdjacentDifferenceInitKernel<AgentDifferenceInitT, InputIteratorT, InputT, OffsetT>,
                  d_input,
                  first_tile_previous,
                  num_tiles,
                  tile_size));
        if (cudaSuccess != error)
        {
          break;
        }

        error = CubDebug(detail::DebugSyncStream(stream));

        if (cudaSuccess != error)
        {
          break;
        }
      }

#ifdef CUB_DEBUG_LOG
      _CubLog("Invoking DeviceAdjacentDifferenceDifferenceKernel"
              "<<<%d, %d, 0, %lld>>>()\n",
              num_tiles,
              AdjacentDifferencePolicyT::BLOCK_THREADS,
              reinterpret_cast<long long>(stream));
#endif // CUB_DEBUG_LOG

      using KernelPolicySelector = detail::adjacent_difference::policy_selector_from_hub<PolicyHub>;
      error                      = CubDebug(
        THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(
          num_tiles, AdjacentDifferencePolicyT::BLOCK_THREADS, 0, stream)
          .doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel < KernelPolicySelector,
                InputIteratorT,
                OutputIteratorT,
                DifferenceOpT,
                OffsetT,
                InputT,
                AliasOpt == MayAlias::Yes,
                ReadOpt == ReadOption::Left >,
                d_input,
                first_tile_previous,
                d_output,
                difference_op,
                num_items));
      if (cudaSuccess != error)
      {
        break;
      }

      error = CubDebug(detail::DebugSyncStream(stream));

      if (cudaSuccess != error)
      {
        break;
      }

    } while (false);

    return error;
  }

  CUB_RUNTIME_FUNCTION static cudaError_t Dispatch(
    void* d_temp_storage,
    size_t& temp_storage_bytes,
    InputIteratorT d_input,
    OutputIteratorT d_output,
    OffsetT num_items,
    DifferenceOpT difference_op,
    cudaStream_t stream)
  {
    cudaError error = cudaSuccess;
    do
    {
      // Get PTX version
      int ptx_version = 0;
      error           = CubDebug(PtxVersion(ptx_version));
      if (cudaSuccess != error)
      {
        break;
      }

      // Create dispatch functor
      DispatchAdjacentDifference dispatch(
        d_temp_storage, temp_storage_bytes, d_input, d_output, num_items, difference_op, stream);

      // Dispatch to chained policy
      error = CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch));
      if (cudaSuccess != error)
      {
        break;
      }
    } while (false);

    return error;
  }
};

namespace detail::adjacent_difference
{
template <MayAlias AliasOpt,
          ReadOption ReadOpt,
          typename InputIteratorT,
          typename OutputIteratorT,
          typename OffsetT,
          typename DifferenceOpT,
          typename PolicySelector        = policy_selector_from_types<InputIteratorT, AliasOpt == MayAlias::Yes>,
          typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
#if _CCCL_HAS_CONCEPTS()
  requires adjacent_difference_policy_selector<PolicySelector>
#endif // _CCCL_HAS_CONCEPTS()
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
  void* d_temp_storage,
  size_t& temp_storage_bytes,
  InputIteratorT d_input,
  OutputIteratorT d_output,
  OffsetT num_items,
  DifferenceOpT difference_op,
  cudaStream_t stream,
  PolicySelector policy_selector         = {},
  KernelLauncherFactory launcher_factory = {})
{
  using InputT = detail::it_value_t<InputIteratorT>;

  ::cuda::compute_capability cc{};
  if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc)))
  {
    return error;
  }

  const adjacent_difference_policy active_policy = policy_selector(cc);
#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)
  NV_IF_TARGET(NV_IS_HOST, ({
                 ::std::stringstream ss;
                 ss << active_policy;
                 _CubLog("Dispatching DeviceAdjacentDifference to compute capability %d.%d with tuning: %s\n",
                         cc.major_cap(),
                         cc.minor_cap(),
                         ss.str().c_str());
               }))
#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)

  const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread;
  const int num_tiles = static_cast<int>(::cuda::ceil_div(num_items, tile_size));

  size_t first_tile_previous_size = (AliasOpt == MayAlias::Yes) * num_tiles * sizeof(InputT);

  void* allocations[1]       = {nullptr};
  size_t allocation_sizes[1] = {(AliasOpt == MayAlias::Yes) * first_tile_previous_size};

  if (const auto error =
        CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
  {
    return error;
  }

  if (d_temp_storage == nullptr)
  {
    if (temp_storage_bytes == 0)
    {
      temp_storage_bytes = 1;
    }
    return cudaSuccess;
  }

  if (num_items == OffsetT{})
  {
    return cudaSuccess;
  }

  auto first_tile_previous = reinterpret_cast<InputT*>(allocations[0]);

  if constexpr (AliasOpt == MayAlias::Yes)
  {
    using AgentDifferenceInitT = AgentDifferenceInit<InputIteratorT, InputT, OffsetT, ReadOpt == ReadOption::Left>;

    constexpr int init_block_size = AgentDifferenceInitT::BLOCK_THREADS;
    const int init_grid_size      = ::cuda::ceil_div(num_tiles, init_block_size);

#ifdef CUB_DEBUG_LOG
    _CubLog("Invoking DeviceAdjacentDifferenceInitKernel"
            "<<<%d, %d, 0, %lld>>>()\n",
            init_grid_size,
            init_block_size,
            reinterpret_cast<long long>(stream));
#endif // CUB_DEBUG_LOG

    if (const auto error = CubDebug(
          THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_block_size, 0, stream)
            .doit(detail::adjacent_difference::
                    DeviceAdjacentDifferenceInitKernel<AgentDifferenceInitT, InputIteratorT, InputT, OffsetT>,
                  d_input,
                  first_tile_previous,
                  num_tiles,
                  tile_size)))
    {
      return error;
    }

    if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
    {
      return error;
    }
  }

#ifdef CUB_DEBUG_LOG
  _CubLog("Invoking DeviceAdjacentDifferenceDifferenceKernel"
          "<<<%d, %d, 0, %lld>>>()\n",
          num_tiles,
          active_policy.threads_per_block,
          reinterpret_cast<long long>(stream));
#endif // CUB_DEBUG_LOG

  if (const auto error = CubDebug(
        THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_tiles, active_policy.threads_per_block, 0, stream)
          .doit(DeviceAdjacentDifferenceDifferenceKernel < PolicySelector,
                InputIteratorT,
                OutputIteratorT,
                DifferenceOpT,
                OffsetT,
                InputT,
                AliasOpt == MayAlias::Yes,
                ReadOpt == ReadOption::Left >,
                d_input,
                first_tile_previous,
                d_output,
                difference_op,
                num_items)))
  {
    return error;
  }

  if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
  {
    return error;
  }

  return cudaSuccess;
}
} // namespace detail::adjacent_difference

CUB_NAMESPACE_END
