1
// 版本 A：纯归约
2
template <typename Input, typename ReductionOp,
3
          typename ValueT, typename AccumT>
4
_CCCL_DEVICE _CCCL_FORCEINLINE AccumT
5
ThreadReduce(const Input& input, ReductionOp reduction_op);
6

7
// 版本 B：带初始值（CUB 叫 prefix）
8
template <typename Input, typename ReductionOp,
9
          typename PrefixT, typename ValueT, typename AccumT>
10
_CCCL_DEVICE _CCCL_FORCEINLINE AccumT
11
ThreadReduce(const Input& input, ReductionOp reduction_op,
12
             PrefixT prefix)
13
{
14
  constexpr int length = static_size_v<Input>;
15
  AccumT array[length + 1];           // 比 input 多 1 格
16
  array[0] = prefix;                  // 第 0 格放 prefix
17
  _CCCL_PRAGMA_UNROLL_FULL()
18
  for (int i = 0; i < length; ++i)
19
    array[i + 1] = input[i];          // input 整体后移一格
20
  return cub::ThreadReduce(array, reduction_op); // 复用版本 A
21
}

1
// cub/thread/thread_reduce.cuh 第 286-301 行
2
template <typename AccumT, typename Input, typename ReductionOp>
3
_CCCL_DEVICE _CCCL_FORCEINLINE AccumT
4
ThreadReduceBinaryTree(const Input& input, ReductionOp reduction_op)
5
{
6
  constexpr auto length = static_size_v<Input>;
7
  auto array            = cub::detail::to_array<AccumT>(input);  // 拷贝到局部数组
8
  _CCCL_PRAGMA_UNROLL_FULL()
9
  for (int i = 1; i < length; i *= 2)              // 外层：步长 1,2,4,...
10
  {
11
    _CCCL_PRAGMA_UNROLL_FULL()
12
    for (int j = 0; j + i < length; j += i * 2)    // 内层：相邻对合并
13
      array[j] = reduction_op(array[j], array[j + i]);
14
  }
15
  return array[0];
16
}

1
// cub/thread/thread_reduce.cuh 第 303-319 行
2
template <typename AccumT, typename Input, typename ReductionOp>
3
_CCCL_DEVICE _CCCL_FORCEINLINE AccumT
4
ThreadReduceTernaryTree(const Input& input, ReductionOp reduction_op)
5
{
6
  constexpr auto length = static_size_v<Input>;
7
  auto array            = cub::detail::to_array<AccumT>(input);
8
  _CCCL_PRAGMA_UNROLL_FULL()
9
  for (int i = 1; i < length; i *= 3)              // 外层：步长 1,3,9,...
10
  {
11
    _CCCL_PRAGMA_UNROLL_FULL()
12
    for (int j = 0; j + i < length; j += i * 3)    // 内层：每 3 个合并
13
    {
14
      auto value = reduction_op(array[j], array[j + i]);
15
      array[j]   = (j + i * 2 < length)            // 边界检查：第三个元素存在吗？
16
                     ? reduction_op(value, array[j + i * 2])
17
                     : value;
18
    }
19
  }
20
  return array[0];
21
}

