Add inclusive_scan with initial value support (warp/block)

cub/cub/block/block_scan.cuh

@@ -2242,6 +2242,63 @@ public:
     }
   }
 
+  //! @rst
+  //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor.
+  //! Each thread contributes an array of consecutive input elements.
+  //!
+  //! - Supports non-commutative scan operators.
+  //! - @blocked
+  //! - @granularity
+  //! - @smemreuse
+  //!
+  //! Snippet
+  //! +++++++
+  //!
+  //! The code snippet below illustrates an inclusive prefix max scan of 128 integer items that
+  //! are partitioned in a :ref:`blocked arrangement <flexible-data-arrangement>` across 64 threads
+  //! where each thread owns 2 consecutive items.
+  //!
+  //! .. literalinclude:: ../../test/catch2_test_block_scan_api.cu
+  //!
+  //!     :language: c++
+  //!     :dedent:
+  //!     :start-after: example-begin inclusive-scan-array-init-value
+  //!     :end-before: example-end inclusive-scan-array-init-value
+  //!
+  //!
+  //! @endrst
+  //!
+  //! @tparam ITEMS_PER_THREAD
+  //!   **[inferred]** The number of consecutive items partitioned onto each thread.
+  //!
+  //! @tparam ScanOp
+  //!   **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)`
+  //!
+  //! @param[in] input
+  //!   Calling thread's input items
+  //!
+  //! @param[out] output
+  //!   Calling thread's output items (may be aliased to `input`)
+  //!
+  //! @param[in] initial_value
+  //!   Initial value to seed the inclusive scan (uniform across block)
+  //!
+  //! @param[in] scan_op
+  //!   Binary scan functor
+  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);
+  }
+
   //! @rst
   //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor.
   //! Each thread contributes an array of consecutive input elements. Also provides every thread
@@ -2325,6 +2382,69 @@ public:
     }
   }
 
+  //! @rst
+  //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor.
+  //! Each thread contributes an array of consecutive input elements. Also provides every thread
+  //! with the block-wide ``block_aggregate`` of all inputs.
+  //!
+  //! - Supports non-commutative scan operators.
+  //! - @blocked
+  //! - @granularity
+  //! - @smemreuse
+  //!
+  //! Snippet
+  //! +++++++
+  //!
+  //! The code snippet below illustrates an inclusive prefix max scan of 128 integer items that
+  //! are partitioned in a :ref:`blocked arrangement <flexible-data-arrangement>` across 64 threads
+  //! where each thread owns 2 consecutive items.
+  //!
+  //! .. literalinclude:: ../../test/catch2_test_block_scan_api.cu
+  //!
+  //!     :language: c++
+  //!     :dedent:
+  //!     :start-after: example-begin inclusive-scan-array-aggregate-init-value
+  //!     :end-before: example-end inclusive-scan-array-aggregate-init-value
+  //!
+  //! The value ``126`` will be stored in ``block_aggregate`` for all threads.
+  //!
+  //! @endrst
+  //!
+  //! @tparam ITEMS_PER_THREAD
+  //!   **[inferred]** The number of consecutive items partitioned onto each thread.
+  //!
+  //! @tparam ScanOp
+  //!   **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)`
+  //!
+  //! @param[in] input
+  //!   Calling thread's input items
+  //!
+  //! @param[out] output
+  //!   Calling thread's output items (may be aliased to `input`)
+  //!
+  //! @param[in] initial_value
+  //!   Initial value to seed the inclusive scan (uniform across block). It is not taken
+  //!   into account for block_aggregate.
+  //!
+  //! @param[in] scan_op
+  //!   Binary scan functor
+  //!
+  //! @param[out] block_aggregate
+  //!   Block-wide aggregate reduction of input items
+  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);
+  }
+
   //! @rst
   //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor.
   //! Each thread contributes an array of consecutive input elements.

cub/cub/warp/warp_scan.cuh

@@ -475,14 +475,65 @@ public:
   //! @param[out] inclusive_output
   //!   Calling thread's output item. May be aliased with `input`
   //!
-  //! @param[in] can_op
+  //! @param[in] scan_op
   //!   Binary scan operator
   template <typename ScanOp>
   _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& inclusive_output, ScanOp scan_op)
   {
     InternalWarpScan(temp_storage).InclusiveScan(input, inclusive_output, scan_op);
   }
 
