Skip to content

Fix tiny problem sizes for warpspeed scan #7921

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 3 commits into from
Mar 9, 2026
Merged

Fix tiny problem sizes for warpspeed scan #7921

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ea03e19
Handle 0 items in warpspeed scan
bernhardmgruber Mar 5, 2026
44de595
Test 0 num_Items
bernhardmgruber Mar 6, 2026
66da4ca
Fix
bernhardmgruber Mar 6, 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
2 changes: 1 addition & 1 deletion cub/cub/detail/warpspeed/squad/load_store.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ squadStoreBulkSync(Squad squad, CpAsyncOobInfo<OutputT> cpAsyncOobInfo, const ::

constexpr ::cuda::std::uint16_t byteMask = 0xFFFF;
const ::cuda::std::uint16_t byteMaskStart = byteMask << cpAsyncOobInfo.smemStartSkipBytes;
const ::cuda::std::uint16_t byteMaskEnd = byteMask >> (16 - cpAsyncOobInfo.smemEndBytesAfter16BBoundary);
const ::cuda::std::uint16_t byteMaskEnd = byteMask >> (16 - cpAsyncOobInfo.smemEndBytesAfter16BBoundary) % 16;
Copy link
Copy Markdown
Contributor Author

@bernhardmgruber bernhardmgruber Mar 6, 2026

Choose a reason for hiding this comment

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

@ahendriksen is there any smarter way to shift byteMask by smemEndBytesAfter16BBoundary, but leave it when it's 16? Would a predicated shift be faster?

Copy link
Copy Markdown
Contributor

@ahendriksen ahendriksen Mar 8, 2026

Choose a reason for hiding this comment

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

Can't think of a quicker way off the top of my head. I always hoped that the compiler would figure out the best way to compute all these values.

A modulo 16 operation is just and AND by ~0xF, which should be quicker than computing the predicate.

Copy link
Copy Markdown
Contributor Author

@bernhardmgruber bernhardmgruber Mar 9, 2026

Choose a reason for hiding this comment

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

Alright, thx!

// byteMaskStart contains zeroes at the left
# if _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2)
const ::cuda::std::uint16_t byteMaskSmall = byteMaskStart & byteMaskEnd;
Expand Down
6 changes: 6 additions & 0 deletions cub/cub/device/dispatch/dispatch_scan.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -410,6 +410,12 @@ struct DispatchScan
template <typename ActivePolicyT>
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t __invoke_lookahead_algorithm(ActivePolicyT)
{
if (num_items == 0)
{
temp_storage_bytes = 1; // just fulfill the contract that CUB always requires some temporary storage
return cudaSuccess;
}

using InputT = ::cuda::std::iter_value_t<InputIteratorT>;
using OutputT = ::cuda::std::iter_value_t<OutputIteratorT>;
using WarpspeedPolicy = typename ActivePolicyT::WarpspeedPolicy;
Expand Down
2 changes: 1 addition & 1 deletion cub/test/catch2_test_device_scan_alignment.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ C2H_TEST("Device scan works with all device interfaces", "[scan][device]", value
constexpr offset_t max_num_items = 8192;

const auto offset = GENERATE_COPY(values({0, 1, 3, 4, 7, 8, 11, 12, 16}), take(3, random(0, max_offset)));
const auto num_items = GENERATE_COPY(values({1, max_num_items}), take(64, random(0, max_num_items)));
const auto num_items = GENERATE_COPY(values({0, 1, max_num_items}), take(64, random(2, max_num_items - 1)));

CAPTURE(num_items, offset);

Expand Down
Loading