[CUB] GPU-to-GPU DeviceReduce with device resident problem size

[NaderAlAwar]

Description

closes #9730
merge first #9722

This extends device resident problem size to GPU-to-GPU determinism.

No SASS changes to any existing benchmarks.

We expect the device-resident problem size to be slightly slower: the host cannot read the problem size, so it must assume it is large, launching a worst-case grid and always taking the two-pass path (there is no single-tile shortcut for small problems). The kernels read the problem size on device and trim the worst-case grid to the size the host would have computed: surplus blocks exit early, and the second pass derives the same block count to consume exactly the partials that were produced. Because the RFA accumulator is partition-independent, results are bitwise identical to the host-provided num_items path (this is verified by the tests, including problem sizes above INT32_MAX, where the two paths chunk the input differently). Compared to the host-provided num_items, the results show an overhead of at most ~7 us, which is most noticeable at the smallest problem sizes, and low single-digit percentages elsewhere.

When the problem-size element is anything other than a signed 32-bit integer we observe larger overheads at mid-to-large sizes (an unsigned 32-bit element cannot guarantee the problem fits a single chunk, so it is widened and takes the same path as 64-bit elements). Unlike the host-provided path, which chunks the input on the host so the kernel only ever uses 32-bit index arithmetic, the deferred kernel performs the equivalent chunking on device. The hot loop remains 32-bit, but carrying the wide problem size across the loop widens the kernel's register footprint relative to the signed-32-bit instantiation (F32: 40 vs. 32 registers, F64: 44 vs. 40 on SM90), which reduces occupancy. Profiling with Nsight Compute confirms this is the cause: achieved occupancy drops from 76% to 56% at F32/2^24, where the kernel is latency-bound, and the effect fades as bandwidth saturates at larger sizes. Users whose problem sizes fit in 31 bits should prefer a signed 32-bit element for the deferred source; a future improvement could use cuda::args::bounds to statically select the fast kernel from a bounded wider source.

GPU-to-GPU reduce

## [0] NVIDIA H200