+  //! @rst
+  //! Computes an inclusive prefix scan using the specified binary scan functor across the
+  //! calling warp.
+  //!
+  //! * @smemwarpreuse
+  //!
+  //! Snippet
+  //! +++++++
+  //!
+  //! The code snippet below illustrates four concurrent warp-wide inclusive prefix sum scans
+  //! within a block of 128 threads (one per each of the 32-thread warps).
+  //!
+  //! .. literalinclude:: ../../../cub/test/catch2_test_warp_scan_api.cu
+  //!     :language: c++
+  //!     :dedent:
+  //!     :start-after: example-begin inclusive-warp-scan-init-value
+  //!     :end-before: example-end inclusive-warp-scan-init-value
+  //!
+  //! Suppose the set of input ``thread_data`` in the first warp is
+  //! ``{0, 1, 2, 3, ..., 31}``, in the second warp is ``{1, 2, 3, 4, ..., 32}`` etc.
+  //!  The corresponding output ``thread_data`` for a max operation in the first
+  //! warp would be ``{3, 3, 3, 3, ..., 31}``, the output for the second warp would be
+  //! ``{3, 3, 3, 4, ..., 32}``, etc.
+  //! @endrst
+  //!
+  //! @tparam ScanOp
+  //!   **[inferred]** Binary scan operator type having member
+  //!   `T operator()(const T &a, const T &b)`
+  //!
+  //! @param[in] input
+  //!   Calling thread's input item
+  //!
+  //! @param[out] inclusive_output
+  //!   Calling thread's output item. May be aliased with `input`
+  //!
+  //! @param[in] initial_value
+  //!   Initial value to seed the inclusive scan (uniform across warp)
+  //!
+  //! @param[in] scan_op
+  //!   Binary scan operator
+  template <typename ScanOp>
+  _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& inclusive_output, T initial_value, ScanOp scan_op)
+  {
+    InternalWarpScan internal(temp_storage);
+
+    T exclusive_output;
+    internal.InclusiveScan(input, inclusive_output, scan_op);
+
+    internal.Update(input, inclusive_output, exclusive_output, scan_op, initial_value, Int2Type<IS_INTEGER>());
+  }
+
   //! @rst
   //! Computes an inclusive prefix scan using the specified binary scan functor across the
   //! calling warp. Also provides every thread with the warp-wide ``warp_aggregate`` of
@@ -544,6 +595,66 @@ public:
     InternalWarpScan(temp_storage).InclusiveScan(input, inclusive_output, scan_op, warp_aggregate);
   }
 
+  //! @rst
+  //! Computes an inclusive prefix scan using the specified binary scan functor across the
+  //! calling warp. Also provides every thread with the warp-wide ``warp_aggregate`` of
+  //! all inputs.
+  //!
+  //! * @smemwarpreuse
+  //!
+  //! Snippet
+  //! +++++++
+  //!
+  //! The code snippet below illustrates four concurrent warp-wide inclusive prefix max scans
+  //! within a block of 128 threads (one scan per warp).
+  //!
+  //! .. literalinclude:: ../../../cub/test/catch2_test_warp_scan_api.cu
+  //!     :language: c++
+  //!     :dedent:
+  //!     :start-after: example-begin inclusive-warp-scan-init-value-aggregate
+  //!     :end-before: example-end inclusive-warp-scan-init-value-aggregate
+  //!
+  //! Suppose the set of input ``thread_data`` across the block of threads is
+  //! ``{1, 1, 1, 1, ..., 1}``. For initial value equal to 3, the corresponding output
+  //! ``thread_data`` for a sum operation in the first warp would be
+  //! ``{4, 5, 6, 7, ..., 35}``, the output for the second warp would be
+  //! ``{4, 5, 6, 7, ..., 35}``, etc.  Furthermore,  ``warp_aggregate`` would be assigned
+  //! ``32`` for threads in each warp.
+  //! @endrst
+  //!
+  //! @tparam ScanOp
+  //!   **[inferred]** Binary scan operator type having member
+  //!   `T operator()(const T &a, const T &b)`
+  //! @param[in] input
+  //!   Calling thread's input item
+  //!
+  //! @param[out] inclusive_output
+  //!   Calling thread's output item. May be aliased with ``input``
+  //!
+  //! @param[in] initial_value
+  //!   Initial value to seed the inclusive scan (uniform across warp). It is not taken
+  //!   into account for warp_aggregate.
+  //!
+  //! @param[in] scan_op
+  //!   Binary scan operator
+  //!
+  //! @param[out] warp_aggregate
+  //!   Warp-wide aggregate reduction of input items.
+  template <typename ScanOp>
+  _CCCL_DEVICE _CCCL_FORCEINLINE void
+  InclusiveScan(T input, T& inclusive_output, T initial_value, ScanOp scan_op, T& warp_aggregate)
+  {
+    InternalWarpScan internal(temp_storage);
+
+    // Perform the inclusive scan operation
+    internal.InclusiveScan(input, inclusive_output, scan_op);
+
+    // Update the inclusive_output and warp_aggregate using the Update function
+    T exclusive_output;
+    internal.Update(
+      input, inclusive_output, exclusive_output, warp_aggregate, scan_op, initial_value, Int2Type<IS_INTEGER>());
+  }
+
   //! @}  end member group
   //! @name Exclusive prefix scans
   //! @{