1
// cub/thread/thread_reduce.cuh 第 410-471 行（简化，只保留核心分支）
2
template <typename Input, typename ReductionOp, typename ValueT, typename AccumT>
3
_CCCL_DEVICE _CCCL_FORCEINLINE AccumT
4
ThreadReduce(const Input& input, ReductionOp reduction_op)
5
{
6
  constexpr auto length = static_size_v<Input>;
7

8
  // 1. 只有 1 个元素：直接返回
9
  if constexpr (length == 1)
10
    return static_cast<AccumT>(input[0]);
11

12
  // 2. 不支持 SIMD 或 元素 >= 8 字节 → 兜底
13
  if constexpr (!is_simd_enabled_cuda_operator<ReductionOp, ValueT>
14
                || sizeof(ValueT) >= 8)
15
    return ThreadReduceSequential<AccumT>(input, reduction_op);
16

17
  // 3. SM90 + int16 + min/max + length >= 10 → DPX 硬件 SIMD
18
  if constexpr (enable_sm90_simd_reduction_v<ValueT, ReductionOp, length>)
19
    NV_IF_TARGET(NV_PROVIDES_SM_90,
20
      (return ThreadReduceSimd(input, reduction_op);))
21

22
  // 4. SM80 + half/bf16 + min/max/sum + length >= 4 → half2 SIMD
23
  if constexpr (enable_sm80_simd_reduction_v<ValueT, ReductionOp, length>)
24
    NV_IF_TARGET(NV_PROVIDES_SM_80,
25
      (return ThreadReduceSimd(input, reduction_op);))
26

27
  // 5. length >= 6 且 SM50+ 整数 plus/bitwise → TernaryTree（ILP 最优）
28
  if constexpr (length >= 6 && enable_ternary_reduction_sm50_v<ValueT, ReductionOp>)
29
    return ThreadReduceTernaryTree<PromT>(input, reduction_op);
30

31
  // 6. 默认：BinaryTree
32
  return ThreadReduceBinaryTree<PromT>(input, reduction_op);
33
}

1
// 一次蝴蝶交换
2
template <typename _Tp, typename ReductionOp>
3
_CCCL_DEVICE _CCCL_FORCEINLINE _Tp
4
ReduceStep(_Tp input, ReductionOp reduction_op, int last_lane, int offset)
5
{
6
  _Tp output = input;
7

8
  // 核心：从 offset 之外的邻居那里"拽"数据过来
9
  _Tp temp = ShuffleDown<LOGICAL_WARP_THREADS>(output, offset, last_lane, member_mask);
10

11
  // 如果这个邻居在有效范围内，就参与归约
12
  if (offset + lane_id <= last_lane)
13
    output = reduction_op(input, temp);
14

15
  return output;
16
}

1
template <bool ALL_LANES_VALID, typename ReductionOp, int STEP>
2
_CCCL_DEVICE _CCCL_FORCEINLINE T
3
ReduceStep(T input, int valid_items, ReductionOp reduction_op, constant_t<STEP>)
4
{
5
  constexpr int OFFSET = 1 << STEP;
6

7
  temp_storage.reduce[lane_id] = input;     // 写共享内存
8
  __syncwarp(member_mask);                  // warp 内同步
9
  // 读邻居的数据
10
  if (lane_id + OFFSET < valid_items)
11
  {
12
    T peer_addend = temp_storage.reduce[lane_id + OFFSET];
13
    input = reduction_op(input, peer_addend);
14
  }
15
  __syncwarp(member_mask);
16
  return ReduceStep<...>(input, ...);        // 模板递归下一步
17
}

1
template <typename T,
2
          int BlockDimX,
3
          BlockReduceAlgorithm Algorithm = BLOCK_REDUCE_WARP_REDUCTIONS,
4
          int BlockDimY = 1,
5
          int BlockDimZ = 1>
6
class BlockReduce
7
{
8
private:
9
  static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ;
10

11
  // 四个具体实现类的别名：
12
  using WarpReductions = detail::BlockReduceWarpReductions<T, BlockDimX, BlockDimY, BlockDimZ>;
13
  using WarpReductionsNondeterministic = detail::BlockReduceWarpReductions<T, BlockDimX, BlockDimY, BlockDimZ, false>;
14
  using RakingCommutativeOnly = detail::BlockReduceRakingCommutativeOnly<T, BlockDimX, BlockDimY, BlockDimZ>;
15
  using Raking = detail::BlockReduceRaking<T, BlockDimX, BlockDimY, BlockDimZ>;
16

17
  // 编译期三级决策树（注意 _If 是 cuda::std:: 的类型萃取器）
18
  using InternalBlockReduce =
19
    ::cuda::std::_If<Algorithm == BLOCK_REDUCE_WARP_REDUCTIONS,
20
                     WarpReductions,
21
                     ::cuda::std::_If<Algorithm == BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC,
22
                                      WarpReductionsNondeterministic,
23
                                      ::cuda::std::_If<Algorithm == BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY,
24
                                                       RakingCommutativeOnly,
25
                                                       Raking>>>;
26
};

