Repository navigation
[CUB] Extract shared numeric and scan reference test helpers #12260
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.
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 1 commit into
NVIDIA:main
from
oleksandr-pavlyk:extract-shared-cub-test-helpers
Oct 27, 2026
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,149 @@ | ||
| // SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
| // SPDX-License-Identifier: BSD-3 | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <cuda/iterator> | ||
| #include <cuda/std/functional> | ||
|
|
||
| #include <cstddef> | ||
| #include <cstdint> | ||
| #include <type_traits> | ||
| #include <utility> | ||
|
|
||
| #include <vector_types.h> | ||
|
|
||
| #include <c2h/custom_type.h> | ||
| #include <c2h/extended_types.h> | ||
|
|
||
| // Comparing results computed on CPU and GPU for extended floating point types is impossible. | ||
| // For instance, when used with a constant iterator of two, the accumulator in sequential reference | ||
| // computation (CPU) bumps into the 4096 limits, which will never change (`4096 + 2 = 4096`). | ||
| // Meanwhile, per-thread aggregates (`2 * 16 = 32`) are accumulated within and among thread blocks, | ||
| // yielding `inf` as a result. No reasonable epsilon can be selected to compare `inf` with `4096`. | ||
| // To make `__half` and `__nv_bfloat16` arithmetic associative, the function object below raises | ||
| // extended floating points to the area of unsigned short integers. This allows us to test large | ||
| // inputs with few code-path differences in device algorithms. | ||
| struct ExtendedFloatSum | ||
| { | ||
| template <class T> | ||
| __host__ __device__ T operator()(T a, T b) const | ||
| { | ||
| T result{}; | ||
| result.__x = a.raw() + b.raw(); | ||
| return result; | ||
| } | ||
|
|
||
| #if TEST_HALF_T() | ||
| __host__ __device__ __half operator()(__half a, __half b) const | ||
| { | ||
| uint16_t result = this->operator()(half_t{a}, half_t(b)).raw(); | ||
| return reinterpret_cast<__half&>(result); | ||
| } | ||
| #endif // TEST_HALF_T() | ||
|
|
||
| #if TEST_BF_T() | ||
| __device__ __nv_bfloat16 operator()(__nv_bfloat16 a, __nv_bfloat16 b) const | ||
| { | ||
| uint16_t result = this->operator()(bfloat16_t{a}, bfloat16_t(b)).raw(); | ||
| return reinterpret_cast<__nv_bfloat16&>(result); | ||
| } | ||
| #endif // TEST_BF_T() | ||
| }; | ||
|
|
||
| template <class It> | ||
| It unwrap_it(It it) | ||
| { | ||
| return it; | ||
| } | ||
|
|
||
| #if TEST_HALF_T() | ||
| inline __half* unwrap_it(half_t* it) | ||
| { | ||
| return reinterpret_cast<__half*>(it); | ||
| } | ||
|
|
||
| template <class OffsetT> | ||
| cuda::constant_iterator<__half, OffsetT> unwrap_it(cuda::constant_iterator<half_t, OffsetT> it) | ||
| { | ||
| half_t wrapped_val = *it; | ||
| __half val = wrapped_val.operator __half(); | ||
| return cuda::constant_iterator<__half, OffsetT>(val); | ||
| } | ||
| #endif // TEST_HALF_T() | ||
|
|
||
| #if TEST_BF_T() | ||
| inline __nv_bfloat16* unwrap_it(bfloat16_t* it) | ||
| { | ||
| return reinterpret_cast<__nv_bfloat16*>(it); | ||
| } | ||
|
|
||
| template <class OffsetT> | ||
| cuda::constant_iterator<__nv_bfloat16, OffsetT> unwrap_it(cuda::constant_iterator<bfloat16_t, OffsetT> it) | ||
| { | ||
| bfloat16_t wrapped_val = *it; | ||
| __nv_bfloat16 val = wrapped_val.operator __nv_bfloat16(); | ||
| return cuda::constant_iterator<__nv_bfloat16, OffsetT>(val); | ||
| } | ||
| #endif // TEST_BF_T() | ||
|
|
||
| template <typename T> | ||
| using unwrap_value_t = std::remove_reference_t<decltype(*unwrap_it(std::declval<T*>()))>; | ||
|
|
||
| template <class WrappedItT, // | ||
| class ItT = decltype(unwrap_it(std::declval<WrappedItT>()))> | ||
| std::integral_constant<bool, !std::is_same_v<WrappedItT, ItT>> // | ||
| inline reference_extended_fp(WrappedItT) | ||
| { | ||
| return {}; | ||
| } | ||
|
|
||
| inline constexpr ExtendedFloatSum unwrap_op(std::true_type /* extended float */, ::cuda::std::plus<>) // | ||
| { | ||
| return {}; | ||
| } | ||
|
|
||
| template <bool V, class OpT> | ||
| inline constexpr OpT unwrap_op(std::integral_constant<bool, V> /* base case */, OpT op) | ||
| { | ||
| return op; | ||
| } | ||
|
|
||
| /** | ||
| * @brief Initializes the given item type with a constant non-zero value. | ||
| */ | ||
| template <typename T> | ||
| inline void init_default_constant(T& val, int element_val = 2) | ||
| { | ||
| val = T{static_cast<T>(element_val)}; | ||
| } | ||
|
|
||
| template <template <typename> class... Policies> | ||
| inline void init_default_constant(c2h::custom_type_t<Policies...>& val, int element_val = 2) | ||
| { | ||
| val.key = static_cast<size_t>(element_val); | ||
| val.val = static_cast<size_t>(element_val); | ||
| } | ||
|
|
||
| inline void init_default_constant(uchar3& val, int element_val = 2) | ||
| { | ||
| const auto element_init = static_cast<unsigned char>(element_val); | ||
| val = uchar3{element_init, element_init, element_init}; | ||
| } | ||
|
|
||
| _CCCL_SUPPRESS_DEPRECATED_PUSH | ||
| _CCCL_SUPPRESS_DEPRECATED_NVRTC_DIAG | ||
| inline void init_default_constant(ulonglong4& val, int element_val = 2) | ||
| { | ||
| const auto element_init = static_cast<unsigned long long>(element_val); | ||
| val = ulonglong4{element_init, element_init, element_init, element_init}; | ||
| } | ||
| _CCCL_SUPPRESS_DEPRECATED_POP | ||
|
|
||
| #if _CCCL_CTK_AT_LEAST(13, 0) | ||
| inline void init_default_constant(ulonglong4_16a& val, int element_val = 2) | ||
| { | ||
| const auto element_init = static_cast<unsigned long long>(element_val); | ||
| val = ulonglong4_16a{element_init, element_init, element_init, element_init}; | ||
| } | ||
| #endif // _CCCL_CTK_AT_LEAST(13, 0) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| // SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. | ||
| // SPDX-License-Identifier: BSD-3 | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <cub/detail/it_traits.cuh> | ||
| #include <cub/detail/type_traits.cuh> | ||
| #include <cub/thread/thread_operators.cuh> | ||
|
|
||
| template <typename InputIt, typename OutputIt, typename InitValueT, typename BinaryOp> | ||
| void compute_exclusive_scan_reference(InputIt first, InputIt last, OutputIt result, InitValueT init, BinaryOp op) | ||
| { | ||
| using value_t = cub::detail::it_value_t<InputIt>; | ||
| using accum_t = ::cuda::std::__accumulator_t<BinaryOp, value_t, InitValueT>; | ||
| using output_t = cub::detail::it_value_t<OutputIt>; | ||
| accum_t acc = static_cast<accum_t>(init); | ||
| for (; first != last; ++first) | ||
| { | ||
| auto v = *first; | ||
| *result++ = static_cast<output_t>(acc); | ||
| acc = op(acc, v); | ||
| } | ||
| } | ||
|
|
||
| template <typename InputIt, typename OutputIt, typename BinaryOp, typename InitValueT> | ||
| void compute_inclusive_scan_reference(InputIt first, InputIt last, OutputIt result, BinaryOp op, InitValueT init) | ||
| { | ||
| using value_t = cub::detail::it_value_t<InputIt>; | ||
| using accum_t = ::cuda::std::__accumulator_t<BinaryOp, value_t, InitValueT>; | ||
| using output_t = cub::detail::it_value_t<OutputIt>; | ||
| accum_t acc = static_cast<accum_t>(init); | ||
| for (; first != last; ++first) | ||
| { | ||
| acc = op(acc, *first); | ||
| *result++ = static_cast<output_t>(acc); | ||
| } | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Q: Is it correct to keep the old date here? Not sure since the code was only moved and not touched.