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

/**
 * @file cub::AgentReduceByKey implements a stateful abstraction of CUDA thread
 *       blocks for participating in device-wide reduce-value-by-key.
 */

#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/single_pass_scan_operators.cuh>
#include <cub/block/block_discontinuity.cuh>
#include <cub/block/block_load.cuh>
#include <cub/block/block_scan.cuh>
#include <cub/block/block_store.cuh>
#include <cub/iterator/cache_modified_input_iterator.cuh>

#include <cuda/std/__functional/operations.h>
#include <cuda/std/__type_traits/conditional.h>
#include <cuda/std/__type_traits/is_pointer.h>
#include <cuda/std/__type_traits/is_same.h>
CUB_NAMESPACE_BEGIN

/******************************************************************************
 * Tuning policy types
 ******************************************************************************/

/**
 * @brief Parameterizable tuning policy type for AgentReduceByKey
 *
 * @tparam BlockThreads
 *   Threads per thread block
 *
 * @tparam ItemsPerThread
 *   Items per thread (per tile of input)
 *
 * @tparam LoadAlgorithm
 *   The BlockLoad algorithm to use
 *
 * @tparam LoadModifier
 *   Cache load modifier for reading input elements
 *
 * @tparam ScanAlgorithm
 *   The BlockScan algorithm to use
 *
 * @tparam DelayConstructorT
 *   Implementation detail, do not specify directly, requirements on the
 *   content of this type are subject to breaking change.
 */
template <int BlockThreads,
          int ItemsPerThread,
          BlockLoadAlgorithm LoadAlgorithm,
          CacheLoadModifier LoadModifier,
          BlockScanAlgorithm ScanAlgorithm,
          typename DelayConstructorT = detail::fixed_delay_constructor_t<350, 450>>
struct AgentReduceByKeyPolicy
{
  ///< Threads per thread block
  static constexpr int BLOCK_THREADS = BlockThreads;

  ///< Items per thread (per tile of input)
  static constexpr int ITEMS_PER_THREAD = ItemsPerThread;

  ///< The BlockLoad algorithm to use
  static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm;

  ///< Cache load modifier for reading input elements
  static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier;

  ///< The BlockScan algorithm to use
  static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm;

  struct detail
  {
    using delay_constructor_t = DelayConstructorT;
  };
};

/******************************************************************************
 * Thread block abstractions
 ******************************************************************************/

namespace detail::reduce
{
/**
 * @brief AgentReduceByKey implements a stateful abstraction of CUDA thread
 *        blocks for participating in device-wide reduce-value-by-key
 *
 * @tparam AgentReduceByKeyPolicyT
 *   Parameterized AgentReduceByKeyPolicy tuning policy type
 *
 * @tparam KeysInputIteratorT
 *   Random-access input iterator type for keys
 *
 * @tparam UniqueOutputIteratorT
 *   Random-access output iterator type for keys
 *
 * @tparam ValuesInputIteratorT
 *   Random-access input iterator type for values
 *
 * @tparam AggregatesOutputIteratorT
 *   Random-access output iterator type for values
 *
 * @tparam NumRunsOutputIteratorT
 *   Output iterator type for recording number of items selected
 *
 * @tparam EqualityOpT
 *   KeyT equality operator type
 *
 * @tparam ReductionOpT
 *   ValueT reduction operator type
 *
 * @tparam OffsetT
 *   Signed integer type for global offsets
 *
 * @tparam AccumT
 *   The type of intermediate accumulator (according to P2322R6)
 */
template <typename AgentReduceByKeyPolicyT,
          typename KeysInputIteratorT,
          typename UniqueOutputIteratorT,
          typename ValuesInputIteratorT,
          typename AggregatesOutputIteratorT,
          typename NumRunsOutputIteratorT,
          typename EqualityOpT,
          typename ReductionOpT,
          typename OffsetT,
          typename AccumT,
          typename StreamingContextT>
struct AgentReduceByKey
{
  // Whether or not this is a streaming invocation (i.e., multiple kernel invocations over partitions of the input)
  static constexpr bool is_streaming_invocation = !::cuda::std::is_same_v<StreamingContextT, NullType>;

  //---------------------------------------------------------------------
  // Types and constants
  //---------------------------------------------------------------------

  // The input keys type
  using KeyInputT = it_value_t<KeysInputIteratorT>;

  // The output keys type
  using KeyOutputT = non_void_value_t<UniqueOutputIteratorT, KeyInputT>;

