/home/runner/work/cccl/cccl/cub/cub/block/block_scan.cuh

File members: /home/runner/work/cccl/cccl/cub/cub/block/block_scan.cuh

/******************************************************************************
 * Copyright (c) 2011, Duane Merrill.  All rights reserved.
 * Copyright (c) 2011-2018, NVIDIA CORPORATION.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the NVIDIA CORPORATION nor the
 *       names of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 ******************************************************************************/

#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/block/specializations/block_scan_raking.cuh>
#include <cub/block/specializations/block_scan_warp_scans.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>

CUB_NAMESPACE_BEGIN

/******************************************************************************
 * Algorithmic variants
 ******************************************************************************/

enum BlockScanAlgorithm
{

  BLOCK_SCAN_RAKING,

  BLOCK_SCAN_RAKING_MEMOIZE,

  BLOCK_SCAN_WARP_SCANS,
};

template <typename T,
          int BLOCK_DIM_X,
          BlockScanAlgorithm ALGORITHM = BLOCK_SCAN_RAKING,
          int BLOCK_DIM_Y              = 1,
          int BLOCK_DIM_Z              = 1,
          int LEGACY_PTX_ARCH          = 0>
class BlockScan
{
private:
  enum
  {
    BLOCK_THREADS = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z,
  };

  static constexpr BlockScanAlgorithm SAFE_ALGORITHM =
    ((ALGORITHM == BLOCK_SCAN_WARP_SCANS) && (BLOCK_THREADS % CUB_WARP_THREADS(0) != 0))
      ? BLOCK_SCAN_RAKING
      : ALGORITHM;

  using WarpScans = BlockScanWarpScans<T, BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z>;
  using Raking =
    BlockScanRaking<T, BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z, (SAFE_ALGORITHM == BLOCK_SCAN_RAKING_MEMOIZE)>;

  using InternalBlockScan = cub::detail::conditional_t<SAFE_ALGORITHM == BLOCK_SCAN_WARP_SCANS, WarpScans, Raking>;

  using _TempStorage = typename InternalBlockScan::TempStorage;

  _TempStorage& temp_storage;

  unsigned int linear_tid;

  _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage()
  {
    __shared__ _TempStorage private_storage;
    return private_storage;
  }

public:
  struct TempStorage : Uninitialized<_TempStorage>
  {};

  _CCCL_DEVICE _CCCL_FORCEINLINE BlockScan()
      : temp_storage(PrivateStorage())
      , linear_tid(RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))
  {}

  _CCCL_DEVICE _CCCL_FORCEINLINE BlockScan(TempStorage& temp_storage)
      : temp_storage(temp_storage.Alias())
      , linear_tid(RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))
  {}

  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output)
  {
    T initial_value{};

    ExclusiveScan(input, output, initial_value, cub::Sum());
  }

  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output, T& block_aggregate)
  {
    T initial_value{};

    ExclusiveScan(input, output, initial_value, cub::Sum(), block_aggregate);
  }

  template <typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output, BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    ExclusiveScan(input, output, cub::Sum(), block_prefix_callback_op);
  }

  template <int ITEMS_PER_THREAD>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD])
  {
    T initial_value{};

    ExclusiveScan(input, output, initial_value, cub::Sum());
  }

  template <int ITEMS_PER_THREAD>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ExclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T& block_aggregate)
  {
    // Reduce consecutive thread items in registers
    T initial_value{};

    ExclusiveScan(input, output, initial_value, cub::Sum(), block_aggregate);
  }

  template <int ITEMS_PER_THREAD, typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(
    T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    ExclusiveScan(input, output, cub::Sum(), block_prefix_callback_op);
  }

  template <typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, T initial_value, ScanOp scan_op)
  {
    InternalBlockScan(temp_storage).ExclusiveScan(input, output, initial_value, scan_op);
  }

  template <typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ExclusiveScan(T input, T& output, T initial_value, ScanOp scan_op, T& block_aggregate)
  {
    InternalBlockScan(temp_storage).ExclusiveScan(input, output, initial_value, scan_op, block_aggregate);
  }

  template <typename ScanOp, typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ExclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op, block_prefix_callback_op);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op)
  {
    // Reduce consecutive thread items in registers
    T thread_prefix = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op);

    // Exclusive scan in registers with prefix as seed
    internal::ThreadScanExclusive(input, output, scan_op, thread_prefix);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(
    T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op, T& block_aggregate)
  {
    // Reduce consecutive thread items in registers
    T thread_prefix = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op, block_aggregate);

    // Exclusive scan in registers with prefix as seed
    internal::ThreadScanExclusive(input, output, scan_op, thread_prefix);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp, typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(
    T (&input)[ITEMS_PER_THREAD],
    T (&output)[ITEMS_PER_THREAD],
    ScanOp scan_op,
    BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    // Reduce consecutive thread items in registers
    T thread_prefix = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_prefix_callback_op);

    // Exclusive scan in registers with prefix as seed
    internal::ThreadScanExclusive(input, output, scan_op, thread_prefix);
  }

