BlockAllReduce: dedicated broadcast slot and a single-warp path

[0z5a]

The broadcast value shared ownership of a union slot with the reduction scratch (issue #3917), so
every participating thread had to re-read it once the reduction completed. Giving the broadcast its
own slot removes that second read, and a single-warp block can skip the cross-warp reduction
entirely, reducing through WarpReduce plus ShuffleIndex<32>.

Multi-warp blocks keep the existing raking path, so only the single-warp configuration changes
behaviour. Measured on an L20 (sm89, CUDA 12.8), ABBA-paired against the unpatched block primitive:
1.80x for a single-warp block, parity for multi-warp blocks. Correctness: 15/15 BlockAllReduce
cases plus the shared-agent regression suite.

Rework in progress (per review)

@fbusato's review asks for a different shape than the one currently pushed, and the rework is the
next commit rather than a tweak:

• drop the new class and the split broadcast slot — BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC
  already carries the broadcast contract, so nothing new is needed to express it;
• extend the cross-warp step in block/specializations/block_reduce_warp_reductions.cuh so every
  thread folds in temp_storage.warp_aggregates[warp_idx] as it is produced, instead of thread 0
  assembling the total and every thread re-reading it afterwards;
• that removes the extra read for every block size, so the patch no longer hinges on the
  block_size == 32 case.

Related: #9322 proposed BlockReduceBroadcast, i.e. a new primitive for the same semantics, and was
closed without prejudice after passing CI. This branch deliberately does not re-land that shape.

Measurement environment

The numbers above were taken on an L20 (sm89) with driver 570.86.10 / CUDA 12.8. The machine has
since been upgraded to driver 595.91.07, so the figures are being re-measured and this section
will carry the updated table; treat the current numbers as driver-570 references.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[fbusato]

I'm not sure if it is worth introducing a new class only to express the broadcast semantic.
The algorithm enum BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC already ensures broadcast logic.

I would try to extend the logic in warp-based block reduction instead of relying on synchronization.

cccl/cub/cub/block/specializations/block_reduce_warp_reductions.cuh

Lines 158 to 170 in 2fcc3f0

	// Update total aggregate in warp 0, lane 0
	if (linear_tid == 0)
	{
	_CCCL_PRAGMA_UNROLL_FULL()
	for (int warp_idx = 1; warp_idx < warps; ++warp_idx)
	{
	if (FullTile || (warp_idx * logical_warp_size < num_valid))
	{
	T addend = temp_storage.warp_aggregates[warp_idx];
	warp_aggregate = reduction_op(warp_aggregate, addend);
	}
	}
	}

The optimization for block_size=32 is nice but a bit narrow as it is more an edge case.

I also want to mention that a previous tentative for the same issue was made in #9322

[0z5a]

I'm not sure if it is worth introducing a new class only to express the broadcast semantic.
The algorithm enum BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC already ensures broadcast logic.

I would try to extend the logic in warp-based block reduction instead of relying on synchronization.

cccl/cub/cub/block/specializations/block_reduce_warp_reductions.cuh

Lines 158 to 170 in 2fcc3f0

	// Update total aggregate in warp 0, lane 0
	if (linear_tid == 0)
	{
	_CCCL_PRAGMA_UNROLL_FULL()
	for (int warp_idx = 1; warp_idx < warps; ++warp_idx)
	{
	if (FullTile || (warp_idx * logical_warp_size < num_valid))
	{
	T addend = temp_storage.warp_aggregates[warp_idx];
	warp_aggregate = reduction_op(warp_aggregate, addend);
	}
	}
	}

The optimization for block_size=32 is nice but a bit narrow as it is more an edge case.

I also want to mention that a previous tentative for the same issue was made in #9322

okay

[0z5a]

Done, and the direction was right — no new primitive was needed.

ApplyWarpAggregates in block_reduce_warp_reductions.cuh now folds the published per-warp
aggregates in every thread, walking them in one fixed order, so the block aggregate is already
identical across the block and no broadcast pass is needed. The new class and the split broadcast
slot are gone; the branch is now a single change to that one file.

One constraint worth flagging: the fold order has to be fixed (all threads start from
warp_aggregates[0]) rather than each thread starting from its own warp aggregate. The latter would
let different warps accumulate floating point partials in different orders, so the result would no
longer be bit-identical across the block — the property that made the broadcast necessary in the
first place. Reading the same shared address per step is served as a hardware broadcast.

Measured on L20 (sm89), driver 595.91.07, counting threads whose value differs from thread 0:

block size	upstream	this change
32	240	0
64	744	0
128	1760	0
256	3800	0
512	7888	0
1024	16072	0

The total is correct in every configuration (wrong=0). A launch-bound microbenchmark shows parity
at BT=128/256 and 0.018-0.024 ms vs 0.022-0.023 ms at BT=1024, so I am claiming the contract change
rather than a measured speedup; a realistic workload that consumes the aggregate (softmax-style
normalisation) is the next measurement if you want a perf number.

Also checked #9322 as suggested: it proposed BlockReduceBroadcast, the same new-primitive shape,
and was closed without prejudice after passing CI — this branch deliberately does not re-land it.

Added the realistic measurement: a block-normalise kernel (every thread divides its own value by
the block sum, i.e. the pattern that forced the caller-side publish/barrier/reload), n = 4M floats,
median of six runs, L20 / driver 595.91.07:

	BT=128	BT=256	BT=512
caller-side broadcast	0.044 ms	0.044 ms	0.048 ms
this change	0.043 ms	0.044 ms	0.061 ms

That is parity — the workload is bandwidth bound (32 MB per iteration at roughly 730 GB/s), so the
removed synchronisation hides behind the memory traffic. To be explicit: I am not claiming a
speedup here. What this change buys is the contract (the aggregate is valid in every thread, for
every block size, without a second pass) and the fact that no new primitive is needed for it.

Regression on the same machine, against the refactored header: 120/120 shared-agent cases and 15/15
BlockAllReduce cases pass, and the deterministic path is exercised by both (block sizes 32 through
1024 in the uniformity test above).

---

Cross-architecture check. Same measurement on an RTX 5090 (sm120, driver 580.82.07) against the
CCCL bundled with CUDA 13.0 — which is a different, older implementation of this same function
(compile-time recursion instead of the loop in 3.6.0). Upstream counts are bit-identical to the L20
column, for block sizes 32/64/128/256/512/1024:

	32	64	128	256	512	1024
L20, sm89, CCCL 3.6.0	240	744	1760	3800	7888	16072
5090, sm120, CUDA 13.0 CCCL	240	744	1760	3800	7888	16072
5090 after the fix	0	0	0	0	0	0

The counts depend only on the block/grid geometry, which is why two architectures and two
implementations agree exactly: "the aggregate is valid in thread 0 only" is a design invariant, not
an sm89 artefact. Porting the fix to the older implementation needed nothing beyond replacing the
if (linear_tid == 0) gate with the same in-every-thread fold in one fixed order starting from
warp_aggregates[0], and the totals stay correct in every configuration. That is also the reason I
expect the change to be version-independent rather than tuned to one CCCL revision.