  // The input values type
  using ValueInputT = it_value_t<ValuesInputIteratorT>;

  // Tuple type for scanning (pairs accumulated segment-value with
  // segment-index)
  using OffsetValuePairT = KeyValuePair<OffsetT, AccumT>;

  // Tuple type for pairing keys and values
  using KeyValuePairT = KeyValuePair<KeyOutputT, AccumT>;

  // Tile status descriptor interface type
  using ScanTileStateT = ReduceByKeyScanTileState<AccumT, OffsetT>;

  // Guarded inequality functor
  template <typename _EqualityOpT>
  struct GuardedInequalityWrapper
  {
    /// Wrapped equality operator
    _EqualityOpT op;

    /// Items remaining
    int num_remaining;

    /// Constructor
    _CCCL_HOST_DEVICE _CCCL_FORCEINLINE GuardedInequalityWrapper(_EqualityOpT op, int num_remaining)
        : op(op)
        , num_remaining(num_remaining)
    {}

    /// Boolean inequality operator, returns <tt>(a != b)</tt>
    template <typename T>
    _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(const T& a, const T& b, int idx) const
    {
      if (idx < num_remaining)
      {
        return !op(a, b); // In bounds
      }

      // Return true if first out-of-bounds item, false otherwise
      return (idx == num_remaining);
    }
  };

  // Constants
  static constexpr int BLOCK_THREADS     = AgentReduceByKeyPolicyT::BLOCK_THREADS;
  static constexpr int ITEMS_PER_THREAD  = AgentReduceByKeyPolicyT::ITEMS_PER_THREAD;
  static constexpr int TILE_ITEMS        = BLOCK_THREADS * ITEMS_PER_THREAD;
  static constexpr int TWO_PHASE_SCATTER = (ITEMS_PER_THREAD > 1);

  // Whether or not the scan operation has a zero-valued identity value (true
  // if we're performing addition on a primitive type)
  static constexpr int HAS_IDENTITY_ZERO =
    (::cuda::std::is_same_v<ReductionOpT, ::cuda::std::plus<>>) && (is_primitive<AccumT>::value);

  // Cache-modified Input iterator wrapper type (for applying cache modifier)
  // for keys Wrap the native input pointer with
  // CacheModifiedValuesInputIterator or directly use the supplied input
  // iterator type
  using WrappedKeysInputIteratorT =
    ::cuda::std::_If<::cuda::std::is_pointer_v<KeysInputIteratorT>,
                     CacheModifiedInputIterator<AgentReduceByKeyPolicyT::LOAD_MODIFIER, KeyInputT, OffsetT>,
                     KeysInputIteratorT>;

  // Cache-modified Input iterator wrapper type (for applying cache modifier)
  // for values Wrap the native input pointer with
  // CacheModifiedValuesInputIterator or directly use the supplied input
  // iterator type
  using WrappedValuesInputIteratorT =
    ::cuda::std::_If<::cuda::std::is_pointer_v<ValuesInputIteratorT>,
                     CacheModifiedInputIterator<AgentReduceByKeyPolicyT::LOAD_MODIFIER, ValueInputT, OffsetT>,
                     ValuesInputIteratorT>;

  // Cache-modified Input iterator wrapper type (for applying cache modifier)
  // for fixup values Wrap the native input pointer with
  // CacheModifiedValuesInputIterator or directly use the supplied input
  // iterator type
  using WrappedFixupInputIteratorT =
    ::cuda::std::_If<::cuda::std::is_pointer_v<AggregatesOutputIteratorT>,
                     CacheModifiedInputIterator<AgentReduceByKeyPolicyT::LOAD_MODIFIER, ValueInputT, OffsetT>,
                     AggregatesOutputIteratorT>;

  // Reduce-value-by-segment scan operator
  using ReduceBySegmentOpT = ReduceBySegmentOp<ReductionOpT>;

  // Parameterized BlockLoad type for keys
  using BlockLoadKeysT =
    BlockLoad<KeyOutputT, BLOCK_THREADS, ITEMS_PER_THREAD, AgentReduceByKeyPolicyT::LOAD_ALGORITHM>;

  // Parameterized BlockLoad type for values
  using BlockLoadValuesT = BlockLoad<AccumT, BLOCK_THREADS, ITEMS_PER_THREAD, AgentReduceByKeyPolicyT::LOAD_ALGORITHM>;

