Skip to content

Productize tuning API for DeviceSegmentedScan #9430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bernhardmgruber merged 2 commits into from
Jun 15, 2026
Merged

Productize tuning API for DeviceSegmentedScan #9430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
516f1c7
productize tuning API for segscan
gonidelis Jun 12, 2026
c65aecc
reviews
gonidelis Jun 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions cub/benchmarks/bench/segmented_scan/base.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,9 @@
template <int ThreadsPerBlock, int ItemsPerThread, int MaxSegmentsPerBlock>
struct policy_selector_t
{
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
-> cub::detail::segmented_scan::segmented_scan_policy
[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const -> cub::SegmentedScanPolicy

@coderabbitai coderabbitai Bot Jun 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

important: Annotation and type-qualification mismatch with upstream contract.

The function uses _CCCL_HOST_DEVICE but the upstream policy selector (tuning_segmented_scan.cuh:106) uses _CCCL_HOST_DEVICE_API. Additionally, the parameter type cuda::compute_capability and return types cub::SegmentedScanPolicy / cub::SegmentedScanBlockPolicy lack the :: prefix required for full qualification starting from the global namespace. As per coding guidelines, type names must be fully qualified except when already declared in the current or enclosing namespace.

Change _CCCL_HOST_DEVICE to _CCCL_HOST_DEVICE_API, use ::cuda::compute_capability for the parameter, and ::cub::SegmentedScanPolicy / ::cub::SegmentedScanBlockPolicy for the types.

Source: Coding guidelines

@bernhardmgruber bernhardmgruber Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could change to _CCCL_HOST_DEVICE_API here, but not required as part of this PR.

@coderabbitai coderabbitai Bot Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: Understood. Would you like me to open a follow-up issue to track switching _CCCL_HOST_DEVICE → _CCCL_HOST_DEVICE_API and adding fully-qualified type names (::cuda::compute_capability, ::cub::SegmentedScanPolicy, ::cub::SegmentedScanBlockPolicy) in this benchmark file?

🧠 Learnings used 
Learnt from: sleeepyjack
Repo: NVIDIA/cccl PR: 8857
File: cudax/include/cuda/experimental/__cuco/__hyperloglog/default_policy.cuh:92-95
Timestamp: 2026-05-06T23:47:10.296Z
Learning: In this CCCL (CUDA C++ Core Libraries) codebase, do NOT treat `_CCCL_API` as host-only. `_CCCL_API` is defined in `libcudacxx/include/cuda/std/__cccl/visibility.h` and always expands to include `_CCCL_HOST_DEVICE` (possibly alongside `_CCCL_TILE`, `_CCCL_VISIBILITY_HIDDEN`, and `_CCCL_EXCLUDE_FROM_EXPLICIT_INSTANTIATION`). So any function annotated with `_CCCL_API` is already valid for both host and device code; reviewers should not flag it as host-only.

Learnt from: sleeepyjack
Repo: NVIDIA/cccl PR: 8857
File: cudax/include/cuda/experimental/__cuco/__hyperloglog/hyperloglog_impl.cuh:426-428
Timestamp: 2026-05-06T23:47:17.597Z
Learning: In the CCCL/CUDA C++ Core Libraries codebase (NVIDIA/cccl), treat `_CCCL_API` as host-and-device callable. `_CCCL_API` expands to `_CCCL_HOST_DEVICE` via `libcudacxx/include/cuda/std/__cccl/visibility.h`, so functions marked with `_CCCL_API` may be invoked from both host and device code. When reviewing device-code contexts, do not flag `_CCCL_API`-annotated functions as if they were host-only.

Learnt from: NaderAlAwar
Repo: NVIDIA/cccl PR: 9171
File: libcudacxx/include/cuda/std/__cccl/diagnostic.h:126-126
Timestamp: 2026-06-03T13:24:36.356Z
Learning: When reviewing code that uses NVIDIA NVCC/NVRTC diagnostic suppression pragmas, treat `nv_diag_suppress` as supporting comma-separated warning numbers in a single directive. For example, `_Pragma("nv_diag_suppress 1444,20199")` (or an equivalent `#pragma nv_diag_suppress 1444,20199`) is valid; do not flag comma-separated warning IDs inside `nv_diag_suppress` pragmas as unsupported.

Learnt from: elstehle
Repo: NVIDIA/cccl PR: 9288
File: cub/benchmarks/bench/segmented_topk/variable/common.cuh:4-4
Timestamp: 2026-06-07T16:13:21.020Z
Learning: In NVIDIA/cccl, headers under cub/benchmarks/ that are .cuh benchmark files use `#pragma once` consistently. During code reviews, do not flag `#pragma once` in these benchmark `.cuh` files as a style violation; the preferred include guard style for this area is `#pragma once` rather than `#ifndef`/`#define`/`#endif`.

Learnt from: elstehle
Repo: NVIDIA/cccl PR: 9288
File: cub/benchmarks/bench/segmented_topk/variable/common.cuh:4-5
Timestamp: 2026-06-07T16:18:35.149Z
Learning: In NVIDIA/cccl, benchmark-internal `.cuh` headers under `cub/benchmarks/` should not use CCCL library-header infrastructure. Do not include `#include <cuda/std/detail/__config>`, do not add system-header pragmas (`#pragma GCC system_header`, `#pragma clang system_header`, `#pragma MSVC system_header`), and do not include the CCCL prologue/epilogue headers (`#include <cuda/std/__cccl/prologue.h>` / `#include <cuda/std/__cccl/epilogue.h>`). Instead, these benchmark headers should start with `#pragma once` and then their includes directly.

{
return cub::detail::segmented_scan::segmented_scan_policy{cub::detail::segmented_scan::block_segmented_scan_policy{
return cub::SegmentedScanPolicy{cub::SegmentedScanBlockPolicy{
ThreadsPerBlock,
ItemsPerThread,
TUNE_BLOCK_LOAD_ALGORITHM,
Expand Down
17 changes: 17 additions & 0 deletions cub/cub/device/device_segmented_scan.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,23 @@ CUB_NAMESPACE_BEGIN
//!
//! @cdp_class{DeviceSegmentedScan}
//!
//! @par Tuning
//! All algorithms in DeviceSegmentedScan that accept an environment can be tuned by passing a custom
//! :ref:`policy selector <cub-policy-selectors>` that returns a @ref SegmentedScanPolicy, as shown in
//! the example below:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_scan_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin segmented-scan-policy-selector
//! :end-before: example-end segmented-scan-policy-selector
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_scan_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin segmented-scan-tuning
//! :end-before: example-end segmented-scan-tuning
//!
//! @endrst
struct DeviceSegmentedScan
{
Expand Down
4 changes: 3 additions & 1 deletion cub/cub/device/dispatch/dispatch_segmented_scan.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,7 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
return error;
}

const segmented_scan_policy active_policy = policy_selector(cc);
const SegmentedScanPolicy active_policy = policy_selector(cc);

#if !_CCCL_COMPILER(NVRTC) && defined(CUB_DEBUG_LOG)
NV_IF_TARGET(
Expand Down Expand Up @@ -174,6 +174,8 @@ CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
constexpr int workers_per_block = 1;
const auto max_segments = active_policy.block.max_segments;
const auto threads_per_block = active_policy.block.threads_per_block;
_CCCL_ASSERT(active_policy.block.threads_per_block > 0, "Policy value for threads_per_block is not positive");
_CCCL_ASSERT(active_policy.block.items_per_thread > 0, "Policy value for items_per_thread is not positive");
_CCCL_ASSERT(max_segments > 0, "Policy value for max segments is not positive");
_CCCL_ASSERT(num_segments_per_worker <= max_segments, "Number of segments per block exceeds maximum value");
return {workers_per_block, threads_per_block, ::cuda::std::min(num_segments_per_worker, max_segments)};
Expand Down
2 changes: 1 addition & 1 deletion cub/cub/device/dispatch/kernels/kernel_segmented_scan.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ namespace detail::segmented_scan
//! of a device-wide segmented prefix scan.
//!
//! @tparam SegmentedScanPolicyGetterT
//! Nullary callable type for getting segmented_scan_policy
//! Nullary callable type for getting SegmentedScanPolicy
//!
//! @tparam InputIteratorT
//! Random-access input iterator type
Expand Down
61 changes: 31 additions & 30 deletions cub/cub/device/dispatch/tuning/tuning_segmented_scan.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,19 @@

CUB_NAMESPACE_BEGIN

namespace detail::segmented_scan
{
struct block_segmented_scan_policy
//! The policy for block-level kernels in @ref DeviceSegmentedScan.
struct SegmentedScanBlockPolicy
{
int threads_per_block;
int items_per_thread;
BlockLoadAlgorithm load_algorithm;
CacheLoadModifier load_modifier;
BlockStoreAlgorithm store_algorithm;
BlockScanAlgorithm scan_algorithm;
int max_segments;
int threads_per_block; //!< Number of threads in a CUDA block
int items_per_thread; //!< Number of items processed per thread
BlockLoadAlgorithm load_algorithm; //!< The @ref BlockLoadAlgorithm used for loading items from global memory
CacheLoadModifier load_modifier; //!< The @ref CacheLoadModifier used for loading items from global memory
BlockStoreAlgorithm store_algorithm; //!< The @ref BlockStoreAlgorithm used for storing items to global memory
BlockScanAlgorithm scan_algorithm; //!< The @ref BlockScanAlgorithm used for block scanning
int max_segments; //!< Maximum number of segments processed per block

[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
operator==(const block_segmented_scan_policy& lhs, const block_segmented_scan_policy& rhs)
operator==(const SegmentedScanBlockPolicy& lhs, const SegmentedScanBlockPolicy& rhs) noexcept
{
return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread
&& lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier
Expand All @@ -47,51 +46,54 @@ struct block_segmented_scan_policy
}

[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
operator!=(const block_segmented_scan_policy& lhs, const block_segmented_scan_policy& rhs)
operator!=(const SegmentedScanBlockPolicy& lhs, const SegmentedScanBlockPolicy& rhs) noexcept
{
return !(lhs == rhs);
}

#if !_CCCL_COMPILER(NVRTC)
friend ::std::ostream& operator<<(::std::ostream& os, const block_segmented_scan_policy& policy)
#if _CCCL_HOSTED()
friend ::std::ostream& operator<<(::std::ostream& os, const SegmentedScanBlockPolicy& policy)
{
return os
<< "block_segmented_scan_policy { .threads_per_block = " << policy.threads_per_block
<< "SegmentedScanBlockPolicy { .threads_per_block = " << policy.threads_per_block
<< ", .items_per_thread = " << policy.items_per_thread << ", .load_algorithm = " << policy.load_algorithm
<< ", .load_modifier = " << policy.load_modifier << ", .store_algorithm = " << policy.store_algorithm
<< ", .scan_algorithm = " << policy.scan_algorithm << ", .max_segments_per_block = " << policy.max_segments
<< " }";
}
#endif // !_CCCL_COMPILER(NVRTC)
#endif // _CCCL_HOSTED()
};

struct segmented_scan_policy
//! The tuning policy for all algorithms in @ref DeviceSegmentedScan.
struct SegmentedScanPolicy
{
block_segmented_scan_policy block;
SegmentedScanBlockPolicy block; //!< Policy for the block-level segmented scan kernel

[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
operator==(const segmented_scan_policy& lhs, const segmented_scan_policy& rhs)
operator==(const SegmentedScanPolicy& lhs, const SegmentedScanPolicy& rhs) noexcept
{
return lhs.block == rhs.block;
}

[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
operator!=(const segmented_scan_policy& lhs, const segmented_scan_policy& rhs)
operator!=(const SegmentedScanPolicy& lhs, const SegmentedScanPolicy& rhs) noexcept
{
return !(lhs == rhs);
}

#if !_CCCL_COMPILER(NVRTC)
friend ::std::ostream& operator<<(::std::ostream& os, const segmented_scan_policy& policy)
#if _CCCL_HOSTED()
friend ::std::ostream& operator<<(::std::ostream& os, const SegmentedScanPolicy& policy)
{
return os << "segmented_scan_policy { .block = " << policy.block << " }";
return os << "SegmentedScanPolicy { .block = " << policy.block << " }";
}
#endif // !_CCCL_COMPILER(NVRTC)
#endif // _CCCL_HOSTED()
};

namespace detail::segmented_scan
{
#if _CCCL_HAS_CONCEPTS()
template <typename T>
concept segmented_scan_policy_selector = policy_selector<T, segmented_scan_policy>;
concept segmented_scan_policy_selector = policy_selector<T, SegmentedScanPolicy>;
#endif // _CCCL_HAS_CONCEPTS()

struct policy_selector
Expand All @@ -100,8 +102,7 @@ struct policy_selector
int accum_size;
int accum_align;

[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const
-> segmented_scan_policy
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> SegmentedScanPolicy
{
constexpr int nominal_threads_per_block = 128;
constexpr int nominal_items_per_thread = 9;
Expand All @@ -125,7 +126,7 @@ struct policy_selector
const auto scan_transposed_blockstore =
large_values ? BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED : BLOCK_STORE_WARP_TRANSPOSE;

return segmented_scan_policy{block_segmented_scan_policy{
return SegmentedScanPolicy{SegmentedScanBlockPolicy{
block_scaled.threads_per_block,
block_scaled.items_per_thread,
scan_transposed_blockload,
Expand All @@ -145,12 +146,12 @@ template <typename AccumT>
struct policy_selector_from_types
{
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const
-> segmented_scan_policy
-> SegmentedScanPolicy
{
constexpr auto accum_size = static_cast<int>(sizeof(AccumT));
constexpr auto accum_align = static_cast<int>(alignof(AccumT));
return policy_selector{accum_size, accum_align}(cc);
};
}
};
} // namespace detail::segmented_scan

Expand Down
45 changes: 42 additions & 3 deletions cub/test/catch2_test_device_segmented_scan_env.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -550,10 +550,9 @@ C2H_TEST("Device segmented inclusive scan init with separate offsets uses enviro
template <unsigned int BlockThreads>
struct segmented_scan_tuning
{
_CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const
-> cub::detail::segmented_scan::segmented_scan_policy
_CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> cub::SegmentedScanPolicy
{
return cub::detail::segmented_scan::segmented_scan_policy{cub::detail::segmented_scan::block_segmented_scan_policy{
return cub::SegmentedScanPolicy{cub::SegmentedScanBlockPolicy{
static_cast<int>(BlockThreads),
1,
cub::BLOCK_LOAD_DIRECT,
Expand Down Expand Up @@ -662,3 +661,43 @@ C2H_TEST("Device segmented inclusive scan init can be tuned", "[segmented_scan][
}

#endif // TEST_LAUNCH != 1

#if _CCCL_COMPILER(GCC, >=, 8) // gcc 7 cannot preserve constexpr-ness from p1 to p2
C2H_TEST("SegmentedScanPolicy", "[segmented_scan][device]")
{
STATIC_REQUIRE(::cuda::std::semiregular<cub::SegmentedScanPolicy>);
STATIC_REQUIRE(::cuda::std::is_aggregate_v<cub::SegmentedScanPolicy>);
STATIC_REQUIRE(::cuda::std::semiregular<cub::SegmentedScanBlockPolicy>);
STATIC_REQUIRE(::cuda::std::is_aggregate_v<cub::SegmentedScanBlockPolicy>);

// aggregate init
constexpr auto block1 = cub::SegmentedScanBlockPolicy{
128,
9,
cub::BLOCK_LOAD_WARP_TRANSPOSE,
cub::LOAD_DEFAULT,
cub::BLOCK_STORE_WARP_TRANSPOSE,
cub::BLOCK_SCAN_WARP_SCANS,
512};
constexpr auto p1 = cub::SegmentedScanPolicy{block1};

# if _CCCL_STD_VER >= 2020
// designated init
constexpr auto block2 = cub::SegmentedScanBlockPolicy{
.threads_per_block = 128,
.items_per_thread = 9,
.load_algorithm = cub::BLOCK_LOAD_WARP_TRANSPOSE,
.load_modifier = cub::LOAD_DEFAULT,
.store_algorithm = cub::BLOCK_STORE_WARP_TRANSPOSE,
.scan_algorithm = cub::BLOCK_SCAN_WARP_SCANS,
.max_segments = 512};
constexpr auto p2 = cub::SegmentedScanPolicy{.block = block2};
# else // _CCCL_STD_VER >= 2020
constexpr auto p2 = p1;
# endif // _CCCL_STD_VER >= 2020

// comparison
STATIC_REQUIRE(p1 == p2);
STATIC_REQUIRE_FALSE(p1 != p2);
}
#endif // _CCCL_COMPILER(GCC, >=, 8)
51 changes: 51 additions & 0 deletions cub/test/catch2_test_device_segmented_scan_env_api.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

#include <thrust/device_vector.h>

#include <cuda/__execution/tune.h>
#include <cuda/devices>
#include <cuda/functional>
#include <cuda/stream>
Expand Down Expand Up @@ -316,3 +317,53 @@ C2H_TEST("cub::DeviceSegmentedScan::InclusiveSegmentedScanInit (separate offsets
REQUIRE(error == cudaSuccess);
REQUIRE(d_out == expected);
}

#if _CCCL_STD_VER >= 2020

// example-begin segmented-scan-policy-selector
struct SegmentedScanPolicySelector
{
__host__ __device__ constexpr auto operator()(cuda::compute_capability cc) const -> cub::SegmentedScanPolicy
{
return {.block = {.threads_per_block = 128,
.items_per_thread = cc > cuda::compute_capability{9, 0} ? 11 : 9,
.load_algorithm = cub::BLOCK_LOAD_WARP_TRANSPOSE,
.load_modifier = cub::LOAD_DEFAULT,
.store_algorithm = cub::BLOCK_STORE_WARP_TRANSPOSE,
.scan_algorithm = cub::BLOCK_SCAN_WARP_SCANS,
.max_segments = 512}};
}
};
// example-end segmented-scan-policy-selector

C2H_TEST("cub::DeviceSegmentedScan::ExclusiveSegmentedScan env-based API with tuning", "[segmented_scan][env]")
{
// example-begin segmented-scan-tuning
::cuda::std::int64_t num_segments = 3;
thrust::device_vector<int> d_offsets = {0, 4, 7, 9};
auto d_offsets_it = thrust::raw_pointer_cast(d_offsets.data());
thrust::device_vector<int> d_in{8, 6, 7, 5, 3, 0, 9, 1, 2};
thrust::device_vector<int> d_out(d_in.size());

const auto error = cub::DeviceSegmentedScan::ExclusiveSegmentedScan(
d_in.begin(),
d_out.begin(),
d_offsets_it,
d_offsets_it + 1,
num_segments,
::cuda::std::plus<>{},
0,
cuda::execution::tune(SegmentedScanPolicySelector{}));
if (error != cudaSuccess)
{
std::cerr << "cub::DeviceSegmentedScan::ExclusiveSegmentedScan failed with status: " << error << '\n';
}

thrust::device_vector<int> expected{0, 8, 14, 21, 0, 3, 3, 0, 1};
// example-end segmented-scan-tuning

REQUIRE(error == cudaSuccess);
REQUIRE(d_out == expected);
}

#endif // _CCCL_STD_VER >= 2020
4 changes: 2 additions & 2 deletions cub/test/catch2_test_device_segmented_scan_multi_segment.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,9 +89,9 @@ template <int ThreadsPerBlock, int ItemsPerThread, int MaxSegmentsPerBlock>
struct policy_selector_t
{
[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const
-> cub::detail::segmented_scan::segmented_scan_policy
-> cub::SegmentedScanPolicy
{
return cub::detail::segmented_scan::segmented_scan_policy{cub::detail::segmented_scan::block_segmented_scan_policy{
return cub::SegmentedScanPolicy{cub::SegmentedScanBlockPolicy{
ThreadsPerBlock,
ItemsPerThread,
cub::BLOCK_LOAD_WARP_TRANSPOSE,
Expand Down
Loading