Fix extracting CUDA stream in `cub::DeviceTransform` #7239

bernhardmgruber · 2026-01-14T23:13:11Z

#6204 changed cub::DeviceTransform APIs from taking cudaStream_t to environments. Special handling preserved support for cudaStream_t. However, user-provided stream types with conversion operators to cudaStream_t were now queried as environments, failing to return a stream.

This PR treats any type that is convertible to a cudaStream_t not as an environment and extracts the underlying stream.

Fixes NVBug 5813928

cub/cub/device/device_transform.cuh

bernhardmgruber · 2026-01-14T23:13:35Z

cub/cub/device/device_transform.cuh

      num_items,
      ::cuda::std::move(transform_op),
-      get_stream(env));
+      ::cuda::std::move(env));


Drive-by fix

bernhardmgruber · 2026-01-14T23:19:45Z

cub/test/catch2_test_device_transform_env.cu

+  auto run = [&](auto streamish) {
+    cub::DeviceTransform::Transform(cuda::std::make_tuple(a, b), result.begin(), num_items, _1 + _2, streamish);
+  };
  SECTION("raw stream")
  {
-    cub::DeviceTransform::Transform(cuda::std::make_tuple(a, b), result.begin(), num_items, _1 + _2, stream);
+    run(stream);
+  }
+  SECTION("custom stream")
+  {
+    run(custom_stream{stream});


I verified that the stream is extracted in the debugger, but I wonder if I could write the unit test in a way to detect if the default stream was taken anywhere. Does anybody know if I can query the stream whether something was really enqueued there?

One solution would be to start a graph capture on a stream and see if anything was captured, but that might have some limitations, not sure if its applicable here

miscco

I would love to better understand where the issue with the conversions from cudaStream_t lies so that we can extent our get_stream CPO

What was the exact case failing? from what I can see it the second and third branch in the if constexpr should already work

Maybe we just need another clause for get_stream

miscco · 2026-01-15T07:44:01Z

cub/cub/device/device_transform.cuh

+    }
+    else
+    {
+      return ::cuda::std::execution::__query_or(env, ::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}).get();


This is very close to what we have in get_stream

I believe the issue is that we are checking whether tis convertible to stream_ref and not cudaStream_t

But at the same time cudaStream_t should be convertible to stream_ref

copy-pr-bot · 2026-01-15T08:52:33Z

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.

bernhardmgruber · 2026-01-15T08:53:16Z

Discussed this PR with @miscco and we concluded that the environment query for cuda::get_stream misses a check whether the environment is convertible to a cudaStream_t. He will add this and then we can replace the get_stream in this PR by just the environment query.

miscco · 2026-01-15T09:26:07Z

/ok to test cc8f29e

miscco · 2026-01-15T09:27:04Z

I have updated our get_stream CPO to ensure that it also works with types that are convertible to ::cudaStream_t as apposed to ::cuda::stream_ref

bernhardmgruber · 2026-01-15T23:19:39Z

I started the backport already from the current state: #7263, since I need to rebase and change this PR again, since #6914 was merged and we have conflicts.

NVIDIA#6204 changed cub::DeviceTransform APIs from taking cudaStream_t to environments. Special handling preserved support for cudaStream_t. However, user-provided stream types with conversion operators to cudaStream_t were now queried as environments, failing to return a stream.

We should not have a special overload for the `get_stream` overload but should ensure that we can pass types that are convertible to `::cudaStream_t` and extract that stream I have expanded the get_stream CPO to accept a `::cudaStream_t __stream` so that we can work with those types

github-actions · 2026-01-16T16:06:49Z

🥳 CI Workflow Results

🟩 Finished in 6h 06m: Pass: 100%/126 | Total: 5d 13h | Max: 5h 42m | Hits: 89%/250239

See results here.

github-actions · 2026-01-16T16:07:16Z

Successfully created backport PR for branch/3.2.x:

[Backport branch/3.2.x] Fix extracting CUDA stream in cub::DeviceTransform #7272