  // Parameterized BlockDiscontinuity type for keys
  using BlockDiscontinuityKeys = BlockDiscontinuity<KeyOutputT, BLOCK_THREADS>;

  // Parameterized BlockScan type
  using BlockScanT = BlockScan<OffsetValuePairT, BLOCK_THREADS, AgentReduceByKeyPolicyT::SCAN_ALGORITHM>;

  // Callback type for obtaining tile prefix during block scan
  using DelayConstructorT = typename AgentReduceByKeyPolicyT::detail::delay_constructor_t;
  using TilePrefixCallbackOpT =
    TilePrefixCallbackOp<OffsetValuePairT, ReduceBySegmentOpT, ScanTileStateT, DelayConstructorT>;

  // Key and value exchange types
  using KeyExchangeT   = KeyOutputT[TILE_ITEMS + 1];
  using ValueExchangeT = AccumT[TILE_ITEMS + 1];

  // Shared memory type for this thread block
  union _TempStorage
  {
    struct ScanStorage
    {
      // Smem needed for tile scanning
      typename BlockScanT::TempStorage scan;

      // Smem needed for cooperative prefix callback
      typename TilePrefixCallbackOpT::TempStorage prefix;

      // Smem needed for discontinuity detection
      typename BlockDiscontinuityKeys::TempStorage discontinuity;
    } scan_storage;

    // Smem needed for loading keys
    typename BlockLoadKeysT::TempStorage load_keys;

    // Smem needed for loading values
    typename BlockLoadValuesT::TempStorage load_values;

    // Smem needed for compacting key value pairs(allows non POD items in this
    // union)
    Uninitialized<KeyValuePairT[TILE_ITEMS + 1]> raw_exchange;
  };

  // Alias wrapper allowing storage to be unioned
  struct TempStorage : Uninitialized<_TempStorage>
  {};

  //---------------------------------------------------------------------
  // Per-thread fields
  //---------------------------------------------------------------------

  /// Reference to temp_storage
  _TempStorage& temp_storage;

  /// Input keys
  WrappedKeysInputIteratorT d_keys_in;

  /// Unique output keys
  UniqueOutputIteratorT d_unique_out;

  /// Input values
  WrappedValuesInputIteratorT d_values_in;

  /// Output value aggregates
  AggregatesOutputIteratorT d_aggregates_out;

  /// Output pointer for total number of segments identified
  NumRunsOutputIteratorT d_num_runs_out;

  /// KeyT equality operator
  EqualityOpT equality_op;

  /// Reduction operator
  ReductionOpT reduction_op;

  /// Reduce-by-segment scan operator
  ReduceBySegmentOpT scan_op;

  /// Streaming context providing context about this partition for streaming invocations
  StreamingContextT streaming_context;

  //---------------------------------------------------------------------
  // Constructor
  //---------------------------------------------------------------------

  /**
   * @param temp_storage
   *   Reference to temp_storage
   *
   * @param d_keys_in
   *   Input keys
   *
   * @param d_unique_out
   *   Unique output keys
   *
   * @param d_values_in
   *   Input values
   *
   * @param d_aggregates_out
   *   Output value aggregates
   *
   * @param d_num_runs_out
   *   Output pointer for total number of segments identified
   *
   * @param equality_op
   *   KeyT equality operator
   *
   * @param reduction_op
   *   ValueT reduction operator
   *
   * @param streaming_context
   *   Streaming context providing context about this partition for streaming invocations
   */
  template <typename StreamingContext>
  _CCCL_DEVICE _CCCL_FORCEINLINE AgentReduceByKey(
    TempStorage& temp_storage,
    KeysInputIteratorT d_keys_in,
    UniqueOutputIteratorT d_unique_out,
    ValuesInputIteratorT d_values_in,
    AggregatesOutputIteratorT d_aggregates_out,
    NumRunsOutputIteratorT d_num_runs_out,
    EqualityOpT equality_op,
    ReductionOpT reduction_op,
    StreamingContext streaming_context)
      : temp_storage(temp_storage.Alias())
      , d_keys_in(d_keys_in)
      , d_unique_out(d_unique_out + streaming_context.num_uniques())
      , d_values_in(d_values_in)
      , d_aggregates_out(d_aggregates_out + streaming_context.num_uniques())
      , d_num_runs_out(d_num_runs_out)
      , equality_op(equality_op)
      , reduction_op(reduction_op)
      , scan_op(reduction_op)
      , streaming_context(streaming_context)
  {}

