[CUB] Make ReduceByKey run-to-run deterministic - #11032
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. 📝 WalkthroughSummary by CodeRabbit
Walkthrough
ChangesReduce-by-key determinism
Assessment against linked issues
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The change makes floating-point ReduceByKey results deterministic by default while preserving an opt-out path and existing behavior for unaffected operations; the supplied verification shows the covered tests and benchmarks pass, so no actionable merge-blocking risk remains. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 69b6565e-8ae2-4392-844b-369f6c27dc16
📒 Files selected for processing (3)
cub/cub/agent/agent_reduce_by_key.cuhcub/cub/device/dispatch/dispatch_reduce_by_key.cuhcub/test/catch2_test_device_reduce_by_key.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
0720444 to
009b23f
Compare
gevtushenko
left a comment
gevtushenko
left a comment
There was a problem hiding this comment.
Thank you for filing the PR! There are a couple of changes that have to be made before we merge:
| static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; | ||
| static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; | ||
| static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; | ||
| static constexpr bool STABLE_REDUCTION_ORDER = StableReductionOrder; |
There was a problem hiding this comment.
important: reduction order stability is not a tuning parameter. The only side-effect of changing a tuning parameter would be performance, while if user adjusts reduction order stability, results are observable in the function output.
Let's keep StableReductionOrder out of agent_reduce_by_key_policy. Instead, use DeviceScan for inspiration on how to pass it directly to dispatch:
cccl/cub/cub/device/device_scan.cuh
Line 167 in b7aaea6
| AccumT, | ||
| streaming_context_t>; | ||
| streaming_context_t, | ||
| /* StableReductionOrder */ true>; |
There was a problem hiding this comment.
critical: performance impact is significant. We have to:
- relax stability requirement when it's not needed: for certain combinations of data types and operators (say, reducing integers with plus operator, or finding min/max values on primitive types) existing implementation provides deterministic result. In this case, we can avoid performance regressions and avoid requirement on stability of reduction tree, see:
cccl/cub/cub/device/device_reduce.cuh
Lines 309 to 310 in b7aaea6
- give users opt out if determinism is not required: as written, determinism is always required. If user didn't need run-to-run determinism, performance regression wouldn't be welcome. We should connect this flag to environments API such that user can opt out of default run-to-run determinism and static assert if user requested gpu-to-gpu determinism and we couldn't provide it, see:
cccl/cub/cub/device/device_reduce.cuh
Lines 297 to 302 in b7aaea6
- modify default reduce by key policy selector to enable tuning: we might want to re-tune the algorithm based on determinism requirement. Let's make a change analogous to how we treat this on scan side, see:
There was a problem hiding this comment.
Done. updated results are in the PR description
| d_values_in, | ||
| d_aggregates_out, | ||
| d_num_runs_out, | ||
| fp64_sum{}, |
There was a problem hiding this comment.
important: it should be cuda::std::plus, otherwise it won't compile when you introduce static assert on the interface side.
There was a problem hiding this comment.
Done. The regression now uses cuda::std::plus.
| c2h::host_vector<std::uint64_t> keys(num_items); | ||
| c2h::host_vector<double> values(num_items); | ||
| std::mt19937_64 rng{42}; | ||
| std::uniform_real_distribution<double> distribution{-1.0, 1.0}; | ||
|
|
||
| for (std::size_t i = 0; i < num_items; ++i) | ||
| { | ||
| keys[i] = i < long_run ? 0 : 1 + (i - long_run) / short_run; | ||
| values[i] = i % 7 == 0 ? distribution(rng) * 1e-2 : 0.0; | ||
| static_cast<void>(rng()); // Match the issue's integer-control RNG consumption. | ||
| } | ||
|
|
||
| c2h::device_vector<std::uint64_t> keys_in = keys; | ||
| c2h::device_vector<double> values_in = values; |
There was a problem hiding this comment.
question: can we simplify this to just c2h::gen(C2H_SEED(2), values);? I assume the test will fail regarless of the pattern.
There was a problem hiding this comment.
I tried the generator approach, but the usable variants either produced NaNs or did not reproduce the failure. I kept the smallest finite pattern that does.
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6bf26672-f5de-49da-94cc-8307076a1198
📒 Files selected for processing (7)
cub/cub/agent/agent_reduce_by_key.cuhcub/cub/device/device_reduce.cuhcub/cub/device/dispatch/dispatch_reduce_by_key.cuhcub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuhcub/test/catch2_test_device_reduce_by_key.cucub/test/catch2_test_device_reduce_env_api.cucub/test/test_device_reduce_by_key_determinism_fail.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
2 participants
@nanan-nvidia @guillaume-michel
Description
closes #9995
DeviceReduce::ReduceByKeycan produce different FP64 result bits across identical calls when a segment spans multiple tiles. Cross-tile partial reductions are combined in an order that can depend on tile completion.This change uses the stable ordering already supported by
TilePrefixCallbackOpfor floating-pointcuda::std::plus. The stable-order requirement is passed separately from the tuning policy. The environment overload defaults torun_to_run, andnot_guaranteedopts out. Integral operations with known CUDA operators and primitive min/max keep the existing path. Unsupported determinism requests fail at compile time. The legacy temporary-storage overload also uses stable ordering for floating-pointcuda::std::plus. The shared kernel defaults to the previous ordering, so RLE and the deprecated dispatcher retain their current behavior. The virtual shared-memory sizing path uses the same flag as the kernel.The regression reproduces the reported 65,536-item input with a 50,000-item first run and compares that aggregate bit-for-bit across 20 identical calls. It uses
cuda::std::plusand runs in one generated test shard. Existing environment tests coverrun_to_runandnot_guaranteed, and one compile-fail test covers the two unsupported determinism cases. Existing parameterized tests cover the run count and complete output correctness.Verification
cuda::std::plusproduces one aggregate bit pattern across all 200 calls on an RTX 5090.Performance
Measurements used CUDA 13.3 on the RTX 5090. The benchmark used one
2^24-item workload with 256-item runs and three order-balanced process pairs.cuda::std::plus, defaultrun_to_runcuda::std::plus,not_guaranteedcuda::std::plus, defaultcuda::minimum, defaultThe last three cases stay within their measured sample variation.
Checklist