|  T{ct}  |  OffsetT{ct}  |  Elements{io}  |   Ref Time |   Ref Noise |   Cmp Time |   Cmp Noise |       Diff |   %Diff |  Status  |
|---------|---------------|----------------|------------|-------------|------------|-------------|------------|---------|----------|
|   F32   |      I32      |      2^4       |   5.506 us |       2.90% |  11.662 us |       4.13% |   6.157 us | 111.83% | 🔴 SLOW  |
|   F32   |      I32      |      2^8       |   6.559 us |       2.49% |  11.422 us |       4.44% |   4.862 us |  74.13% | 🔴 SLOW  |
|   F32   |      I32      |      2^12      |  11.308 us |       4.21% |  11.621 us |       3.55% |   0.313 us |   2.77% | 🔵 SAME  |
|   F32   |      I32      |      2^16      |  11.902 us |       3.75% |  12.741 us |       3.53% |   0.838 us |   7.04% | 🔴 SLOW  |
|   F32   |      I32      |      2^20      |  14.240 us |       3.13% |  14.965 us |       2.92% |   0.725 us |   5.09% | 🔴 SLOW  |
|   F32   |      I32      |      2^24      |  49.129 us |       1.21% |  50.498 us |       1.19% |   1.369 us |   2.79% | 🔴 SLOW  |
|   F32   |      I32      |      2^28      | 345.331 us |       0.36% | 345.729 us |       0.41% |   0.398 us |   0.12% | 🔵 SAME  |
|   F32   |      I64      |      2^4       |   5.638 us |       2.72% |  12.048 us |       4.08% |   6.411 us | 113.71% | 🔴 SLOW  |
|   F32   |      I64      |      2^8       |   6.281 us |       2.50% |  11.339 us |       4.06% |   5.058 us |  80.53% | 🔴 SLOW  |
|   F32   |      I64      |      2^12      |  11.096 us |       4.06% |  11.563 us |       3.95% |   0.467 us |   4.21% | 🔴 SLOW  |
|   F32   |      I64      |      2^16      |  12.029 us |       4.28% |  13.003 us |       3.96% |   0.973 us |   8.09% | 🔴 SLOW  |
|   F32   |      I64      |      2^20      |  14.276 us |       2.98% |  15.091 us |       2.94% |   0.815 us |   5.71% | 🔴 SLOW  |
|   F32   |      I64      |      2^24      |  48.844 us |       1.32% |  57.319 us |       1.19% |   8.475 us |  17.35% | 🔴 SLOW  |
|   F32   |      I64      |      2^28      | 345.633 us |       0.34% | 367.383 us |       0.27% |  21.750 us |   6.29% | 🔴 SLOW  |
|   F32   |      I64      |      2^31      |   2.716 ms |       4.80% |   2.787 ms |       3.70% |  70.327 us |   2.59% | 🔵 SAME  |
|   F32   |      I64      |      2^32      |   5.542 ms |       2.72% |   5.795 ms |       3.53% | 253.009 us |   4.57% | 🔴 SLOW  |
|   F64   |      I32      |      2^4       |   5.728 us |       3.93% |  13.085 us |       4.01% |   7.357 us | 128.43% | 🔴 SLOW  |
|   F64   |      I32      |      2^8       |   6.802 us |       2.27% |  12.960 us |       3.45% |   6.158 us |  90.53% | 🔴 SLOW  |
|   F64   |      I32      |      2^12      |  11.267 us |       4.62% |  13.733 us |       3.81% |   2.466 us |  21.88% | 🔴 SLOW  |
|   F64   |      I32      |      2^16      |  12.059 us |       3.84% |  14.390 us |       2.92% |   2.331 us |  19.33% | 🔴 SLOW  |
|   F64   |      I32      |      2^20      |  21.962 us |       2.26% |  23.997 us |       2.03% |   2.035 us |   9.26% | 🔴 SLOW  |
|   F64   |      I32      |      2^24      | 100.249 us |       0.75% | 101.898 us |       0.78% |   1.649 us |   1.65% | 🔴 SLOW  |
|   F64   |      I32      |      2^28      | 602.301 us |       0.74% | 601.516 us |       0.82% |  -0.784 us |  -0.13% | 🔵 SAME  |
|   F64   |      I64      |      2^4       |   5.937 us |       2.60% |  12.204 us |       3.94% |   6.267 us | 105.57% | 🔴 SLOW  |
|   F64   |      I64      |      2^8       |   6.839 us |       2.24% |  12.110 us |       3.29% |   5.271 us |  77.07% | 🔴 SLOW  |
|   F64   |      I64      |      2^12      |  10.686 us |       4.31% |  12.705 us |       2.92% |   2.019 us |  18.90% | 🔴 SLOW  |
|   F64   |      I64      |      2^16      |  11.856 us |       3.56% |  13.320 us |       3.46% |   1.464 us |  12.35% | 🔴 SLOW  |
|   F64   |      I64      |      2^20      |  21.846 us |       2.23% |  23.211 us |       2.02% |   1.365 us |   6.25% | 🔴 SLOW  |
|   F64   |      I64      |      2^24      | 100.489 us |       0.71% |  91.983 us |       0.82% |  -8.507 us |  -8.47% | 🟢 FAST  |
|   F64   |      I64      |      2^28      | 601.346 us |       0.86% | 617.782 us |       0.65% |  16.436 us |   2.73% | 🔴 SLOW  |
|   F64   |      I64      |      2^31      |   4.677 ms |       3.67% |   4.893 ms |       2.77% | 216.016 us |   4.62% | 🔴 SLOW  |
|   F64   |      I64      |      2^32      |   9.424 ms |       3.49% |   9.670 ms |       3.04% | 245.889 us |   2.61% | 🔵 SAME  |

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[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.

[NaderAlAwar]

/ok to test 52aa029

[NaderAlAwar]

/ok to test 4808d52

[NaderAlAwar]

/ok to test 40ee0fc

[NaderAlAwar]

@coderabbitai full review

[coderabbitai]

✅ Action performed

Full review finished.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e6abc1ab-b1ee-4bdd-b3e4-a7c76dbf6921

📥 Commits

Reviewing files that changed from the base of the PR and between d47ca8e and a2147b2.

📒 Files selected for processing (7)

• cub/benchmarks/bench/reduce/deferred_deterministic.cu
• cub/benchmarks/bench/reduce/deterministic.cu
• cub/cub/device/device_reduce.cuh
• cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh
• cub/cub/device/dispatch/kernels/kernel_reduce_deterministic.cuh
• cub/test/catch2_test_device_reduce_deferred.cu
• cub/test/catch2_test_device_reduce_env.cu

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added deterministic reductions with deferred problem sizes (including larger element counts) for GPU-to-GPU execution.
  • Expanded deterministic reduction benchmarking and coverage for float and double.

• Bug Fixes

  • Fixed deterministic behavior for deferred-size edge cases, including empty inputs and very large counts.
  • Improved benchmark robustness by skipping unsupported configurations and out-of-memory cases.

• Tests

  • Added validation for deferred-size deterministic dispatch, immediate-vs-deferred consistency, large-count handling, and deterministic CUDA graph capture/replay.

Walkthrough

Adds deferred num_items support to GPU-to-GPU deterministic reduce by wiring device-resident problem sizes through DeviceReduce, dispatch, kernels, tests, and benchmarks.

Changes

Deferred deterministic reduce support

Layer / File(s)	Summary
DeviceReduce entrypoint removes deferred rejection cub/cub/device/device_reduce.cuh	Removes the deferred num_items rejection for gpu_to_gpu determinism and forwards the num_items dispatch-argument type into the RFA dispatch call.
Deterministic reduce kernels handle deferred num_items cub/cub/device/dispatch/kernels/kernel_reduce_deterministic.cuh	Adds deferred problem-size helpers, updates deterministic kernel problem-size handling and 64-bit chunked accumulation, and adds a deferred single-tile kernel entry point.
Dispatch orchestration for deferred num_items cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh	Adds deferred-parameter support and kernel-side num_items wiring, reworks host-side sizing, selects immediate versus deferred single-tile kernels, updates the default policy selector, and bypasses the single-tile shortcut for deferred problem sizes.
Tests for deferred deterministic dispatch cub/test/catch2_test_device_reduce_deferred.cu, cub/test/catch2_test_device_reduce_env.cu	Adds fixed-grid deterministic helpers and tests for deferred boundary handling, immediate-vs-deferred equality, large deferred counts, and CUDA graph replay; updates environment expectations for the deterministic kernel signature.
Benchmarks for deferred and non-deferred deterministic reduce cub/benchmarks/bench/reduce/deferred_deterministic.cu, cub/benchmarks/bench/reduce/deterministic.cu	Adds a deferred deterministic reduce benchmark and updates the existing benchmark with OffsetT handling, deferred-size guards, execution-tag changes, and expanded size axes.

Assessment against linked issues

Objective	Addressed	Explanation
Implement device-resident problem size support for GPU-to-GPU deterministic reduce [#9730]	✅
Ensure correctness parity between deferred and immediate num_items paths, including large (>32-bit) counts [#9730]	✅
Support CUDA graph capture/replay with deferred counts [#9730]	✅

Suggested reviewers: fbusato, jrhemstad

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6bb7f349-d5b1-402e-af08-6c5790279f17

📥 Commits

Reviewing files that changed from the base of the PR and between d47ca8e and a2147b2.

📒 Files selected for processing (7)

• cub/benchmarks/bench/reduce/deferred_deterministic.cu
• cub/benchmarks/bench/reduce/deterministic.cu
• cub/cub/device/device_reduce.cuh
• cub/cub/device/dispatch/dispatch_reduce_deterministic.cuh
• cub/cub/device/dispatch/kernels/kernel_reduce_deterministic.cuh
• cub/test/catch2_test_device_reduce_deferred.cu
• cub/test/catch2_test_device_reduce_env.cu

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 2h 36m: Pass: 100%/287 | Total: 12d 18h | Max: 2h 35m | Hits: 20%/976945

See results here.