[CUB] Refactor cub::DeviceFor to always take an environment - #11013
[CUB] Refactor cub::DeviceFor to always take an environment#11013macdonaldezra wants to merge 7 commits into
cub::DeviceFor to always take an environment#11013Conversation
There was a problem hiding this comment.
Pull request overview
This PR continues the CUB device-scope API migration away from a trailing cudaStream_t parameter by refactoring cub::DeviceFor to consistently accept an execution environment (const Env&) for both single-phase and two-phase overloads. It also tightens environment/stream querying behavior to avoid silent fallback to the default stream for env types that are only non-const convertible to cudaStream_t, and fixes ForEachInExtents behavior for empty extents plus a mislabeled NVTX range.
Changes:
- Replaced
cudaStream_tparameters inDeviceFortwo-phase overloads withconst Env&, and removed the bare-stream single-phase overloads/SFINAE. - Added a compile-time guard to reject env types whose stream conversion cannot be queried through a
const&(preventing silent default-stream launches). - Fixed
ForEachInExtentsempty-extent handling (avoidsfast_div_modassertion) and corrected theForEachInLayoutNVTX range name; expanded/updated Catch2 coverage accordingly.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
cub/cub/device/device_for.cuh |
Refactors DeviceFor APIs to use environments everywhere, adds const-query static_assert guard, fixes empty-extents path, and corrects NVTX labeling/docs. |
cub/test/catch2_test_device_for_env.cu |
Adds broad environment-kind acceptance tests, stream-routing verification, and two-phase tuning propagation coverage for the updated env-based APIs. |
cub/test/catch2_test_device_for_api.cu |
Updates unambiguity guard tests and adds size-query guards for extents/layout plus a runtime test for empty-extents handling. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesDeviceFor operations now use execution environments for stream selection. Legacy ChangesDeviceFor execution environments
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The API refactor and associated fixes have no demonstrated merge-blocking correctness, security, availability, or deployment risk at the current head; the remaining item is limited to a localized const-qualification style cleanup. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cub/cub/device/device_for.cuh (1)
1278-1289: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winsuggestion:
extents,shape,sub_sizes_div_array,extents_div_array, andop_wrapperare never modified after initialization. Declare themconst.- auto extents = layout_mapping.extents(); + const auto extents = layout_mapping.extents(); using ShapeT = implicit_prom_t<extent_index_type>; - auto shape = static_cast<ShapeT>(cub::detail::size(extents)); + const auto shape = static_cast<ShapeT>(cub::detail::size(extents)); // must precede the fast_div_mod arrays below, whose constructor asserts a positive divisor if (shape == 0) { return cudaSuccess; } - fast_mod_array_t sub_sizes_div_array = cub::detail::sub_sizes_fast_div_mod<is_layout_right>(extents, seq); - fast_mod_array_t extents_div_array = cub::detail::extents_fast_div_mod(extents, seq); - for_each::op_wrapper_extents_t<OpType, extents_type, is_layout_right, fast_mod_array_t> op_wrapper{ + const fast_mod_array_t sub_sizes_div_array = cub::detail::sub_sizes_fast_div_mod<is_layout_right>(extents, seq); + const fast_mod_array_t extents_div_array = cub::detail::extents_fast_div_mod(extents, seq); + const for_each::op_wrapper_extents_t<OpType, extents_type, is_layout_right, fast_mod_array_t> op_wrapper{ op, extents, sub_sizes_div_array, extents_div_array};As per coding guidelines: "All variables that are not modified must be declared
const, including cast results, function return values, and loop-invariant computations."Source: Coding guidelines
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 38b95cf6-a477-4750-b9e7-9fe3ee2c9994
📒 Files selected for processing (3)
cub/cub/device/device_for.cuhcub/test/catch2_test_device_for_api.cucub/test/catch2_test_device_for_env.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
| // `env` is queried through a const reference, so a conversion operator or accessor that is not | ||
| // const-qualified is unreachable and would silently fall back to the default stream. | ||
| static_assert(!(::cuda::std::is_convertible_v<EnvT, cudaStream_t> | ||
| && !::cuda::std::__is_callable_v<::cuda::get_stream_t, const EnvT&>), | ||
| "a type convertible to cudaStream_t must have a const-qualified conversion operator to be usable " | ||
| "as a DeviceFor environment"); |
There was a problem hiding this comment.
This is unfortunately not valid. We allow user to provide environments without a stream and are falling back to a default stream
There was a problem hiding this comment.
Ah kk, I removed the assert and added a test pinning the default-stream fallback for these types instead here.
| template <class RandomAccessIteratorT, class NumItemsT, class OpT> | ||
| [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __for_each_n( |
There was a problem hiding this comment.
Critical: This overload cannot be removed. The intention here is to query for the necessary temporary storage size.
The overload needs to be updated to take a const _Env&
There was a problem hiding this comment.
Oof, yup, I mistakenly thought this was dead code. I updated the function and added a test case for the overload update in this commit.
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: cbd4e601-3346-4d85-9f7c-1677215a9eb4
📒 Files selected for processing (2)
cub/cub/device/device_for.cuhcub/test/catch2_test_device_for_env.cu
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Description
Part of #9875 (the DeviceFor row).
Replaces the trailing
cudaStream_tparameter of the two-phaseBulk,ForEachN,ForEach,ForEachCopyN,ForEachCopy,ForEachInExtents, andForEachInLayoutoverloads with an execution environment, and removes the bare-stream single-phase overloads and their SFINAE, which are no longer needed now that environments are taken byconst&(#10105).Notes:
cudaStream_tis not const-qualified — previously served by the removed overloads, these would otherwise silently launch on the default stream. Flagged for discussion: happy to generalize it intocuda::get_stream/env_dispatch.cuhas a follow-up, or drop it for exact parity with the sibling migrations.fast_div_modassertion failure for empty extents inForEachInExtents, the mislabeledForEachInLayoutNVTX range, and its param docs.Checklist