Repository navigation
[CUB] Make ReduceByKey run-to-run deterministic - #11032
ahmedtaha100 wants to merge 11 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (8)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 SummarySummary by CodeRabbit
Walkthrough
ChangesReduce-by-key determinism
Assessment against linked issues
Suggested reviewers: Merge Risk: ⚪ Minimal · up to ReduceByKey now provides stable ordering for supported floating-point additions while preserving existing paths for other supported operations. Multi-tile deterministic coverage and compile-time validation are included, with no current merge-blocking risk identified. 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.
| auto* const d_keys_in = thrust::raw_pointer_cast(keys_in.data()); | ||
| auto* const d_values_in = thrust::raw_pointer_cast(values_in.data()); | ||
| auto* const d_unique_out = thrust::raw_pointer_cast(unique_out.data()); | ||
| auto* const d_aggregates_out = thrust::raw_pointer_cast(aggregates_out.data()); | ||
| auto* const d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data()); |
There was a problem hiding this comment.
nit, for better readability
| auto* const d_keys_in = thrust::raw_pointer_cast(keys_in.data()); | |
| auto* const d_values_in = thrust::raw_pointer_cast(values_in.data()); | |
| auto* const d_unique_out = thrust::raw_pointer_cast(unique_out.data()); | |
| auto* const d_aggregates_out = thrust::raw_pointer_cast(aggregates_out.data()); | |
| auto* const d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data()); | |
| auto* d_keys_in = thrust::raw_pointer_cast(keys_in.data()); | |
| auto* d_values_in = thrust::raw_pointer_cast(values_in.data()); | |
| auto* d_unique_out = thrust::raw_pointer_cast(unique_out.data()); | |
| auto* d_aggregates_out = thrust::raw_pointer_cast(aggregates_out.data()); | |
| auto* d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data()); |
There was a problem hiding this comment.
applies to other occurrences as well
There was a problem hiding this comment.
Alternatively, as the pointer nature is clear from the cast:
| auto* const d_keys_in = thrust::raw_pointer_cast(keys_in.data()); | |
| auto* const d_values_in = thrust::raw_pointer_cast(values_in.data()); | |
| auto* const d_unique_out = thrust::raw_pointer_cast(unique_out.data()); | |
| auto* const d_aggregates_out = thrust::raw_pointer_cast(aggregates_out.data()); | |
| auto* const d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data()); | |
| const auto d_keys_in = thrust::raw_pointer_cast(keys_in.data()); | |
| const auto d_values_in = thrust::raw_pointer_cast(values_in.data()); | |
| const auto d_unique_out = thrust::raw_pointer_cast(unique_out.data()); | |
| const auto d_aggregates_out = thrust::raw_pointer_cast(aggregates_out.data()); | |
| const auto d_num_runs_out = thrust::raw_pointer_cast(num_runs_out.data()); |
This way one does not have to think about if the const could be about the pointee.
There was a problem hiding this comment.
used const auto for all six pointer declarations in this test, including temporary storage, following pauleonix’s suggestion
| #include <cub/device/device_reduce.cuh> | ||
|
|
||
| #include <cuda/std/__functional/operations.h> | ||
| #include <cuda/std/bit> |
There was a problem hiding this comment.
| #include <cuda/std/bit> | |
| #include <cuda/std/__bit/bit_cast.h> |
| constexpr std::size_t num_items = 65'536; | ||
| constexpr std::size_t long_run = 50'000; | ||
| constexpr std::size_t short_run = 64; | ||
| constexpr int num_repetitions = 20; |
There was a problem hiding this comment.
maybe 10 repetitions would be sufficient
There was a problem hiding this comment.
the test with 10 repetitions missed it in one baseline run, even 20 can occasionally miss it. I kept 20 for more chances to catch the bug
| }; | ||
|
|
||
| template <class ReductionOpT, class AccumT, class KeyT> | ||
| template <class ReductionOpT, class AccumT, class KeyT, bool StableReductionOrder> |
There was a problem hiding this comment.
Important: The new template parameter seems unused. Why is it needed? I think it can be removed again.
There was a problem hiding this comment.
I added it following @gevtushenko’s earlier request to support tuning based on determinism
| { | ||
| using vsmem_helper_t = detail::vsmem_helper_default_fallback_policy_t< | ||
| using vsmem_helper_t = detail::reduce_by_key::vsmem_helper_t< | ||
| false, |
There was a problem hiding this comment.
Q: why are we passing false here and not StableReductionOrder?
There was a problem hiding this comment.
this is the deprecated dispatcher, which has no StableReductionOrder parameter, false preserves its existing ordering, the modern dispatcher forwards the flag
There was a problem hiding this comment.
@ahmedtaha100 please use an inline comment to indicate the name of the template parameter /*StableReductionOrder=*/false instead of a bare literal
| AccumT, | ||
| streaming_context_t>); | ||
| streaming_context_t, | ||
| false>); |
| * Whether to use a stable reduction order across tiles | ||
| */ | ||
| template <typename AgentReduceByKeyPolicyT, | ||
| bool StableReductionOrder, |
There was a problem hiding this comment.
Nit: let's please keep the ordering of the template parameters consistent, so this should go at the end of the template parameter list
| { | ||
| using vsmem_helper_t = detail::vsmem_helper_default_fallback_policy_t< | ||
| using vsmem_helper_t = detail::reduce_by_key::vsmem_helper_t< | ||
| false, |
There was a problem hiding this comment.
@ahmedtaha100 please use an inline comment to indicate the name of the template parameter /*StableReductionOrder=*/false instead of a bare literal
| "gpu_to_gpu deterministic reduce-by-key requires integral types with known operators or " | ||
| "primitive types with min/max"); | ||
|
|
||
| static constexpr bool stable_reduction_order = is_run_to_run_required && is_fp_plus_op; |
There was a problem hiding this comment.
| static constexpr bool stable_reduction_order = is_run_to_run_required && is_fp_plus_op; | |
| constexpr bool stable_reduction_order = is_run_to_run_required && is_fp_plus_op; |
No need to make it static (and below as well)
pauleonix
left a comment
pauleonix
left a comment
There was a problem hiding this comment.
Thank you for your contribution! Up to nits this looks good!
@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. RLE and the deprecated dispatcher explicitly select the previous ordering and retain their current behavior. The virtual shared-memory sizing path uses the same flag as the kernel.The regression uses a 262,144-item input with a 200,000-item first run, checks the run count, and compares all output keys and aggregate bits 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.types_0only).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