You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[CUB] Support stream-aware two-phase launch wrappers - #12379
This PR is a focused part of #11374 that modifies launch helper macros to support passing execution stream as last argument of two-phase API.
Infer trailing stream-like arguments through cuda::get_stream and preserve ordering with test setup across all launch modes:
TEST_LAUNCH=0: forward the supplied stream to the host API and synchronize it before returning. Preserve omitted stream arguments.
TEST_LAUNCH=1: enqueue the parent kernel on the supplied host stream, preserving ordering with preparation work; invoke the API on the device side default stream; and wait for the parent and its child kernels to complete before returning.
TEST_LAUNCH=2: capture on the supplied non-default stream, or a temporary stream when the default stream is requested. Launch the graph on the requested stream and synchronize it before returning.
Use checked C2H buffers for temporary storage and scalar results. Fix the device-side memory resource to return the allocation's starting offset.
Add stream-forwarding and explicit default-stream tests, and update the launch-wrapper usage documentation.
Expand catch2_test_launch_wrapper.cu to test both two-phase and environment-based wrappers using
a shared synthetic CUB-style algorithm. Both interfaces reuse storage sizing, execution and stream validation,
graph-capture checks, and kernel dispatch.
Environment-wrapper tests cover host, device, and graph launches with a non-default stream or no stream property,
plus explicit default streams for host and device launches. They exercise ordinary and templated wrappers,
allocation accounting, and execution across multiple blocks.
Replace device vectors and Thrust-based result verification with checked C2H device buffers and host-side
comparisons. Factor allocation and asynchronous preparation into a shared fixture, centralize result checks,
and restore expected-stream state with a scope guard. Check stream completion before readback so verification does not mask incomplete work.
Checklist
New or existing tests cover these changes.
The documentation is up to date with these changes.
Infer trailing stream-like arguments through cuda::get_stream and preserve
ordering with test setup across all launch modes:
- TEST_LAUNCH=0: forward the supplied stream to the host API and synchronize
it before returning. Preserve omitted stream arguments.
- TEST_LAUNCH=1: synchronize the supplied host stream before launching the
parent kernel, invoke the API on the device-side default stream, and
wait for device execution to complete.
- TEST_LAUNCH=2: capture on the supplied non-default stream, or a temporary
stream when the default stream is requested. Launch the graph on the
requested stream and synchronize it before returning.
Use checked C2H buffers for temporary storage and scalar results. Fix the
device-side memory resource to return the allocation's starting offset.
Add stream-forwarding and explicit default-stream tests, and update the
launch-wrapper usage documentation.
No actionable comments were generated in the recent review.
Summary
Summary by CodeRabbit
Documentation
Updated test-launch guidance for stream selection and ordering across host, device, and graph launches. Clarified environment-based memory resources, allocation sizing, wrapper completion, and input-preparation ordering when an environment does not provide a stream.
Tests
Expanded coverage of stream and environment forwarding, default-stream behavior, launch ordering, memory allocation, completion, and output validation across wrapper variants. Added checks for stream-capture errors and preparation/completion ordering.
Launch helpers now forward optional stream or environment arguments and coordinate ordinary, device-side, graph, and environment launches. Tests cover stream selection, completion, and output. The memory-resource allocation offset and launch-wrapper guidance are also updated.
Changes: Stream-aware launch behavior
Layer / File(s)
Summary
Argument forwarding and stream selection cub/test/catch2_test_launch_helper.h
Wrapper call operators forward arguments. Dispatch identifies stream or environment arguments and selects a fallback stream when needed.
Ordinary, device, and graph launches cub/test/catch2_test_launch_helper.h
Launch paths invoke APIs with selected streams and check completion. Graph launches execute and destroy graph handles through a shared helper.
Environment launch ordering and allocation tracking cub/test/catch2_test_launch_helper.h, cub/test/catch2_test_memory_resources.h
Environment launches use existing or fallback streams and retain allocation tracking. The device-side memory resource returns allocations at the counter’s previous offset.
Tests cover stream forwarding, launch targets, environment wrappers, completion, and output. The documentation describes stream handling, launch ordering, and allocation requirements.
Possibly related issues
NVIDIA/cccl#11613: The documentation adds the expected_allocation_size(...) requirement for device-launch tests.
NVIDIA/cccl#9875: The changes add and test environment-aware launch wrapper behavior, including stream handling.
Suggested reviewers:jacobfaib
Change: Feature
Merge Risk: Low · up to 10b50
Custom dual-interface environments may lose properties or fail to compile. The change is mergeable with awareness of this narrow regression.
- [ ] Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Tip
🎯 Customize your risk assessments
Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.
✅ Automatic approvals, with your rules
Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.
Reviewing files that changed from the base of the PR and between a2bef4a and 10b5013.
Files selected for processing: 1
cub/test/catch2_test_launch_helper.h
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
suggestion: Reject default-stream capture in the graph test.
The explicit-default test does not enforce the documented temporary non-default capture stream. Add a focused check that rejects a default stream in the graph-mode API call.
Reviewing files that changed from the base of the PR and between e4e7d3c and 4c1b48e.
Files selected for processing: 1
docs/cub/developer/test_overview.rst
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
1. since capture_stream is always non-default, make_device_buffer
on it needs no device specified and capture_stream.device()
can be used in simpler API
2. Use nameless template parameter in launch_impl in Graph Mode,
as HasStream template parameter is unused there.
3. Improve readability by replacing capture_stream.get() with
use of `capture_sh` variable populated with capture stream handle
We had a work-around specifically for earlier versions of clang
that do not support CUDA dynamic parallelism. Since we have been
using clang-22 for a while now, removed that work-around
1. Launch wrappers for DeviceLaunch (TEST_LAUNCH == 1) enqueue
device_side_api_launcher_kernel using cuda::launch on
the supplied stream avoiding call to `stream.sync()`;
Both for two-phase invocation as well as for _ENV single
phase invocation.
2. Documentation is updated.
… algorithm
Add environment overloads to the synthetic CUB-style algorithm and exercise
their wrappers across host, device, and graph launches.
Share storage sizing, validation, and kernel dispatch between API forms.
Use checked device buffers, common preparation and result-checking helpers,
and scoped stream expectations. Check stream completion before readback.
When two-phase launch wrappers are given an environment, the
wrapper would only override stream for device launchers and
for temporary graph captures.
A test is added for property preservation for two-phase launch
wrappers.
Documentation is updated as well
Also got rid of cudaStreamCreate single usage by factoring
out make_owning_stream utility and reusing it in launch and
launch_env to create stream as needed.
Reviewing files that changed from the base of the PR and between 2bebb86 and 6784178.
Files selected for processing: 2
cub/test/catch2_test_launch_helper.h
docs/cub/developer/test_overview.rst
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
…pers
Add tests that reproduce failure before the fix. Add the fix by making
wrapped API arguments signature part of wrapper's signature with use
of trailing return type. This permits use of is_invocable_v on the
ActionT to distinguish whether wrapped API would accept supplied
arguments `args...` followed by stream handle.
If it does not, we override_stream in the last argument instead of
appending it after `args...` pack.
Reviewing files that changed from the base of the PR and between 77451a5 and a2bef4a.
Files selected for processing: 3
cub/test/catch2_test_launch_helper.h
cub/test/catch2_test_launch_wrapper.cu
docs/cub/developer/test_overview.rst
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Passing cuda::stream_ref{stream} ended up passing that type to wrapped
API which expected cudaStream_t. With no implicit conversion defined
that broken the build.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Description
This PR is a focused part of #11374 that modifies launch helper macros to support passing execution stream as last argument of two-phase API.
Infer trailing stream-like arguments through
cuda::get_streamand preserve ordering with test setup across all launch modes:TEST_LAUNCH=0: forward the supplied stream to the host API and synchronize it before returning. Preserve omitted stream arguments.TEST_LAUNCH=1: enqueue the parent kernel on the supplied host stream, preserving ordering with preparation work; invoke the API on the device side default stream; and wait for the parent and its child kernels to complete before returning.TEST_LAUNCH=2: capture on the supplied non-default stream, or a temporary stream when the default stream is requested. Launch the graph on the requested stream and synchronize it before returning.Use checked C2H buffers for temporary storage and scalar results. Fix the device-side memory resource to return the allocation's starting offset.
Add stream-forwarding and explicit default-stream tests, and update the launch-wrapper usage documentation.
Expand
catch2_test_launch_wrapper.cuto test both two-phase and environment-based wrappers usinga shared synthetic CUB-style algorithm. Both interfaces reuse storage sizing, execution and stream validation,
graph-capture checks, and kernel dispatch.
Environment-wrapper tests cover host, device, and graph launches with a non-default stream or no stream property,
plus explicit default streams for host and device launches. They exercise ordinary and templated wrappers,
allocation accounting, and execution across multiple blocks.
Replace device vectors and Thrust-based result verification with checked C2H device buffers and host-side
comparisons. Factor allocation and asynchronous preparation into a shared fixture, centralize result checks,
and restore expected-stream state with a scope guard. Check stream completion before readback so verification does not mask incomplete work.
Checklist