From 98092ec2d37a742d2769dd84c2d10154091df274 Mon Sep 17 00:00:00 2001 From: Bernhard Manfred Gruber Date: Wed, 13 May 2026 20:33:10 +0200 Subject: [PATCH] Use signed offset type for DevicePartition The env-overloads in cub::DevicePartition used unsigned offset types, while the non-env-overloads used signed ones. This was presumably an oversight in #7459. Fixes: #8968 --- cub/cub/device/device_partition.cuh | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/cub/cub/device/device_partition.cuh b/cub/cub/device/device_partition.cuh index 4ce2184d87c..cfc5073383b 100644 --- a/cub/cub/device/device_partition.cuh +++ b/cub/cub/device/device_partition.cuh @@ -321,7 +321,14 @@ public: { _CCCL_NVTX_RANGE_SCOPE("cub::DevicePartition::Flagged"); - using offset_t = detail::choose_offset_t; + using ChooseOffsetT = detail::choose_signed_offset; + using offset_t = typename ChooseOffsetT::type; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const cudaError_t error = ChooseOffsetT::is_exceeding_offset_type(num_items)) + { + return error; + } // Dispatch with environment - handles all boilerplate return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { @@ -565,7 +572,14 @@ public: { _CCCL_NVTX_RANGE_SCOPE("cub::DevicePartition::If"); - using offset_t = detail::choose_offset_t; + using ChooseOffsetT = detail::choose_signed_offset; + using offset_t = typename ChooseOffsetT::type; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const cudaError_t error = ChooseOffsetT::is_exceeding_offset_type(num_items)) + { + return error; + } // Dispatch with environment - handles all boilerplate return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) {