* Fix extracting CUDA stream in cub::DeviceTransform #6204 changed cub::DeviceTransform APIs from taking cudaStream_t to environments. Special handling preserved support for cudaStream_t. However, user-provided stream types with conversion operators to cudaStream_t were now queried as environments, failing to return a stream. * Properly use `get_stream` in device transform We should not have a special overload for the `get_stream` overload but should ensure that we can pass types that are convertible to `::cudaStream_t` and extract that stream I have expanded the get_stream CPO to accept a `::cudaStream_t __stream` so that we can work with those types Co-authored-by: Michael Schellenberger Costa <[email protected]> (cherry picked from commit 31f8a13)

bernhardmgruber requested a review from a team as a code owner January 14, 2026 23:13

bernhardmgruber requested a review from NaderAlAwar January 14, 2026 23:13

github-project-automation bot added this to CCCL Jan 14, 2026

github-project-automation bot moved this to Todo in CCCL Jan 14, 2026

cccl-authenticator-app bot moved this from Todo to In Review in CCCL Jan 14, 2026

bernhardmgruber commented Jan 14, 2026

View reviewed changes

cub/cub/device/device_transform.cuh Show resolved Hide resolved

bernhardmgruber commented Jan 14, 2026

View reviewed changes

bernhardmgruber added the backport branch/3.2.x label Jan 14, 2026

bernhardmgruber commented Jan 14, 2026

View reviewed changes

pciolkosz approved these changes Jan 15, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

elstehle approved these changes Jan 15, 2026

View reviewed changes

miscco reviewed Jan 15, 2026

View reviewed changes

bernhardmgruber marked this pull request as draft January 15, 2026 08:52

cccl-authenticator-app bot moved this from In Review to In Progress in CCCL Jan 15, 2026

miscco marked this pull request as ready for review January 15, 2026 09:26

miscco requested a review from a team as a code owner January 15, 2026 09:26

miscco requested a review from griwes January 15, 2026 09:26

cccl-authenticator-app bot moved this from In Progress to In Review in CCCL Jan 15, 2026

bernhardmgruber enabled auto-merge (squash) January 15, 2026 12:36

This comment has been minimized.

Sign in to view

bernhardmgruber mentioned this pull request Jan 15, 2026

[Backport to 3.2] Fix extracting CUDA stream in cub::DeviceTransform #7263

Merged

bernhardmgruber and others added 2 commits January 16, 2026 00:28

bernhardmgruber force-pushed the fix_custom_stream_transform branch from cc8f29e to 8ff0cdd Compare January 15, 2026 23:29

This comment has been minimized.

Sign in to view

miscco added 2 commits January 16, 2026 09:16

Fix issue with __query_or

65d8ee9

Avoid unused variable warning

b87896f

bernhardmgruber merged commit 31f8a13 into NVIDIA:main Jan 16, 2026
139 checks passed

github-project-automation bot moved this from In Review to Done in CCCL Jan 16, 2026

github-actions bot mentioned this pull request Jan 16, 2026

[Backport branch/3.2.x] Fix extracting CUDA stream in cub::DeviceTransform #7272

Closed

bernhardmgruber deleted the fix_custom_stream_transform branch January 16, 2026 17:46

Fix extracting CUDA stream in cub::DeviceTransform #7239

Fix extracting CUDA stream in cub::DeviceTransform #7239

Conversation

bernhardmgruber commented Jan 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

bernhardmgruber Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

bernhardmgruber Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

pciolkosz Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

This comment has been minimized.

miscco left a comment

Choose a reason for hiding this comment

Uh oh!

miscco Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

copy-pr-bot bot commented Jan 15, 2026

Uh oh!

bernhardmgruber commented Jan 15, 2026

Uh oh!

miscco commented Jan 15, 2026

Uh oh!

miscco commented Jan 15, 2026

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

bernhardmgruber commented Jan 15, 2026

Uh oh!

This comment has been minimized.

This comment has been minimized.

github-actions bot commented Jan 16, 2026

🥳 CI Workflow Results

🟩 Finished in 6h 06m: Pass: 100%/126 | Total: 5d 13h | Max: 5h 42m | Hits: 89%/250239

Uh oh!

Uh oh!

github-actions bot commented Jan 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Fix extracting CUDA stream in `cub::DeviceTransform` #7239

Fix extracting CUDA stream in `cub::DeviceTransform` #7239

bernhardmgruber commented Jan 14, 2026 •

edited

Loading