  _CCCL_DEVICE _CCCL_FORCEINLINE AgentReduceByKey(
    TempStorage& temp_storage,
    KeysInputIteratorT d_keys_in,
    UniqueOutputIteratorT d_unique_out,
    ValuesInputIteratorT d_values_in,
    AggregatesOutputIteratorT d_aggregates_out,
    NumRunsOutputIteratorT d_num_runs_out,
    EqualityOpT equality_op,
    ReductionOpT reduction_op,
    NullType streaming_context)
      : temp_storage(temp_storage.Alias())
      , d_keys_in(d_keys_in)
      , d_unique_out(d_unique_out)
      , d_values_in(d_values_in)
      , d_aggregates_out(d_aggregates_out)
      , d_num_runs_out(d_num_runs_out)
      , equality_op(equality_op)
      , reduction_op(reduction_op)
      , scan_op(reduction_op)
      , streaming_context(streaming_context)
  {}

  //---------------------------------------------------------------------
  // Scatter utility methods
  //---------------------------------------------------------------------

  /**
   * Directly scatter flagged items to output offsets
   */
  _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterDirect(
    KeyValuePairT (&scatter_items)[ITEMS_PER_THREAD],
    OffsetT (&segment_flags)[ITEMS_PER_THREAD],
    OffsetT (&segment_indices)[ITEMS_PER_THREAD])
  {
    // Scatter flagged keys and values
    _CCCL_PRAGMA_UNROLL_FULL()
    for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
    {
      if (segment_flags[ITEM])
      {
        d_unique_out[segment_indices[ITEM]]     = scatter_items[ITEM].key;
        d_aggregates_out[segment_indices[ITEM]] = scatter_items[ITEM].value;
      }
    }
  }

  /**
   * 2-phase scatter flagged items to output offsets
   *
   * The exclusive scan causes each head flag to be paired with the previous
   * value aggregate: the scatter offsets must be decremented for value
   * aggregates
   */
  _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterTwoPhase(
    KeyValuePairT (&scatter_items)[ITEMS_PER_THREAD],
    OffsetT (&segment_flags)[ITEMS_PER_THREAD],
    OffsetT (&segment_indices)[ITEMS_PER_THREAD],
    OffsetT num_tile_segments,
    OffsetT num_tile_segments_prefix)
  {
    __syncthreads();

    // Compact and scatter pairs
    _CCCL_PRAGMA_UNROLL_FULL()
    for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
    {
      if (segment_flags[ITEM])
      {
        temp_storage.raw_exchange.Alias()[segment_indices[ITEM] - num_tile_segments_prefix] = scatter_items[ITEM];
      }
    }

    __syncthreads();

    for (int item = threadIdx.x; item < num_tile_segments; item += BLOCK_THREADS)
    {
      KeyValuePairT pair                                = temp_storage.raw_exchange.Alias()[item];
      d_unique_out[num_tile_segments_prefix + item]     = pair.key;
      d_aggregates_out[num_tile_segments_prefix + item] = pair.value;
    }
  }

  /**
   * Scatter flagged items
   */
  _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter(
    KeyValuePairT (&scatter_items)[ITEMS_PER_THREAD],
    OffsetT (&segment_flags)[ITEMS_PER_THREAD],
    OffsetT (&segment_indices)[ITEMS_PER_THREAD],
    OffsetT num_tile_segments,
    OffsetT num_tile_segments_prefix)
  {
    // Do a one-phase scatter if (a) two-phase is disabled or (b) the average
    // number of selected items per thread is less than one
    if (TWO_PHASE_SCATTER && (num_tile_segments > BLOCK_THREADS))
    {
      ScatterTwoPhase(scatter_items, segment_flags, segment_indices, num_tile_segments, num_tile_segments_prefix);
    }
    else
    {
      ScatterDirect(scatter_items, segment_flags, segment_indices);
    }
  }

  //---------------------------------------------------------------------
  // Cooperatively scan a device-wide sequence of tiles with other CTAs
  //---------------------------------------------------------------------

