Tracing reduce, scan, sort in CUB: Looking for Algorithm-Level Insights

[Wilson1211]

After tracing the code to understand how the reduce function is implemented, I ran into some difficulties.

I found that the Reduce function is located in cub/device/device_reduce.cuh, which calls
DispatchReduce<InputIteratorT, OutputIteratorT, OffsetT, ReductionOpT, T>::Dispatch(...).

This, in turn, is defined in cub/device/dispatch/dispatch_reduce.cuh.
Inside Dispatch, it eventually calls:

CubDebug(max_policy.Invoke(ptx_version, dispatch));

The Invoke function looks like this:

template <typename ActivePolicyT>
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) {
  auto wrapped_policy = detail::reduce::MakeReducePolicyWrapper(active_policy);
  if (num_items <= static_cast<OffsetT>(
        wrapped_policy.SingleTile().BlockThreads() * wrapped_policy.SingleTile().ItemsPerThread())) {
    // Small, single tile size
    return InvokeSingleTile(kernel_source.SingleTileKernel(), wrapped_policy);
  } else {
    // Regular size
    return InvokePasses(kernel_source.ReductionKernel(), kernel_source.SingleTileSecondKernel(), wrapped_policy);
  }
}

I kept tracing further, but I haven’t found a clear or intuitive view of how the reduce algorithm is implemented from an high-level perspective.

I understand that many optimizations are applied, but I was wondering if there’s a central piece of the implementation that outlines the core algorithm used for reduction.
Alternatively, is there any documentation or write-up (especially in CCCL) that describes how reduce, scan, and sort are implemented in CUB?

If I’ve missed something important or traced the wrong direction, I’d really appreciate your guidance.

Thank you!

[elstehle]

Thank you for expressing your interest in exploring CUB's codebase, @Wilson1211!

Implementation Details

I can understand that the CUB dispatch mechanism can be overwhelming at times to developers that are new to CUB's codebase. You will find the implementation details in the Agent* of the respective algorithm, e.g., AgentReduce in cub/agent/agent_reduce.cuh.

Call Trace

I think the next step in the call trace is:

The dispatch object's Invoke function is then called with the best policy for the device's PTX version:

The call trace of CUB's dispatch mechanism is outlined in the Dispatch Layer section of CUB's Developer Overview - that may be a good starting point for you.

For DeviceReduce some of the invocations that follow the ones you had already outlined are:

• DispatchReduce::Invoke
• DispatchReduce::InvokePasses
• ... (some setup/init work)
• DeviceReduceKernel
• AgentReduce::ConsumeTiles