From 682bcd22822db7560268997fcfa78ec61919da39 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 14 Feb 2025 16:13:29 +0000 Subject: [PATCH 01/18] Move unique_by_key kernels to NVRTC compilable header --- cub/cub/agent/agent_unique_by_key.cuh | 7 +- .../dispatch/dispatch_unique_by_key.cuh | 130 +-------------- .../device/dispatch/kernels/unique_by_key.cuh | 150 ++++++++++++++++++ cub/test/catch2_test_nvrtc.cu | 1 + 4 files changed, 155 insertions(+), 133 deletions(-) create mode 100644 cub/cub/device/dispatch/kernels/unique_by_key.cuh diff --git a/cub/cub/agent/agent_unique_by_key.cuh b/cub/cub/agent/agent_unique_by_key.cuh index 787f8443c61b..cdf18bdcf038 100644 --- a/cub/cub/agent/agent_unique_by_key.cuh +++ b/cub/cub/agent/agent_unique_by_key.cuh @@ -49,8 +49,7 @@ #include #include -#include -#include +#include CUB_NAMESPACE_BEGIN @@ -136,8 +135,8 @@ struct AgentUniqueByKey //--------------------------------------------------------------------- // The input key and value type - using KeyT = typename std::iterator_traits::value_type; - using ValueT = typename std::iterator_traits::value_type; + using KeyT = typename ::cuda::std::iterator_traits::value_type; + using ValueT = typename ::cuda::std::iterator_traits::value_type; // Tile status descriptor interface type using ScanTileStateT = ScanTileState; diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 2cc5a2c918e6..e30c33a546d9 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -42,143 +42,15 @@ # pragma system_header #endif // no system header -#include #include +#include #include #include #include #include -#include - CUB_NAMESPACE_BEGIN -/****************************************************************************** - * Kernel entry points - *****************************************************************************/ - -namespace detail::unique_by_key -{ -/** - * @brief Unique by key kernel entry point (multi-block) - * - * @tparam KeyInputIteratorT - * Random-access input iterator type for keys - * - * @tparam ValueInputIteratorT - * Random-access input iterator type for values - * - * @tparam KeyOutputIteratorT - * Random-access output iterator type for keys - * - * @tparam ValueOutputIteratorT - * Random-access output iterator type for values - * - * @tparam NumSelectedIteratorT - * Output iterator type for recording the number of items selected - * - * @tparam ScanTileStateT - * Tile status interface type - * - * @tparam EqualityOpT - * Equality operator type - * - * @tparam OffsetT - * Signed integer type for global offsets - * - * @param[in] d_keys_in - * Pointer to the input sequence of keys - * - * @param[in] d_values_in - * Pointer to the input sequence of values - * - * @param[out] d_keys_out - * Pointer to the output sequence of selected data items - * - * @param[out] d_values_out - * Pointer to the output sequence of selected data items - * - * @param[out] d_num_selected_out - * Pointer to the total number of items selected - * (i.e., length of @p d_keys_out or @p d_values_out) - * - * @param[in] tile_state - * Tile status interface - * - * @param[in] equality_op - * Equality operator - * - * @param[in] num_items - * Total number of input items - * (i.e., length of @p d_keys_in or @p d_values_in) - * - * @param[in] num_tiles - * Total number of tiles for the entire problem - * - * @param[in] vsmem - * Memory to support virtual shared memory - */ -template -__launch_bounds__(int( - vsmem_helper_default_fallback_policy_t< - typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, - AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>::agent_policy_t::BLOCK_THREADS)) - CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceUniqueByKeySweepKernel( - KeyInputIteratorT d_keys_in, - ValueInputIteratorT d_values_in, - KeyOutputIteratorT d_keys_out, - ValueOutputIteratorT d_values_out, - NumSelectedIteratorT d_num_selected_out, - ScanTileStateT tile_state, - EqualityOpT equality_op, - OffsetT num_items, - int num_tiles, - vsmem_t vsmem) -{ - using VsmemHelperT = vsmem_helper_default_fallback_policy_t< - typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, - AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>; - - using AgentUniqueByKeyPolicyT = typename VsmemHelperT::agent_policy_t; - - // Thread block type for selecting data from input tiles - using AgentUniqueByKeyT = typename VsmemHelperT::agent_t; - - // Static shared memory allocation - __shared__ typename VsmemHelperT::static_temp_storage_t static_temp_storage; - - // Get temporary storage - typename AgentUniqueByKeyT::TempStorage& temp_storage = - VsmemHelperT::get_temp_storage(static_temp_storage, vsmem, (blockIdx.x * gridDim.y) + blockIdx.y); - - // Process tiles - AgentUniqueByKeyT(temp_storage, d_keys_in, d_values_in, d_keys_out, d_values_out, equality_op, num_items) - .ConsumeRange(num_tiles, tile_state, d_num_selected_out); - - // If applicable, hints to discard modified cache lines for vsmem - VsmemHelperT::discard_temp_storage(temp_storage); -} -} // namespace detail::unique_by_key /****************************************************************************** * Dispatch ******************************************************************************/ diff --git a/cub/cub/device/dispatch/kernels/unique_by_key.cuh b/cub/cub/device/dispatch/kernels/unique_by_key.cuh new file mode 100644 index 000000000000..2a8cd5858a4c --- /dev/null +++ b/cub/cub/device/dispatch/kernels/unique_by_key.cuh @@ -0,0 +1,150 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#pragma once + +#include + +#include "cub/util_namespace.cuh" + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Kernel entry points + *****************************************************************************/ + +namespace detail::unique_by_key +{ +/** + * @brief Unique by key kernel entry point (multi-block) + * + * @tparam KeyInputIteratorT + * Random-access input iterator type for keys + * + * @tparam ValueInputIteratorT + * Random-access input iterator type for values + * + * @tparam KeyOutputIteratorT + * Random-access output iterator type for keys + * + * @tparam ValueOutputIteratorT + * Random-access output iterator type for values + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording the number of items selected + * + * @tparam ScanTileStateT + * Tile status interface type + * + * @tparam EqualityOpT + * Equality operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @param[in] d_keys_in + * Pointer to the input sequence of keys + * + * @param[in] d_values_in + * Pointer to the input sequence of values + * + * @param[out] d_keys_out + * Pointer to the output sequence of selected data items + * + * @param[out] d_values_out + * Pointer to the output sequence of selected data items + * + * @param[out] d_num_selected_out + * Pointer to the total number of items selected + * (i.e., length of @p d_keys_out or @p d_values_out) + * + * @param[in] tile_state + * Tile status interface + * + * @param[in] equality_op + * Equality operator + * + * @param[in] num_items + * Total number of input items + * (i.e., length of @p d_keys_in or @p d_values_in) + * + * @param[in] num_tiles + * Total number of tiles for the entire problem + * + * @param[in] vsmem + * Memory to support virtual shared memory + */ +template +__launch_bounds__(int( + vsmem_helper_default_fallback_policy_t< + typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, + AgentUniqueByKey, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>::agent_policy_t::BLOCK_THREADS)) + CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceUniqueByKeySweepKernel( + KeyInputIteratorT d_keys_in, + ValueInputIteratorT d_values_in, + KeyOutputIteratorT d_keys_out, + ValueOutputIteratorT d_values_out, + NumSelectedIteratorT d_num_selected_out, + ScanTileStateT tile_state, + EqualityOpT equality_op, + OffsetT num_items, + int num_tiles, + vsmem_t vsmem) +{ + using VsmemHelperT = vsmem_helper_default_fallback_policy_t< + typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, + AgentUniqueByKey, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>; + + using AgentUniqueByKeyPolicyT = typename VsmemHelperT::agent_policy_t; + + // Thread block type for selecting data from input tiles + using AgentUniqueByKeyT = typename VsmemHelperT::agent_t; + + // Static shared memory allocation + __shared__ typename VsmemHelperT::static_temp_storage_t static_temp_storage; + + // Get temporary storage + typename AgentUniqueByKeyT::TempStorage& temp_storage = + VsmemHelperT::get_temp_storage(static_temp_storage, vsmem, (blockIdx.x * gridDim.y) + blockIdx.y); + + // Process tiles + AgentUniqueByKeyT(temp_storage, d_keys_in, d_values_in, d_keys_out, d_values_out, equality_op, num_items) + .ConsumeRange(num_tiles, tile_state, d_num_selected_out); + + // If applicable, hints to discard modified cache lines for vsmem + VsmemHelperT::discard_temp_storage(temp_storage); +} +} // namespace detail::unique_by_key + +CUB_NAMESPACE_END diff --git a/cub/test/catch2_test_nvrtc.cu b/cub/test/catch2_test_nvrtc.cu index 9ebfade56141..0dffb0bf944b 100644 --- a/cub/test/catch2_test_nvrtc.cu +++ b/cub/test/catch2_test_nvrtc.cu @@ -60,6 +60,7 @@ TEST_CASE("Test nvrtc", "[test][nvrtc]") #include #include #include + #include #include #include From 1d4e74f0575aed60343c47e7985d5bb1f05e2693 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 14 Feb 2025 18:39:34 +0000 Subject: [PATCH 02/18] Add dynamic dispatch for unique_by_key --- .../dispatch/dispatch_unique_by_key.cuh | 168 ++++++++++++------ .../dispatch/tuning/tuning_unique_by_key.cuh | 36 +++- cub/cub/util_vsmem.cuh | 17 ++ 3 files changed, 159 insertions(+), 62 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index e30c33a546d9..c3d754a5ddd2 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -43,6 +43,7 @@ #endif // no system header #include +#include #include #include #include @@ -51,6 +52,43 @@ CUB_NAMESPACE_BEGIN +namespace detail::unique_by_key +{ +template + +struct DeviceUniqueByKeyKernelSource +{ + CUB_DEFINE_KERNEL_GETTER(CompactInitKernel, + detail::scan::DeviceCompactInitKernel); + + CUB_DEFINE_KERNEL_GETTER( + UniqueByKeySweepKernel, + DeviceUniqueByKeySweepKernel< + MaxPolicyT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + EqualityOpT, + OffsetT>); + + CUB_RUNTIME_FUNCTION ScanTileStateT TileState() + { + return ScanTileStateT(); + } +}; +} // namespace detail::unique_by_key + /****************************************************************************** * Dispatch ******************************************************************************/ @@ -79,15 +117,36 @@ CUB_NAMESPACE_BEGIN * @tparam OffsetT * Signed integer type for global offsets */ -template , detail::value_t>> +template < + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename EqualityOpT, + typename OffsetT, + typename PolicyHub = + detail::unique_by_key::policy_hub, detail::value_t>, + typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< + typename PolicyHub::MaxPolicy, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + NumSelectedIteratorT, + ScanTileState, + EqualityOpT, + OffsetT>, + typename KernelLauncherFactory = detail::TripleChevronFactory, + typename VSMemHelperPolicyT = detail::vsmem_helper_default_fallback_policy_t< + typename PolicyHub::MaxPolicy::UniqueByKeyPolicyT, + detail::unique_by_key::AgentUniqueByKey, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>> struct DispatchUniqueByKey { /****************************************************************************** @@ -103,9 +162,6 @@ struct DispatchUniqueByKey using KeyT = typename std::iterator_traits::value_type; using ValueT = typename std::iterator_traits::value_type; - // Tile status descriptor interface type - using ScanTileStateT = ScanTileState; - /// Device-accessible allocation of temporary storage. When nullptr, the required allocation size /// is written to `temp_storage_bytes` and no work is done. void* d_temp_storage; @@ -138,6 +194,12 @@ struct DispatchUniqueByKey /// **[optional]** CUDA stream to launch kernels within. Default is stream0. cudaStream_t stream; + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + VSMemHelperPolicyT vsmem_helper; + /** * @param[in] d_temp_storage * Device-accessible allocation of temporary storage. @@ -183,7 +245,10 @@ struct DispatchUniqueByKey NumSelectedIteratorT d_num_selected_out, EqualityOpT equality_op, OffsetT num_items, - cudaStream_t stream) + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + VSMemHelperPolicyT vsmem_helper = {}) : d_temp_storage(d_temp_storage) , temp_storage_bytes(temp_storage_bytes) , d_keys_in(d_keys_in) @@ -194,27 +259,19 @@ struct DispatchUniqueByKey , equality_op(equality_op) , num_items(num_items) , stream(stream) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + , vsmem_helper(vsmem_helper) {} /****************************************************************************** * Dispatch entrypoints ******************************************************************************/ - template - CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(InitKernel init_kernel, ScanKernel scan_kernel) + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t + Invoke(InitKernelT init_kernel, UniqueByKeySweepKernelT sweep_kernel, ActivePolicyT policy = {}) { - using Policy = typename ActivePolicyT::UniqueByKeyPolicyT; - - using VsmemHelperT = cub::detail::vsmem_helper_default_fallback_policy_t< - Policy, - detail::unique_by_key::AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>; - cudaError error = cudaSuccess; do { @@ -227,17 +284,19 @@ struct DispatchUniqueByKey } // Number of input tiles - constexpr auto block_threads = VsmemHelperT::agent_policy_t::BLOCK_THREADS; - constexpr auto items_per_thread = VsmemHelperT::agent_policy_t::ITEMS_PER_THREAD; - int tile_size = block_threads * items_per_thread; - int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); - const auto vsmem_size = num_tiles * VsmemHelperT::vsmem_per_block; + const auto block_threads = vsmem_helper.BlockThreads(policy.UniqueByKey()); + const auto items_per_thread = vsmem_helper.ItemsPerThread(policy.UniqueByKey()); + int tile_size = block_threads * items_per_thread; + int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + const auto vsmem_size = num_tiles * vsmem_helper.VSMemPerBlock(); // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0, vsmem_size}; + auto tile_state = kernel_source.TileState(); + // Bytes needed for tile status descriptors - error = CubDebug(ScanTileStateT::AllocationSize(num_tiles, allocation_sizes[0])); + error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0])); if (cudaSuccess != error) { break; @@ -258,8 +317,6 @@ struct DispatchUniqueByKey break; } - // Construct the tile status interface - ScanTileStateT tile_state; error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0])); if (cudaSuccess != error) { @@ -275,7 +332,7 @@ struct DispatchUniqueByKey #endif // CUB_DEBUG_LOG // Invoke init_kernel to initialize tile descriptors - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream) .doit(init_kernel, tile_state, num_tiles, d_num_selected_out); // Check for failure to launch @@ -312,7 +369,7 @@ struct DispatchUniqueByKey scan_grid_size.y = ::cuda::ceil_div(num_tiles, max_dim_x); scan_grid_size.x = CUB_MIN(num_tiles, max_dim_x); -// Log select_if_kernel configuration + // Log select_if_kernel configuration #ifdef CUB_DEBUG_LOG { // Get SM occupancy for unique_by_key_kernel @@ -339,8 +396,8 @@ struct DispatchUniqueByKey // Invoke select_if_kernel error = - THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(scan_grid_size, block_threads, 0, stream) - .doit(scan_kernel, + launcher_factory(scan_grid_size, block_threads, 0, stream) + .doit(sweep_kernel, d_keys_in, d_values_in, d_keys_out, @@ -371,21 +428,12 @@ struct DispatchUniqueByKey } template - CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke() + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) { - // Ensure kernels are instantiated. - return Invoke( - detail::scan::DeviceCompactInitKernel, - detail::unique_by_key::DeviceUniqueByKeySweepKernel< - typename PolicyHub::MaxPolicy, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - NumSelectedIteratorT, - ScanTileStateT, - EqualityOpT, - OffsetT>); + // TODO: implement this + auto wrapped_policy = detail::unique_by_key::MakeUniqueByKeyPolicyWrapper(active_policy); + + return Invoke(kernel_source.CompactInitKernel(), kernel_source.UniqueByKeySweepKernel(), wrapped_policy); } /** @@ -425,6 +473,7 @@ struct DispatchUniqueByKey * **[optional]** CUDA stream to launch kernels within. * Default is stream0. */ + template CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( void* d_temp_storage, size_t& temp_storage_bytes, @@ -435,7 +484,11 @@ struct DispatchUniqueByKey NumSelectedIteratorT d_num_selected_out, EqualityOpT equality_op, OffsetT num_items, - cudaStream_t stream) + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}, + VSMemHelperPolicyT vsmem_helper = {}) { cudaError_t error; do @@ -459,10 +512,13 @@ struct DispatchUniqueByKey d_num_selected_out, equality_op, num_items, - stream); + stream, + kernel_source, + launcher_factory, + vsmem_helper); // Dispatch to chained policy - error = CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch)); + error = CubDebug(max_policy.Invoke(ptx_version, dispatch)); if (cudaSuccess != error) { break; diff --git a/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh b/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh index e7c02c91db13..ad502abc2c78 100644 --- a/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh @@ -47,10 +47,7 @@ CUB_NAMESPACE_BEGIN -namespace detail -{ - -namespace unique_by_key +namespace detail::unique_by_key { enum class primitive_key @@ -770,6 +767,34 @@ struct sm100_tuning +struct UniqueByKeyPolicyWrapper : PolicyT +{ + CUB_RUNTIME_FUNCTION UniqueByKeyPolicyWrapper(PolicyT base) + : PolicyT(base) + {} +}; + +template +struct UniqueByKeyPolicyWrapper> + : StaticPolicyT +{ + CUB_RUNTIME_FUNCTION UniqueByKeyPolicyWrapper(StaticPolicyT base) + : StaticPolicyT(base) + {} + + CUB_RUNTIME_FUNCTION static constexpr PolicyWrapper UniqueByKey() + { + return cub::detail::MakePolicyWrapper(typename StaticPolicyT::UniqueByKeyPolicyT()); + } +}; + +template +CUB_RUNTIME_FUNCTION UniqueByKeyPolicyWrapper MakeUniqueByKeyPolicyWrapper(PolicyT policy) +{ + return UniqueByKeyPolicyWrapper{policy}; +} + template struct policy_hub { @@ -843,8 +868,7 @@ struct policy_hub using MaxPolicy = Policy1000; }; -} // namespace unique_by_key -} // namespace detail +} // namespace detail::unique_by_key template using DeviceUniqueByKeyPolicy CCCL_DEPRECATED_BECAUSE("This class is considered an implementation detail and it will " diff --git a/cub/cub/util_vsmem.cuh b/cub/cub/util_vsmem.cuh index c282fb443cae..1a9514f1c588 100644 --- a/cub/cub/util_vsmem.cuh +++ b/cub/cub/util_vsmem.cuh @@ -96,6 +96,23 @@ public: // The amount of global memory-backed virtual shared memory needed, padded to an integer multiple of 128 bytes static constexpr ::cuda::std::size_t vsmem_per_block = needs_vsmem ? (required_smem + padding_bytes) : 0; + template + _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) + { + return AgentPolicyT::BLOCK_THREADS; + } + + template + _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) + { + return AgentPolicyT::ITEMS_PER_THREAD; + } + + static constexpr ::cuda::std::size_t VSMemPerBlock() + { + return vsmem_per_block; + } + /** * @brief Used from within the device algorithm's kernel to get the temporary storage that can be * passed to the agent, specialized for the case when we can use native shared memory as temporary From bc363a8c12a23053a2476591ce0949e68a9634f6 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 14 Feb 2025 18:56:31 +0000 Subject: [PATCH 03/18] Add missing host device qualifier --- cub/cub/util_vsmem.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/util_vsmem.cuh b/cub/cub/util_vsmem.cuh index 1a9514f1c588..43f8601dc6c8 100644 --- a/cub/cub/util_vsmem.cuh +++ b/cub/cub/util_vsmem.cuh @@ -108,7 +108,7 @@ public: return AgentPolicyT::ITEMS_PER_THREAD; } - static constexpr ::cuda::std::size_t VSMemPerBlock() + _CCCL_HOST_DEVICE static constexpr ::cuda::std::size_t VSMemPerBlock() { return vsmem_per_block; } From 9704eb98f6568f91aedfcd1f30077e5ef1c6250d Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 14 Feb 2025 19:59:54 +0000 Subject: [PATCH 04/18] Fix issue where argument name was not changed --- cub/cub/device/dispatch/dispatch_unique_by_key.cuh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index c3d754a5ddd2..229ce277d2fd 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -373,9 +373,9 @@ struct DispatchUniqueByKey #ifdef CUB_DEBUG_LOG { // Get SM occupancy for unique_by_key_kernel - int scan_sm_occupancy; - error = CubDebug(MaxSmOccupancy(scan_sm_occupancy, // out - scan_kernel, + int sweep_sm_occupancy; + error = CubDebug(MaxSmOccupancy(sweep_sm_occupancy, // out + sweep_kernel, block_threads)); if (cudaSuccess != error) { From 00da18f360c9e94b3fec90bcc23effddbb1313c5 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 17 Feb 2025 21:46:11 +0000 Subject: [PATCH 05/18] Forgot to rename local variable --- cub/cub/device/dispatch/dispatch_unique_by_key.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 229ce277d2fd..0af6887323db 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -390,7 +390,7 @@ struct DispatchUniqueByKey block_threads, (long long) stream, items_per_thread, - scan_sm_occupancy); + sweep_sm_occupancy); } #endif // CUB_DEBUG_LOG From 440716eee88a803bd47cc8af29a2cb4f20144bdc Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 18 Feb 2025 14:43:56 +0000 Subject: [PATCH 06/18] Remove old TODO comment --- cub/cub/device/dispatch/dispatch_unique_by_key.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 0af6887323db..7d233839c3bd 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -430,7 +430,6 @@ struct DispatchUniqueByKey template CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) { - // TODO: implement this auto wrapped_policy = detail::unique_by_key::MakeUniqueByKeyPolicyWrapper(active_policy); return Invoke(kernel_source.CompactInitKernel(), kernel_source.UniqueByKeySweepKernel(), wrapped_policy); From 2a171334cd59c5dbbd479117d3bd52530a3d6868 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 19 Feb 2025 18:47:21 +0000 Subject: [PATCH 07/18] Make VSMemHelper be a template parameter of invoke(). This is needed because this allows us to define it such that we can use the ActivePolicyT template parameter of invoke() as its policy, rather than the policy hub's MaxPolicy. ActivePolicyT will be the correct policy for the current ptx version, since calling invoke() from dispatch() causes us to move down the chained policies until we get to the correct one. --- .../dispatch/dispatch_unique_by_key.cuh | 88 +++++++++---------- cub/cub/util_vsmem.cuh | 4 +- 2 files changed, 45 insertions(+), 47 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 7d233839c3bd..7c4b1621dc5e 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -117,36 +117,26 @@ struct DeviceUniqueByKeyKernelSource * @tparam OffsetT * Signed integer type for global offsets */ -template < - typename KeyInputIteratorT, - typename ValueInputIteratorT, - typename KeyOutputIteratorT, - typename ValueOutputIteratorT, - typename NumSelectedIteratorT, - typename EqualityOpT, - typename OffsetT, - typename PolicyHub = - detail::unique_by_key::policy_hub, detail::value_t>, - typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< - typename PolicyHub::MaxPolicy, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - NumSelectedIteratorT, - ScanTileState, - EqualityOpT, - OffsetT>, - typename KernelLauncherFactory = detail::TripleChevronFactory, - typename VSMemHelperPolicyT = detail::vsmem_helper_default_fallback_policy_t< - typename PolicyHub::MaxPolicy::UniqueByKeyPolicyT, - detail::unique_by_key::AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>> +template , detail::value_t>, + typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< + typename PolicyHub::MaxPolicy, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + NumSelectedIteratorT, + ScanTileState, + EqualityOpT, + OffsetT>, + typename KernelLauncherFactory = detail::TripleChevronFactory> struct DispatchUniqueByKey { /****************************************************************************** @@ -198,8 +188,6 @@ struct DispatchUniqueByKey KernelLauncherFactory launcher_factory; - VSMemHelperPolicyT vsmem_helper; - /** * @param[in] d_temp_storage * Device-accessible allocation of temporary storage. @@ -247,8 +235,7 @@ struct DispatchUniqueByKey OffsetT num_items, cudaStream_t stream, KernelSource kernel_source = {}, - KernelLauncherFactory launcher_factory = {}, - VSMemHelperPolicyT vsmem_helper = {}) + KernelLauncherFactory launcher_factory = {}) : d_temp_storage(d_temp_storage) , temp_storage_bytes(temp_storage_bytes) , d_keys_in(d_keys_in) @@ -261,16 +248,18 @@ struct DispatchUniqueByKey , stream(stream) , kernel_source(kernel_source) , launcher_factory(launcher_factory) - , vsmem_helper(vsmem_helper) {} /****************************************************************************** * Dispatch entrypoints ******************************************************************************/ - template - CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t - Invoke(InitKernelT init_kernel, UniqueByKeySweepKernelT sweep_kernel, ActivePolicyT policy = {}) + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke( + InitKernelT init_kernel, + UniqueByKeySweepKernelT sweep_kernel, + ActivePolicyT policy = {}, + VSMemHelperPolicyT vsmem_helper = {}) { cudaError error = cudaSuccess; do @@ -427,12 +416,23 @@ struct DispatchUniqueByKey return error; } - template - CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) + template > + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t + Invoke(ActivePolicyT active_policy = {}, VSMemHelperPolicyT vsmem_helper = {}) { auto wrapped_policy = detail::unique_by_key::MakeUniqueByKeyPolicyWrapper(active_policy); - return Invoke(kernel_source.CompactInitKernel(), kernel_source.UniqueByKeySweepKernel(), wrapped_policy); + return Invoke( + kernel_source.CompactInitKernel(), kernel_source.UniqueByKeySweepKernel(), wrapped_policy, vsmem_helper); } /** @@ -486,8 +486,7 @@ struct DispatchUniqueByKey cudaStream_t stream, KernelSource kernel_source = {}, KernelLauncherFactory launcher_factory = {}, - MaxPolicyT max_policy = {}, - VSMemHelperPolicyT vsmem_helper = {}) + MaxPolicyT max_policy = {}) { cudaError_t error; do @@ -513,8 +512,7 @@ struct DispatchUniqueByKey num_items, stream, kernel_source, - launcher_factory, - vsmem_helper); + launcher_factory); // Dispatch to chained policy error = CubDebug(max_policy.Invoke(ptx_version, dispatch)); diff --git a/cub/cub/util_vsmem.cuh b/cub/cub/util_vsmem.cuh index 43f8601dc6c8..b2204c55f19d 100644 --- a/cub/cub/util_vsmem.cuh +++ b/cub/cub/util_vsmem.cuh @@ -99,13 +99,13 @@ public: template _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) { - return AgentPolicyT::BLOCK_THREADS; + return AgentT::BLOCK_THREADS; } template _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) { - return AgentPolicyT::ITEMS_PER_THREAD; + return AgentT::ITEMS_PER_THREAD; } _CCCL_HOST_DEVICE static constexpr ::cuda::std::size_t VSMemPerBlock() From adc4acc56ab84e3ba44f347dee7dc1886e9ccd75 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 20 Feb 2025 15:04:04 +0000 Subject: [PATCH 08/18] Revert "Make VSMemHelper be a template parameter of invoke(). This is needed because this allows us to define it such that we can use the ActivePolicyT template parameter of invoke() as its policy, rather than the policy hub's MaxPolicy. ActivePolicyT will be the correct policy for the current ptx version, since calling invoke() from dispatch() causes us to move down the chained policies until we get to the correct one." This reverts commit 2a171334cd59c5dbbd479117d3bd52530a3d6868. --- .../dispatch/dispatch_unique_by_key.cuh | 88 ++++++++++--------- cub/cub/util_vsmem.cuh | 4 +- 2 files changed, 47 insertions(+), 45 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 7c4b1621dc5e..7d233839c3bd 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -117,26 +117,36 @@ struct DeviceUniqueByKeyKernelSource * @tparam OffsetT * Signed integer type for global offsets */ -template , detail::value_t>, - typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< - typename PolicyHub::MaxPolicy, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - NumSelectedIteratorT, - ScanTileState, - EqualityOpT, - OffsetT>, - typename KernelLauncherFactory = detail::TripleChevronFactory> +template < + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename EqualityOpT, + typename OffsetT, + typename PolicyHub = + detail::unique_by_key::policy_hub, detail::value_t>, + typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< + typename PolicyHub::MaxPolicy, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + NumSelectedIteratorT, + ScanTileState, + EqualityOpT, + OffsetT>, + typename KernelLauncherFactory = detail::TripleChevronFactory, + typename VSMemHelperPolicyT = detail::vsmem_helper_default_fallback_policy_t< + typename PolicyHub::MaxPolicy::UniqueByKeyPolicyT, + detail::unique_by_key::AgentUniqueByKey, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>> struct DispatchUniqueByKey { /****************************************************************************** @@ -188,6 +198,8 @@ struct DispatchUniqueByKey KernelLauncherFactory launcher_factory; + VSMemHelperPolicyT vsmem_helper; + /** * @param[in] d_temp_storage * Device-accessible allocation of temporary storage. @@ -235,7 +247,8 @@ struct DispatchUniqueByKey OffsetT num_items, cudaStream_t stream, KernelSource kernel_source = {}, - KernelLauncherFactory launcher_factory = {}) + KernelLauncherFactory launcher_factory = {}, + VSMemHelperPolicyT vsmem_helper = {}) : d_temp_storage(d_temp_storage) , temp_storage_bytes(temp_storage_bytes) , d_keys_in(d_keys_in) @@ -248,18 +261,16 @@ struct DispatchUniqueByKey , stream(stream) , kernel_source(kernel_source) , launcher_factory(launcher_factory) + , vsmem_helper(vsmem_helper) {} /****************************************************************************** * Dispatch entrypoints ******************************************************************************/ - template - CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke( - InitKernelT init_kernel, - UniqueByKeySweepKernelT sweep_kernel, - ActivePolicyT policy = {}, - VSMemHelperPolicyT vsmem_helper = {}) + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t + Invoke(InitKernelT init_kernel, UniqueByKeySweepKernelT sweep_kernel, ActivePolicyT policy = {}) { cudaError error = cudaSuccess; do @@ -416,23 +427,12 @@ struct DispatchUniqueByKey return error; } - template > - CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t - Invoke(ActivePolicyT active_policy = {}, VSMemHelperPolicyT vsmem_helper = {}) + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) { auto wrapped_policy = detail::unique_by_key::MakeUniqueByKeyPolicyWrapper(active_policy); - return Invoke( - kernel_source.CompactInitKernel(), kernel_source.UniqueByKeySweepKernel(), wrapped_policy, vsmem_helper); + return Invoke(kernel_source.CompactInitKernel(), kernel_source.UniqueByKeySweepKernel(), wrapped_policy); } /** @@ -486,7 +486,8 @@ struct DispatchUniqueByKey cudaStream_t stream, KernelSource kernel_source = {}, KernelLauncherFactory launcher_factory = {}, - MaxPolicyT max_policy = {}) + MaxPolicyT max_policy = {}, + VSMemHelperPolicyT vsmem_helper = {}) { cudaError_t error; do @@ -512,7 +513,8 @@ struct DispatchUniqueByKey num_items, stream, kernel_source, - launcher_factory); + launcher_factory, + vsmem_helper); // Dispatch to chained policy error = CubDebug(max_policy.Invoke(ptx_version, dispatch)); diff --git a/cub/cub/util_vsmem.cuh b/cub/cub/util_vsmem.cuh index b2204c55f19d..43f8601dc6c8 100644 --- a/cub/cub/util_vsmem.cuh +++ b/cub/cub/util_vsmem.cuh @@ -99,13 +99,13 @@ public: template _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) { - return AgentT::BLOCK_THREADS; + return AgentPolicyT::BLOCK_THREADS; } template _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) { - return AgentT::ITEMS_PER_THREAD; + return AgentPolicyT::ITEMS_PER_THREAD; } _CCCL_HOST_DEVICE static constexpr ::cuda::std::size_t VSMemPerBlock() From 5b210eca8c29bef416f1bd761720422b5a5d141b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 20 Feb 2025 15:10:51 +0000 Subject: [PATCH 09/18] Make vsmem_helper a template on the dispatch class so that we can pass it from the c.parallel layer. Also we don't need it to have any run-time state --- .../dispatch/dispatch_unique_by_key.cuh | 18 ++++------ cub/cub/util_vsmem.cuh | 36 ++++++++++++------- 2 files changed, 30 insertions(+), 24 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 7d233839c3bd..27a3523180c9 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -198,8 +198,6 @@ struct DispatchUniqueByKey KernelLauncherFactory launcher_factory; - VSMemHelperPolicyT vsmem_helper; - /** * @param[in] d_temp_storage * Device-accessible allocation of temporary storage. @@ -247,8 +245,7 @@ struct DispatchUniqueByKey OffsetT num_items, cudaStream_t stream, KernelSource kernel_source = {}, - KernelLauncherFactory launcher_factory = {}, - VSMemHelperPolicyT vsmem_helper = {}) + KernelLauncherFactory launcher_factory = {}) : d_temp_storage(d_temp_storage) , temp_storage_bytes(temp_storage_bytes) , d_keys_in(d_keys_in) @@ -261,7 +258,6 @@ struct DispatchUniqueByKey , stream(stream) , kernel_source(kernel_source) , launcher_factory(launcher_factory) - , vsmem_helper(vsmem_helper) {} /****************************************************************************** @@ -284,11 +280,11 @@ struct DispatchUniqueByKey } // Number of input tiles - const auto block_threads = vsmem_helper.BlockThreads(policy.UniqueByKey()); - const auto items_per_thread = vsmem_helper.ItemsPerThread(policy.UniqueByKey()); + const auto block_threads = VSMemHelperPolicyT::BlockThreads(policy.UniqueByKey()); + const auto items_per_thread = VSMemHelperPolicyT::ItemsPerThread(policy.UniqueByKey()); int tile_size = block_threads * items_per_thread; int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); - const auto vsmem_size = num_tiles * vsmem_helper.VSMemPerBlock(); + const auto vsmem_size = num_tiles * VSMemHelperPolicyT::VSMemPerBlock(); // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0, vsmem_size}; @@ -486,8 +482,7 @@ struct DispatchUniqueByKey cudaStream_t stream, KernelSource kernel_source = {}, KernelLauncherFactory launcher_factory = {}, - MaxPolicyT max_policy = {}, - VSMemHelperPolicyT vsmem_helper = {}) + MaxPolicyT max_policy = {}) { cudaError_t error; do @@ -513,8 +508,7 @@ struct DispatchUniqueByKey num_items, stream, kernel_source, - launcher_factory, - vsmem_helper); + launcher_factory); // Dispatch to chained policy error = CubDebug(max_policy.Invoke(ptx_version, dispatch)); diff --git a/cub/cub/util_vsmem.cuh b/cub/cub/util_vsmem.cuh index 43f8601dc6c8..a23af409c04a 100644 --- a/cub/cub/util_vsmem.cuh +++ b/cub/cub/util_vsmem.cuh @@ -96,18 +96,6 @@ public: // The amount of global memory-backed virtual shared memory needed, padded to an integer multiple of 128 bytes static constexpr ::cuda::std::size_t vsmem_per_block = needs_vsmem ? (required_smem + padding_bytes) : 0; - template - _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) - { - return AgentPolicyT::BLOCK_THREADS; - } - - template - _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) - { - return AgentPolicyT::ITEMS_PER_THREAD; - } - _CCCL_HOST_DEVICE static constexpr ::cuda::std::size_t VSMemPerBlock() { return vsmem_per_block; @@ -232,6 +220,18 @@ struct vsmem_helper_with_fallback_impl : public vsmem_helper_impl { using agent_t = DefaultAgentT; using agent_policy_t = DefaultAgentPolicyT; + + template + _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) + { + return AgentPolicyT::BLOCK_THREADS; + } + + template + _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) + { + return AgentPolicyT::ITEMS_PER_THREAD; + } }; template struct vsmem_helper_with_fallback_impl @@ -239,6 +239,18 @@ struct vsmem_helper_with_fallback_impl + _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) + { + return FallbackAgentPolicyT::BLOCK_THREADS; + } + + template + _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) + { + return FallbackAgentPolicyT::ITEMS_PER_THREAD; + } }; /** From dff4361cf1cb8ab8cac3b425acabff66baa9e207 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 20 Feb 2025 20:33:16 +0000 Subject: [PATCH 10/18] Refactor how we instantiate vsmem_helper --- .../dispatch/dispatch_unique_by_key.cuh | 45 ++++++++++++++----- cub/cub/util_vsmem.cuh | 29 ------------ 2 files changed, 33 insertions(+), 41 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index 27a3523180c9..dce03ddd433c 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -87,6 +87,34 @@ struct DeviceUniqueByKeyKernelSource return ScanTileStateT(); } }; + +template +struct VSMemHelper +{ + template + CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(ActivePolicyT /*policy*/) + { + return cub::detail::vsmem_helper_default_fallback_policy_t::agent_policy_t::BLOCK_THREADS; + } + + template + CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(ActivePolicyT /*policy*/) + { + return cub::detail::vsmem_helper_default_fallback_policy_t::agent_policy_t::ITEMS_PER_THREAD; + } + + template + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t VSMemPerBlock(ActivePolicyT /*policy*/) + { + return cub::detail::vsmem_helper_default_fallback_policy_t::vsmem_per_block; + } +}; } // namespace detail::unique_by_key /****************************************************************************** @@ -138,15 +166,8 @@ template < EqualityOpT, OffsetT>, typename KernelLauncherFactory = detail::TripleChevronFactory, - typename VSMemHelperPolicyT = detail::vsmem_helper_default_fallback_policy_t< - typename PolicyHub::MaxPolicy::UniqueByKeyPolicyT, - detail::unique_by_key::AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>> + typename VSMemHelperT = detail::unique_by_key:: + VSMemHelper> struct DispatchUniqueByKey { /****************************************************************************** @@ -280,11 +301,11 @@ struct DispatchUniqueByKey } // Number of input tiles - const auto block_threads = VSMemHelperPolicyT::BlockThreads(policy.UniqueByKey()); - const auto items_per_thread = VSMemHelperPolicyT::ItemsPerThread(policy.UniqueByKey()); + const auto block_threads = VSMemHelperT::BlockThreads(policy.UniqueByKey()); + const auto items_per_thread = VSMemHelperT::ItemsPerThread(policy.UniqueByKey()); int tile_size = block_threads * items_per_thread; int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); - const auto vsmem_size = num_tiles * VSMemHelperPolicyT::VSMemPerBlock(); + const auto vsmem_size = num_tiles * VSMemHelperT::VSMemPerBlock(policy.UniqueByKey()); // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0, vsmem_size}; diff --git a/cub/cub/util_vsmem.cuh b/cub/cub/util_vsmem.cuh index a23af409c04a..c282fb443cae 100644 --- a/cub/cub/util_vsmem.cuh +++ b/cub/cub/util_vsmem.cuh @@ -96,11 +96,6 @@ public: // The amount of global memory-backed virtual shared memory needed, padded to an integer multiple of 128 bytes static constexpr ::cuda::std::size_t vsmem_per_block = needs_vsmem ? (required_smem + padding_bytes) : 0; - _CCCL_HOST_DEVICE static constexpr ::cuda::std::size_t VSMemPerBlock() - { - return vsmem_per_block; - } - /** * @brief Used from within the device algorithm's kernel to get the temporary storage that can be * passed to the agent, specialized for the case when we can use native shared memory as temporary @@ -220,18 +215,6 @@ struct vsmem_helper_with_fallback_impl : public vsmem_helper_impl { using agent_t = DefaultAgentT; using agent_policy_t = DefaultAgentPolicyT; - - template - _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) - { - return AgentPolicyT::BLOCK_THREADS; - } - - template - _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) - { - return AgentPolicyT::ITEMS_PER_THREAD; - } }; template struct vsmem_helper_with_fallback_impl @@ -239,18 +222,6 @@ struct vsmem_helper_with_fallback_impl - _CCCL_HOST_DEVICE static constexpr int BlockThreads(AgentPolicyT /* policy */) - { - return FallbackAgentPolicyT::BLOCK_THREADS; - } - - template - _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(AgentPolicyT /* policy */) - { - return FallbackAgentPolicyT::ITEMS_PER_THREAD; - } }; /** From caed37f2a037716235c876f41a760d82f2537ca1 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 21 Feb 2025 17:06:01 +0000 Subject: [PATCH 11/18] Pass VSMemHelper as template parameter to unique by key kernel --- .../dispatch/dispatch_unique_by_key.cuh | 28 -------- .../device/dispatch/kernels/unique_by_key.cuh | 71 +++++++++++-------- 2 files changed, 43 insertions(+), 56 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index dce03ddd433c..ae30f8b92e33 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -87,34 +87,6 @@ struct DeviceUniqueByKeyKernelSource return ScanTileStateT(); } }; - -template -struct VSMemHelper -{ - template - CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(ActivePolicyT /*policy*/) - { - return cub::detail::vsmem_helper_default_fallback_policy_t::agent_policy_t::BLOCK_THREADS; - } - - template - CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(ActivePolicyT /*policy*/) - { - return cub::detail::vsmem_helper_default_fallback_policy_t::agent_policy_t::ITEMS_PER_THREAD; - } - - template - CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t VSMemPerBlock(ActivePolicyT /*policy*/) - { - return cub::detail::vsmem_helper_default_fallback_policy_t::vsmem_per_block; - } -}; } // namespace detail::unique_by_key /****************************************************************************** diff --git a/cub/cub/device/dispatch/kernels/unique_by_key.cuh b/cub/cub/device/dispatch/kernels/unique_by_key.cuh index 2a8cd5858a4c..81a66b4afca0 100644 --- a/cub/cub/device/dispatch/kernels/unique_by_key.cuh +++ b/cub/cub/device/dispatch/kernels/unique_by_key.cuh @@ -26,6 +26,33 @@ CUB_NAMESPACE_BEGIN namespace detail::unique_by_key { + +template +struct VSMemHelper +{ + template + using VSMemHelperDefaultFallbackPolicyT = + vsmem_helper_default_fallback_policy_t; + + template + CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(ActivePolicyT /*policy*/) + { + return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::BLOCK_THREADS; + } + + template + CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(ActivePolicyT /*policy*/) + { + return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::ITEMS_PER_THREAD; + } + + template + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t VSMemPerBlock(ActivePolicyT /*policy*/) + { + return VSMemHelperDefaultFallbackPolicyT::vsmem_per_block; + } +}; + /** * @brief Unique by key kernel entry point (multi-block) * @@ -85,25 +112,20 @@ namespace detail::unique_by_key * @param[in] vsmem * Memory to support virtual shared memory */ -template -__launch_bounds__(int( - vsmem_helper_default_fallback_policy_t< - typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, - AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>::agent_policy_t::BLOCK_THREADS)) +template < + typename ChainedPolicyT, + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename ScanTileStateT, + typename EqualityOpT, + typename OffsetT, + typename VSMemHelperT = + VSMemHelper> +__launch_bounds__(int(VSMemHelperT::template VSMemHelperDefaultFallbackPolicyT< + typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT>::agent_policy_t::BLOCK_THREADS)) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceUniqueByKeySweepKernel( KeyInputIteratorT d_keys_in, ValueInputIteratorT d_values_in, @@ -116,15 +138,8 @@ __launch_bounds__(int( int num_tiles, vsmem_t vsmem) { - using VsmemHelperT = vsmem_helper_default_fallback_policy_t< - typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, - AgentUniqueByKey, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - EqualityOpT, - OffsetT>; + using VsmemHelperT = typename VSMemHelperT::template VSMemHelperDefaultFallbackPolicyT< + typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT>; using AgentUniqueByKeyPolicyT = typename VsmemHelperT::agent_policy_t; From f37206b5370536cd9a510647a6ff2bfc94205b1e Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 24 Feb 2025 20:16:36 +0000 Subject: [PATCH 12/18] Move the template parameters of VSMemHelper to its methods to workaround an nvcc 12.0 template instantiation issue with vsmem_helper_fallback_policy_t --- .../dispatch/dispatch_unique_by_key.cuh | 77 ++++++++++++------- .../device/dispatch/kernels/unique_by_key.cuh | 56 ++++++++------ 2 files changed, 82 insertions(+), 51 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index ae30f8b92e33..85ab1c9c50bf 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -117,29 +117,27 @@ struct DeviceUniqueByKeyKernelSource * @tparam OffsetT * Signed integer type for global offsets */ -template < - typename KeyInputIteratorT, - typename ValueInputIteratorT, - typename KeyOutputIteratorT, - typename ValueOutputIteratorT, - typename NumSelectedIteratorT, - typename EqualityOpT, - typename OffsetT, - typename PolicyHub = - detail::unique_by_key::policy_hub, detail::value_t>, - typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< - typename PolicyHub::MaxPolicy, - KeyInputIteratorT, - ValueInputIteratorT, - KeyOutputIteratorT, - ValueOutputIteratorT, - NumSelectedIteratorT, - ScanTileState, - EqualityOpT, - OffsetT>, - typename KernelLauncherFactory = detail::TripleChevronFactory, - typename VSMemHelperT = detail::unique_by_key:: - VSMemHelper> +template , detail::value_t>, + typename KernelSource = detail::unique_by_key::DeviceUniqueByKeyKernelSource< + typename PolicyHub::MaxPolicy, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + NumSelectedIteratorT, + ScanTileState, + EqualityOpT, + OffsetT>, + typename KernelLauncherFactory = detail::TripleChevronFactory, + typename VSMemHelperT = detail::unique_by_key::VSMemHelper> struct DispatchUniqueByKey { /****************************************************************************** @@ -273,11 +271,34 @@ struct DispatchUniqueByKey } // Number of input tiles - const auto block_threads = VSMemHelperT::BlockThreads(policy.UniqueByKey()); - const auto items_per_thread = VSMemHelperT::ItemsPerThread(policy.UniqueByKey()); - int tile_size = block_threads * items_per_thread; - int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); - const auto vsmem_size = num_tiles * VSMemHelperT::VSMemPerBlock(policy.UniqueByKey()); + const auto block_threads = VSMemHelperT::template BlockThreads< + typename ActivePolicyT::UniqueByKeyPolicyT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>(policy.UniqueByKey()); + const auto items_per_thread = VSMemHelperT::template ItemsPerThread< + typename ActivePolicyT::UniqueByKeyPolicyT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>(policy.UniqueByKey()); + int tile_size = block_threads * items_per_thread; + int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + const auto vsmem_size = + num_tiles + * VSMemHelperT::template VSMemPerBlock< + typename ActivePolicyT::UniqueByKeyPolicyT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>(policy.UniqueByKey()); // Specify temporary storage allocation requirements size_t allocation_sizes[2] = {0, vsmem_size}; diff --git a/cub/cub/device/dispatch/kernels/unique_by_key.cuh b/cub/cub/device/dispatch/kernels/unique_by_key.cuh index 81a66b4afca0..4a1ae82653b7 100644 --- a/cub/cub/device/dispatch/kernels/unique_by_key.cuh +++ b/cub/cub/device/dispatch/kernels/unique_by_key.cuh @@ -27,29 +27,28 @@ CUB_NAMESPACE_BEGIN namespace detail::unique_by_key { -template struct VSMemHelper { - template + template using VSMemHelperDefaultFallbackPolicyT = vsmem_helper_default_fallback_policy_t; - template + template CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(ActivePolicyT /*policy*/) { - return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::BLOCK_THREADS; + return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::BLOCK_THREADS; } - template + template CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(ActivePolicyT /*policy*/) { - return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::ITEMS_PER_THREAD; + return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::ITEMS_PER_THREAD; } - template + template CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t VSMemPerBlock(ActivePolicyT /*policy*/) { - return VSMemHelperDefaultFallbackPolicyT::vsmem_per_block; + return VSMemHelperDefaultFallbackPolicyT::vsmem_per_block; } }; @@ -112,20 +111,25 @@ struct VSMemHelper * @param[in] vsmem * Memory to support virtual shared memory */ -template < - typename ChainedPolicyT, - typename KeyInputIteratorT, - typename ValueInputIteratorT, - typename KeyOutputIteratorT, - typename ValueOutputIteratorT, - typename NumSelectedIteratorT, - typename ScanTileStateT, - typename EqualityOpT, - typename OffsetT, - typename VSMemHelperT = - VSMemHelper> -__launch_bounds__(int(VSMemHelperT::template VSMemHelperDefaultFallbackPolicyT< - typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT>::agent_policy_t::BLOCK_THREADS)) +template +__launch_bounds__(int( + VSMemHelperT::template VSMemHelperDefaultFallbackPolicyT< + typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>::agent_policy_t::BLOCK_THREADS)) CUB_DETAIL_KERNEL_ATTRIBUTES void DeviceUniqueByKeySweepKernel( KeyInputIteratorT d_keys_in, ValueInputIteratorT d_values_in, @@ -139,7 +143,13 @@ __launch_bounds__(int(VSMemHelperT::template VSMemHelperDefaultFallbackPolicyT< vsmem_t vsmem) { using VsmemHelperT = typename VSMemHelperT::template VSMemHelperDefaultFallbackPolicyT< - typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT>; + typename ChainedPolicyT::ActivePolicy::UniqueByKeyPolicyT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyOutputIteratorT, + ValueOutputIteratorT, + EqualityOpT, + OffsetT>; using AgentUniqueByKeyPolicyT = typename VsmemHelperT::agent_policy_t; From d1c885c6ada218d2a8e65ff3de2e3e8ae89173ec Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 24 Feb 2025 22:42:16 +0000 Subject: [PATCH 13/18] Make vsmem helper functions host device --- cub/cub/device/dispatch/kernels/unique_by_key.cuh | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/cub/cub/device/dispatch/kernels/unique_by_key.cuh b/cub/cub/device/dispatch/kernels/unique_by_key.cuh index 81a66b4afca0..6a81a8d7b0e5 100644 --- a/cub/cub/device/dispatch/kernels/unique_by_key.cuh +++ b/cub/cub/device/dispatch/kernels/unique_by_key.cuh @@ -18,6 +18,11 @@ #include #include +#define XSTR(x) STR(x) +#define STR(x) #x + +#pragma message "The value of ABC: " XSTR(CUB_PTX_ARCH) + CUB_NAMESPACE_BEGIN /****************************************************************************** @@ -35,19 +40,19 @@ struct VSMemHelper vsmem_helper_default_fallback_policy_t; template - CUB_RUNTIME_FUNCTION static constexpr int BlockThreads(ActivePolicyT /*policy*/) + _CCCL_HOST_DEVICE static constexpr int BlockThreads(ActivePolicyT /*policy*/) { return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::BLOCK_THREADS; } template - CUB_RUNTIME_FUNCTION static constexpr int ItemsPerThread(ActivePolicyT /*policy*/) + _CCCL_HOST_DEVICE static constexpr int ItemsPerThread(ActivePolicyT /*policy*/) { return VSMemHelperDefaultFallbackPolicyT::agent_policy_t::ITEMS_PER_THREAD; } template - CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t VSMemPerBlock(ActivePolicyT /*policy*/) + _CCCL_HOST_DEVICE static constexpr ::cuda::std::size_t VSMemPerBlock(ActivePolicyT /*policy*/) { return VSMemHelperDefaultFallbackPolicyT::vsmem_per_block; } From c6ef86f90921e4a87519506d94f01b3ec934e7bf Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 24 Feb 2025 22:54:50 +0000 Subject: [PATCH 14/18] Remove unused macro --- cub/cub/device/dispatch/kernels/unique_by_key.cuh | 5 ----- 1 file changed, 5 deletions(-) diff --git a/cub/cub/device/dispatch/kernels/unique_by_key.cuh b/cub/cub/device/dispatch/kernels/unique_by_key.cuh index 3ca67513eea4..92350a118278 100644 --- a/cub/cub/device/dispatch/kernels/unique_by_key.cuh +++ b/cub/cub/device/dispatch/kernels/unique_by_key.cuh @@ -18,11 +18,6 @@ #include #include -#define XSTR(x) STR(x) -#define STR(x) #x - -#pragma message "The value of ABC: " XSTR(CUB_PTX_ARCH) - CUB_NAMESPACE_BEGIN /****************************************************************************** From 4d0372c4e5f9737b74881ebfec682b462dc558c3 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 25 Feb 2025 01:18:22 +0000 Subject: [PATCH 15/18] Make KeyT and ValueT templates for DispatchUniqueByKey --- cub/cub/device/dispatch/dispatch_unique_by_key.cuh | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index e3601a8fe3eb..f1b13e9a8bd2 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -138,7 +138,9 @@ template < EqualityOpT, OffsetT>, typename KernelLauncherFactory = detail::TripleChevronFactory, - typename VSMemHelperT = detail::unique_by_key::VSMemHelper> + typename VSMemHelperT = detail::unique_by_key::VSMemHelper, + typename KeyT = detail::it_value_t, + typename ValueT = detail::it_value_t> struct DispatchUniqueByKey { /****************************************************************************** @@ -150,10 +152,6 @@ struct DispatchUniqueByKey INIT_KERNEL_THREADS = 128, }; - // The input key and value type - using KeyT = detail::it_value_t; - using ValueT = detail::it_value_t; - /// Device-accessible allocation of temporary storage. When nullptr, the required allocation size /// is written to `temp_storage_bytes` and no work is done. void* d_temp_storage; From 895852dd8fac6e915a7cf2265f5ebe704ccefa6b Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Tue, 25 Feb 2025 23:15:36 +0000 Subject: [PATCH 16/18] Fix SFINAE for UniqueByKeyPolicyWrapper --- cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh b/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh index cd570373933a..c3e146f11360 100644 --- a/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh +++ b/cub/cub/device/dispatch/tuning/tuning_unique_by_key.cuh @@ -776,7 +776,8 @@ struct UniqueByKeyPolicyWrapper : PolicyT }; template -struct UniqueByKeyPolicyWrapper> +struct UniqueByKeyPolicyWrapper> : StaticPolicyT { CUB_RUNTIME_FUNCTION UniqueByKeyPolicyWrapper(StaticPolicyT base) From fa0f687a1ebe72a815b7a5dfeb31dbffce2eed67 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Wed, 26 Feb 2025 16:09:54 +0000 Subject: [PATCH 17/18] Add comment explaining vsmem design --- cub/cub/device/dispatch/kernels/unique_by_key.cuh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cub/cub/device/dispatch/kernels/unique_by_key.cuh b/cub/cub/device/dispatch/kernels/unique_by_key.cuh index 92350a118278..3a9b31e83d77 100644 --- a/cub/cub/device/dispatch/kernels/unique_by_key.cuh +++ b/cub/cub/device/dispatch/kernels/unique_by_key.cuh @@ -27,6 +27,9 @@ CUB_NAMESPACE_BEGIN namespace detail::unique_by_key { +// TODO: this class should be templated on `typename... Ts` to avoid repetition, +// but due to an issue with NVCC 12.0 we currently template each member function +// individually instead. struct VSMemHelper { template From 89a51d5362fec2c99426d52ae9f91d0af70da0ff Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Fri, 28 Feb 2025 16:04:41 +0000 Subject: [PATCH 18/18] Get PTX version with launcher factory to avoid unsupported toolchain error --- cub/cub/device/dispatch/dispatch_unique_by_key.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh index f1b13e9a8bd2..d17b184737ff 100644 --- a/cub/cub/device/dispatch/dispatch_unique_by_key.cuh +++ b/cub/cub/device/dispatch/dispatch_unique_by_key.cuh @@ -502,7 +502,7 @@ struct DispatchUniqueByKey { // Get PTX version int ptx_version = 0; - error = CubDebug(PtxVersion(ptx_version)); + error = CubDebug(launcher_factory.PtxVersion(ptx_version)); if (cudaSuccess != error) { break;