1
// 只有前 valid_items 个"线程槽位"有数据
2
int thread_offset = lane_id;
3

4
// 第一个有效元素直接赋值（不能用 reduction_op）
5
if (thread_offset < valid_items)
6
{
7
    thread_aggregate = transform_op(d_in[block_offset + thread_offset]);
8
    thread_offset += NumThreads;
9
}
10

11
// 后续有效元素用 reduction_op 累加
12
while (thread_offset < valid_items)
13
{
14
    thread_aggregate = reduction_op(thread_aggregate, transform_op(d_in[block_offset + thread_offset]));
15
    thread_offset += NumThreads;
16
}

1
                    N 个元素（可能 1 亿）
2
                         │
3
          ┌──────────────┼──────────────┐
4
          │              │              │
5
       Block 0       Block 1  ...   Block M-1    ← 第一趟（并行）
6
      归约 [0..K)   归约 [K..2K)     归约 [...N)
7
          │              │              │
8
      部分结果 0     部分结果 1     部分结果 M-1
9
          │              │              │
10
          └──────────────┼──────────────┘
11
                         │
12
                    1 个 Block                      ← 第二趟（串行）
13
                    归约 M 个部分结果
14
                         │
15
                    最终结果 1 个

1
// dispatch_reduce.cuh — 顶层入口
2
auto dispatch(...)
3
{
4
    reduce_policy active_policy = policy_selector(compute_capability);
5

6
    // 一个 tile 就装下了 → 直接单 block 搞定
7
    if (num_items <= single_tile_capacity)
8
        return invoke_single_tile();      // DeviceReduceSingleTileKernel<<<1, N, 0>>>
9

10
    // 大问题：两趟
11
    return invoke_passes();
12
}
13

14
// 第一趟：多 block 并行归约
15
auto invoke_passes(...)
16
{
17
    // 1. 算需要多少 block
18
    int max_blocks = sm_occupancy * sm_count * subscription_factor;
19

20
    // 2. GridEvenShare 把数据均匀分给各 block
21
    GridEvenShare even_share;
22
    even_share.DispatchInit(num_items, max_blocks, tile_size);
23
    int grid_size = even_share.grid_size;
24

25
    // 3. 分配临时存储：每个 block 放一个部分结果
26
    AccumT* d_block_reductions = allocate(max_blocks * sizeof(AccumT));
27

28
    // 4. 发射第一趟 kernel
29
    DeviceReduceKernel<<<grid_size, block_threads, 0, stream>>>(
30
        d_in, d_block_reductions, num_items, even_share, reduction_op, transform_op);
31

32
    // ===== 第二趟：单 block 汇总 M 个部分结果 =====
33
    DeviceReduceSingleTileKernel<<<1, block_threads, 0, stream>>>(
34
        d_block_reductions, d_out, grid_size, reduction_op, init, identity);
35
}

1
struct reduce_policy
2
{
3
    agent_reduce_policy reduce;                // 正常数据量
4
    agent_reduce_policy single_tile;           // 小数据量（一个 tile 能装下）
5
    agent_reduce_policy reduce_nondeterministic; // 非确定性快速路径
6
};
7

8
// agent_reduce_policy — 就是 AgentReducePolicy 的运行期版本
9
struct agent_reduce_policy
10
{
11
    int block_threads;        // 每个 block 多少线程
12
    int items_per_thread;     // 每线程处理多少元素
13
    int vec_size;             // 向量化加载宽度
14
    BlockReduceAlgorithm block_algorithm;  // BLOCK_REDUCE_RAKING 还是 WARP_REDUCTIONS
15
    CacheLoadModifier load_modifier;       // LOAD_LDG（绕过 L1）还是 LOAD_DEFAULT
16
};