[cudax] Isolate and synchronize reduce scratch

[tpn]

Why this is needed

coop::reduce keeps its temporary storage internally. The mapped warp-group
overload previously indexed that storage using a warp's rank within its group.
That rank is local to each group, but sibling groups have independent
synchronizers and may execute concurrently.

For example, partitioning a seven-warp block into groups of three produces:

physical block warp:  0 1 2 | 3 4 5 | 6
rank within group:    0 1 2 | 0 1 2 | invalid
old scratch slot:     0 1 2 | 0 1 2 | -

Physical warps 0 and 3 could therefore use the same CUB
WarpReduce::TempStorage at the same time. The groups also shared partial and
broadcast storage. A barrier within one group does not synchronize the other
group.

The block, cluster, and grid implementations had a related lifetime problem.
Their scratch storage persists across calls, but some paths returned before
every participant had finished consuming it. A subsequent reduction could
overwrite that storage while the previous root was still reading it.

Mapped groups were also passed by value. Some group objects own synchronizer
state, so destroying the parameter copy could tear down copied synchronization
state independently of the original owner.

Examples

Independently synchronized warp groups shared the same slots

The following pattern creates two valid three-warp groups in one block. The
final warp does not participate.

template <class Config>
__device__ void grouped_reduce(Config config, int* output)
{
  constexpr int group_warps = 3;
  constexpr int group_count = 2;

  cudax::this_block block{config};

  using Barriers =
    cuda::barrier<cuda::thread_scope_block>[group_count];

  __shared__
    cuda::std::aligned_storage_t<sizeof(Barriers), alignof(Barriers)>
      barriers_storage;

  auto& barriers = reinterpret_cast<Barriers&>(barriers_storage);

  cudax::group group{
    cuda::warp,
    block,
    cudax::group_by<group_warps, false>{},
    cudax::barrier_synchronizer{barriers}};

  if (!cuda::gpu_thread.is_part_of(group))
  {
    return;
  }

  const int group_rank = group.template rank_as<int>(block);
  int values[1]        = {group_rank + 1};

  const auto result =
    cudax::coop::reduce(group, values, cuda::std::plus<>{});

  if (result.has_value())
  {
    output[group_rank] = *result;
  }
}

With a 224-thread launch, both groups execute the same reduce
specialization:

const auto config =
  cuda::make_config(cuda::grid_dims<1>(), cuda::block_dims<224>());

Each group numbers its warps from zero. Before this change, physical warp 0 and
physical warp 3 both selected temporary slot 0, physical warps 1 and 4 selected
slot 1, and so on. Since the groups synchronize independently, CUB temporary
storage and partial results could be modified concurrently.

Back-to-back reductions could reuse scratch too early

Ordinary repeated use was enough to expose the missing terminal barriers:

template <class Group>
__device__ void repeated_reduce(const Group& group, int* output)
{
  for (int iteration = 0; iteration < 16; ++iteration)
  {
    int values[1] = {iteration + 1};

    const auto result =
      cudax::coop::reduce(group, values, cuda::std::plus<>{});

    if (result.has_value())
    {
      output[iteration] = *result;
    }
  }
}

This pattern applies to this_block, this_cluster, this_grid, and mapped
groups. Without a terminal collective, non-root participants could enter the
next iteration and overwrite scratch while the previous root was still
consuming it.

Block broadcast could reactivate aliased storage too soon

The block implementation uses a union because the CUB temporary storage and
broadcast value do not need to coexist after the collective finishes:

union Scratch
{
  typename BlockReduce::TempStorage block_reduce;
  T broadcast;
};

The old sequence effectively did this:

const auto result =
  BlockReduce{scratch.block_reduce}.Reduce(values, reduce_op);

if (is_root)
{
  // This aliases block_reduce, but other threads may still be using it.
  scratch.broadcast = result;
}

block.sync_aligned();
return scratch.broadcast;

The fixed sequence separates the storage lifetimes and protects the returned
value from the next invocation:

const auto result =
  BlockReduce{scratch.block_reduce}.Reduce(values, reduce_op);

// Every thread has finished using the CUB storage.
block.sync_aligned();

if (is_root)
{
  scratch.broadcast = result;
}

// Publish the broadcast value.
block.sync_aligned();

const T value = scratch.broadcast;

// Every thread has copied the value before scratch can be reused.
block.sync_aligned();

return value;

What changed

Mapped warp-group reductions now assign storage by physical block-warp rank:

• Each physical warp receives its own CUB temporary-storage slot.
• Partial and broadcast values use separate, correctly aligned raw storage.
• Independent groups can reuse local ranks without sharing storage.
• The group is passed by const&, preserving the owning synchronizer's
  lifetime.

The implementation requires a statically sized, contiguous mapping. It uses the
exact physical warp count when the block extents are static. For runtime block
extents, it uses the architectural bound of 32 warps per block, derived from
CUDA's maximum of 1024 threads per block. Larger thread counts associated with
some architectures describe residency per SM, not the size of one thread
block.

The collective paths now synchronize before scratch is repurposed or reused:

• Block broadcast waits for all CUB users before activating the aliased
  broadcast member.
• Broadcast results are copied into registers before the final
  synchronization.
• Root-only block, cluster, grid, and mapped-group reductions perform a
  terminal collective before returning.
• Mapped groups retain the barrier that protects partial reads and the final
  barrier that protects the next invocation.

These barriers add synchronization cost, but they are required while reduce
owns implicit shared or global scratch.

Follow-ups

Support runtime block extents in this_warp

The separate this_warp overload still computes its shared-memory array bound
directly from static_extent. A runtime block configuration therefore fails
during compilation:

const auto config =
  cuda::make_config(cuda::grid_dims<1>(), cuda::block_dims(128));

template <class Config>
__device__ void warp_reduce(Config config)
{
  cudax::this_warp warp{config};
  int values[1] = {1};

  // The warp size is static, but the containing block extents are dynamic.
  const auto result =
    cudax::coop::reduce(warp, values, cuda::std::plus<>{});
}

This defect predates the mapped-group change. It should be fixed separately
after defining the behavior of a terminal partial warp.

Make temporary storage caller-provided

reduce.cuh already notes that implicit scratch is temporary API design.
Caller-provided storage would make ownership explicit and could remove the
terminal barriers added here.

It should also replace the global grid-partial array. Two concurrent launches
of the same grid-reduction specialization currently refer to the same global
storage:

cuda::launch(stream1, config, grid_reduce_kernel{}, input1, output1);
cuda::launch(stream2, config, grid_reduce_kernel{}, input2, output2);

The streams may execute concurrently, allowing the two grids to overwrite each
other's partial results.

Define the contract for future strided mappings

The currently supported contiguous mappings assign unit ranks in ascending
physical-warp order. The physical-slot calculation relies on that property.

Before adding strided or permuted mappings, the mapping API should expose a
physical base rank or a dedicated order-preserving trait rather than treating
contiguity alone as sufficient.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.