  /**
   * @brief Process a tile of input (dynamic chained scan)
   *
   * @tparam IS_LAST_TILE
   *   Whether the current tile is the last tile
   *
   * @param num_remaining
   *   Number of global input items remaining (including this tile)
   *
   * @param tile_idx
   *   Tile index
   *
   * @param tile_offset
   *   Tile offset
   *
   * @param tile_state
   *   Global tile state descriptor
   */
  template <bool IS_LAST_TILE>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ConsumeTile(OffsetT num_remaining, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state)
  {
    // Tile keys
    KeyOutputT keys[ITEMS_PER_THREAD];

    // Tile keys shuffled up
    KeyOutputT prev_keys[ITEMS_PER_THREAD];

    // Tile values
    AccumT values[ITEMS_PER_THREAD];

    // Segment head flags
    OffsetT head_flags[ITEMS_PER_THREAD];

    // Segment indices
    OffsetT segment_indices[ITEMS_PER_THREAD];

    // Zipped values and segment flags|indices
    OffsetValuePairT scan_items[ITEMS_PER_THREAD];

    // Zipped key value pairs for scattering
    KeyValuePairT scatter_items[ITEMS_PER_THREAD];

    // Load keys
    if (IS_LAST_TILE)
    {
      BlockLoadKeysT(temp_storage.load_keys).Load(d_keys_in + tile_offset, keys, num_remaining);
    }
    else
    {
      BlockLoadKeysT(temp_storage.load_keys).Load(d_keys_in + tile_offset, keys);
    }

    // Load tile predecessor key in first thread
    KeyOutputT tile_predecessor;
    if (threadIdx.x == 0)
    {
      // if (tile_idx == 0)
      //   first tile gets repeat of first item (thus first item will not
      //   be flagged as a head)
      // else
      //   Subsequent tiles get last key from previous tile
      if constexpr (is_streaming_invocation)
      {
        tile_predecessor = (tile_idx == 0) ? streaming_context.predecessor_key() : d_keys_in[tile_offset - 1];
      }
      else
      {
        tile_predecessor = (tile_idx == 0) ? keys[0] : d_keys_in[tile_offset - 1];
      }
    }

    __syncthreads();

    // Load values
    if (IS_LAST_TILE)
    {
      BlockLoadValuesT(temp_storage.load_values).Load(d_values_in + tile_offset, values, num_remaining);
    }
    else
    {
      BlockLoadValuesT(temp_storage.load_values).Load(d_values_in + tile_offset, values);
    }

    __syncthreads();

    // Initialize head-flags and shuffle up the previous keys
    if (IS_LAST_TILE)
    {
      // Use custom flag operator to additionally flag the first out-of-bounds item
      GuardedInequalityWrapper<EqualityOpT> flag_op(equality_op, num_remaining);
      BlockDiscontinuityKeys(temp_storage.scan_storage.discontinuity)
        .FlagHeads(head_flags, keys, prev_keys, flag_op, tile_predecessor);
    }
    else
    {
      InequalityWrapper<EqualityOpT> flag_op(equality_op);
      BlockDiscontinuityKeys(temp_storage.scan_storage.discontinuity)
        .FlagHeads(head_flags, keys, prev_keys, flag_op, tile_predecessor);
    }

    // Reset head-flag on the very first item to make sure we don't start a new run for data where
    // (key[0] == key[0]) is false (e.g., when key[0] is NaN)
    if constexpr (is_streaming_invocation)
    {
      if (streaming_context.is_first_partition() && threadIdx.x == 0 && tile_idx == 0)
      {
        head_flags[0] = 0;
      }
    }
    else
    {
      if (threadIdx.x == 0 && tile_idx == 0)
      {
        head_flags[0] = 0;
      }
    }

    // Zip values and head flags
    _CCCL_PRAGMA_UNROLL_FULL()
    for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
    {
      scan_items[ITEM].value = values[ITEM];
      scan_items[ITEM].key   = head_flags[ITEM];
    }

    // Perform exclusive tile scan
    // Inclusive block-wide scan aggregate
    OffsetValuePairT block_aggregate;

    // Number of segments prior to this tile
    OffsetT num_segments_prefix;

    // The tile prefix folded with block_aggregate
    OffsetValuePairT total_aggregate;

    if (tile_idx == 0)
    {
      // Scan first tile
      // First partition does not need to account for preceding partitions
      if constexpr (is_streaming_invocation)
      {
        if (streaming_context.is_first_partition())
        {
          BlockScanT(temp_storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, scan_op, block_aggregate);
          num_segments_prefix = 0;
          total_aggregate     = block_aggregate;
        }
        // Subsequent partitions need to account for preceding partitions
        else
        {
          auto init_value = OffsetValuePairT{0, streaming_context.prefix()};
          BlockScanT(temp_storage.scan_storage.scan)
            .ExclusiveScan(scan_items, scan_items, init_value, scan_op, block_aggregate);
          num_segments_prefix = 0;
          // note, block_aggregate does not include the prefix
          block_aggregate = scan_op(init_value, block_aggregate);
          total_aggregate = block_aggregate;
        }
      }
      else
      {
        BlockScanT(temp_storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, scan_op, block_aggregate);
        num_segments_prefix = 0;
        total_aggregate     = block_aggregate;
      }

      // Update tile status if there are successor tiles
      if ((!IS_LAST_TILE) && (threadIdx.x == 0))
      {
        tile_state.SetInclusive(0, block_aggregate);
      }
    }
    else
    {
      // Scan non-first tile
      TilePrefixCallbackOpT prefix_op(tile_state, temp_storage.scan_storage.prefix, scan_op, tile_idx);
      BlockScanT(temp_storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, scan_op, prefix_op);

      block_aggregate     = prefix_op.GetBlockAggregate();
      num_segments_prefix = prefix_op.GetExclusivePrefix().key;
      total_aggregate     = prefix_op.GetInclusivePrefix();
    }

    // Rezip scatter items and segment indices
    _CCCL_PRAGMA_UNROLL_FULL()
    for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
    {
      scatter_items[ITEM].key   = prev_keys[ITEM];
      scatter_items[ITEM].value = scan_items[ITEM].value;
      segment_indices[ITEM]     = scan_items[ITEM].key;
    }

    // At this point, each flagged segment head has:
    //  - The key for the previous segment
    //  - The reduced value from the previous segment
    //  - The segment index for the reduced value

    // Scatter flagged keys and values
    OffsetT num_tile_segments = block_aggregate.key;
    Scatter(scatter_items, head_flags, segment_indices, num_tile_segments, num_segments_prefix);

    // Last thread in last tile will output final count (and last pair, if necessary)
    if ((IS_LAST_TILE) && (threadIdx.x == BLOCK_THREADS - 1))
    {
      OffsetT num_segments = num_segments_prefix + num_tile_segments;

      // If the last tile is a full tile, we need to write out the run ending with the last item
      // If this was not a full tile, we already have flagged the head of one-past-the-last-item
      if (num_remaining == TILE_ITEMS)
      {
        if constexpr (is_streaming_invocation)
        {
          if (streaming_context.is_last_partition())
          {
            d_unique_out[num_segments]     = keys[ITEMS_PER_THREAD - 1];
            d_aggregates_out[num_segments] = total_aggregate.value;
            num_segments++;
          }
          else
          {
            // Write the prefix aggregate of this partition as context for the subsequent partition
            streaming_context.write_prefix(total_aggregate.value);
          }
        }
        else
        {
          d_unique_out[num_segments]     = keys[ITEMS_PER_THREAD - 1];
          d_aggregates_out[num_segments] = total_aggregate.value;
          num_segments++;
        }
      }

      if constexpr (is_streaming_invocation)
      {
        // Add the number of unique items in this partition to the global aggregate
        auto total_uniques = streaming_context.add_num_uniques(num_segments);

        // If this is the last partition, write out the number of unique items
        if (streaming_context.is_last_partition())
        {
          *d_num_runs_out = total_uniques;
        }
      }
      else
      {
        *d_num_runs_out = num_segments;
      }
    }
  }

  /**
   * @brief Scan tiles of items as part of a dynamic chained scan
   *
   * @param num_items
   *   Total number of input items
   *
   * @param tile_state
   *   Global tile state descriptor
   *
   * @param start_tile
   *   The starting tile for the current grid
   */
  _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeRange(OffsetT num_items, ScanTileStateT& tile_state, int start_tile)
  {
    // Blocks are launched in increasing order, so just assign one tile per
    // block

    // Current tile index
    int tile_idx = start_tile + blockIdx.x;

    // Global offset for the current tile
    OffsetT tile_offset = OffsetT(TILE_ITEMS) * tile_idx;

    // Remaining items (including this tile)
    OffsetT num_remaining = num_items - tile_offset;

    if (num_remaining > TILE_ITEMS)
    {
      // Not last tile
      ConsumeTile<false>(num_remaining, tile_idx, tile_offset, tile_state);
    }
    else if (num_remaining > 0)
    {
      // Last tile
      ConsumeTile<true>(num_remaining, tile_idx, tile_offset, tile_state);
    }
  }
};
} // namespace detail::reduce

CUB_NAMESPACE_END
