[CUB][device] Add a env-based overload of the device segmented reductions primitives

[rbourgeois33]

Description

closes #6673

This Draft PR implements a env-based overload of the device segmented reductions primitives.

@gevtushenko, following our discussion:

Consider adding a static assert that if determinism is specified - it's run-to-run or not-guaranteed.

I have checked that the required determinism, if provided, must not be gpu_to_gpu. As far as I understand, it is equivalent to check that it's run-to-run or not-guaranteed.

I made this check list of what this PR should verify before being merged below. Let me know what I should add to the Sum example to tick the first box :).

Checklist

• We have agreed on the way to implement the env-based device segmented reductions starting with the first one:Sum.
• I have extended the overload to other relevant primitives.
• The documentation is up to date with these changes.
• If requirements are the same as the CUDALibrarySamples repo, I have updated the licence headers and signed my commits.
• I have squashed commits

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

[srinivasyadav18]

Thank you for trying to contribute to CCCL! Apologies for very early review on the draft. I have some suggestions which might help you to get better context. Feel free to ask any questions you have.

[rbourgeois33]

Hello @srinivasyadav18 . Thanks a lot for the feedback :)

I added several tests in the new file cub/test/catch2_test_device_segmented_reduce_env_api.cu as you suggested.

These include three tests to verify that the new API correctly accepts a stream and the two non-gpu_to_gpu determinism arguments from the env.

In the comment header for the new function, I linked to the segmented-reduce-sum-env-stream example, since the determinism arguments are currently unused. I believe this is more useful for the documentation. Perhaps the fact that the determinism arguments are unused should be clearly stated.

I removed my example code because could not stay here long-term. I also didn’t add a tests for passing an mr argument since memory resources are still experimental.

Regarding the complete boiler plate to extract useful info from env, I think it's a great idea ! Less code duplicate is always better, in particular since the env extraction is quite verbose. I subscribed the PR and make changes accordingly in time.

Cheers !

On a side note, looking at

cccl/cub/test/catch2_test_device_reduce_env_api.cu

Line 87 in 38162fa

C2H_TEST("cub::DeviceReduce::Sum accepts stream", "[reduce][env]")

it seems to me that it is an incomplete duplicate of:

cccl/cub/test/catch2_test_device_reduce_env_api.cu

Line 148 in 38162fa

C2H_TEST("cub::DeviceReduce::Sum accepts stream", "[reduce][env]")

since it has the same title and does not check correctness of the result. Perhaps it should be removed.

[srinivasyadav18]

Thanks a lot for working on the feedback. The example test file looks great now.
I left some minor comments. But rest all Looks Good To Me.

[srinivasyadav18]

These two comments are totally optional, feel free to look into.
Please let me know if you have questions. Thanks!

[srinivasyadav18]

/ok to test 05194d7

[rbourgeois33]

Just added another unit test. Let me know if you think it would be useful or not to extend the work to other segmented reductions.

[srinivasyadav18]

Thanks! The test looks good. Require minor changes.
The PR is coming close to merge. Once we have the changes and CI is green we can ship it.

[NaderAlAwar]

Great work @rbourgeois33! Left a comment regarding the API tests, will approve once that is addressed.

[srinivasyadav18]

Just added another unit test. Let me know if you think it would be useful or not to extend the work to other segmented reductions.

I would prefer to just have Sum() in the current PR and other overloads like Reduce() Min() Max() etc.. in the next one, as this would simplify the review process and also helps to reduce boilerplate code as this PR #6622 will be merged soon.

[rbourgeois33]

Hi @NaderAlAwar , @srinivasyadav18 !

I think we have resolved all comments/suggestions and that the PR is ready to be merged.

Let me know if I should squash / sign my commits and/or update some licence headers (I had to for another contribution to the CUDALibrarySamples repo).

Also, thanks for your help ! It's super nice that external contribution to cccl are encouraged like this. I learned a lot !

[NaderAlAwar]

/ok to test a20271b

[srinivasyadav18]

@rbourgeois33 Thank you for working on this.

Let me know if I should squash / sign my commits and/or update some licence headers

Signing commits is optional (not required). We squash the commits during merge, so not required to squash in the PR itself.

[NaderAlAwar]

Great work @rbourgeois33! We still need to create the env based overloads for the other primitives.

[github-actions]

🥳 CI Workflow Results

🟩 Finished in 1h 12m: Pass: 100%/93 | Total: 1d 06h | Max: 1h 01m | Hits: 99%/92903

See results here.