#ifndef DOXYGEN_SHOULD_SKIP_THIS // Do not document no-initial-value scans

  template <typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op)
  {
    InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op);
  }

  template <typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate)
  {
    InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op, block_aggregate);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op)
  {
    // Reduce consecutive thread items in registers
    T thread_partial = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_partial, thread_partial, scan_op);

    // Exclusive scan in registers with prefix
    internal::ThreadScanExclusive(input, output, scan_op, thread_partial, (linear_tid != 0));
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op, T& block_aggregate)
  {
    // Reduce consecutive thread items in registers
    T thread_partial = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_partial, thread_partial, scan_op, block_aggregate);

    // Exclusive scan in registers with prefix
    internal::ThreadScanExclusive(input, output, scan_op, thread_partial, (linear_tid != 0));
  }

#endif // DOXYGEN_SHOULD_SKIP_THIS  // Do not document no-initial-value scans

  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output)
  {
    InclusiveScan(input, output, cub::Sum());
  }

  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output, T& block_aggregate)
  {
    InclusiveScan(input, output, cub::Sum(), block_aggregate);
  }

  template <typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output, BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    InclusiveScan(input, output, cub::Sum(), block_prefix_callback_op);
  }

  template <int ITEMS_PER_THREAD>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD])
  {
    if (ITEMS_PER_THREAD == 1)
    {
      InclusiveSum(input[0], output[0]);
    }
    else
    {
      // Reduce consecutive thread items in registers
      Sum scan_op;
      T thread_prefix = internal::ThreadReduce(input, scan_op);

      // Exclusive thread block-scan
      ExclusiveSum(thread_prefix, thread_prefix);

      // Inclusive scan in registers with prefix as seed
      internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
    }
  }

  template <int ITEMS_PER_THREAD>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  InclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T& block_aggregate)
  {
    if (ITEMS_PER_THREAD == 1)
    {
      InclusiveSum(input[0], output[0], block_aggregate);
    }
    else
    {
      // Reduce consecutive thread items in registers
      Sum scan_op;
      T thread_prefix = internal::ThreadReduce(input, scan_op);

      // Exclusive thread block-scan
      ExclusiveSum(thread_prefix, thread_prefix, block_aggregate);

      // Inclusive scan in registers with prefix as seed
      internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
    }
  }

  template <int ITEMS_PER_THREAD, typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(
    T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    if (ITEMS_PER_THREAD == 1)
    {
      InclusiveSum(input[0], output[0], block_prefix_callback_op);
    }
    else
    {
      // Reduce consecutive thread items in registers
      Sum scan_op;
      T thread_prefix = internal::ThreadReduce(input, scan_op);

      // Exclusive thread block-scan
      ExclusiveSum(thread_prefix, thread_prefix, block_prefix_callback_op);

      // Inclusive scan in registers with prefix as seed
      internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
    }
  }

  template <typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op)
  {
    InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op);
  }

  template <typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate)
  {
    InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op, block_aggregate);
  }

  template <typename ScanOp, typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  InclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op, block_prefix_callback_op);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op)
  {
    if (ITEMS_PER_THREAD == 1)
    {
      InclusiveScan(input[0], output[0], scan_op);
    }
    else
    {
      // Reduce consecutive thread items in registers
      T thread_prefix = internal::ThreadReduce(input, scan_op);

      // Exclusive thread block-scan
      ExclusiveScan(thread_prefix, thread_prefix, scan_op);

      // Inclusive scan in registers with prefix as seed (first thread does not seed)
      internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
    }
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op)
  {
    // Reduce consecutive thread items in registers
    T thread_prefix = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op);

    // Exclusive scan in registers with prefix as seed
    internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void
  InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op, T& block_aggregate)
  {
    if (ITEMS_PER_THREAD == 1)
    {
      InclusiveScan(input[0], output[0], scan_op, block_aggregate);
    }
    else
    {
      // Reduce consecutive thread items in registers
      T thread_prefix = internal::ThreadReduce(input, scan_op);

      // Exclusive thread block-scan (with no initial value)
      ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_aggregate);

      // Inclusive scan in registers with prefix as seed (first thread does not seed)
      internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
    }
  }

  template <int ITEMS_PER_THREAD, typename ScanOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(
    T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op, T& block_aggregate)
  {
    // Reduce consecutive thread items in registers
    T thread_prefix = internal::ThreadReduce(input, scan_op);

    // Exclusive thread block-scan
    ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op, block_aggregate);

    // Exclusive scan in registers with prefix as seed
    internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
  }

  template <int ITEMS_PER_THREAD, typename ScanOp, typename BlockPrefixCallbackOp>
  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(
    T (&input)[ITEMS_PER_THREAD],
    T (&output)[ITEMS_PER_THREAD],
    ScanOp scan_op,
    BlockPrefixCallbackOp& block_prefix_callback_op)
  {
    if (ITEMS_PER_THREAD == 1)
    {
      InclusiveScan(input[0], output[0], scan_op, block_prefix_callback_op);
    }
    else
    {
      // Reduce consecutive thread items in registers
      T thread_prefix = internal::ThreadReduce(input, scan_op);

      // Exclusive thread block-scan
      ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_prefix_callback_op);

      // Inclusive scan in registers with prefix as seed
      internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
    }
  }

};

CUB_NAMESPACE_END