diff --git a/cub/cub/agent/agent_topk.cuh b/cub/cub/agent/agent_topk.cuh index 2193a82b2bec..505a66510d1a 100644 --- a/cub/cub/agent/agent_topk.cuh +++ b/cub/cub/agent/agent_topk.cuh @@ -23,7 +23,10 @@ #include #include +#include + #include +#include #include CUB_NAMESPACE_BEGIN @@ -47,11 +50,15 @@ namespace detail::topk //! @tparam ScanAlgorithm //! The BlockScan algorithm to use //! +//! @tparam VecSize +//! Number of keys per vectorized load of contiguous input with BLOCK_LOAD_VECTORIZE, 0 to load through BlockLoad +//! template + BlockScanAlgorithm ScanAlgorithm, + int VecSize> struct agent_topk_policy { static constexpr int threads_per_block = ThreadsPerBlock; @@ -59,6 +66,7 @@ struct agent_topk_policy static constexpr int bits_per_pass = BitsPerPass; static constexpr BlockLoadAlgorithm load_algorithm = LoadAlgorithm; static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + static constexpr int vec_size = VecSize; }; template @@ -258,6 +266,22 @@ struct AgentTopK // Utility methods for device topK //--------------------------------------------------------------------- + // Load a full tile in a blocked arrangement + template + _CCCL_DEVICE _CCCL_FORCEINLINE void load_tile(InputItT block_in, key_in_t (&thread_data)[items_per_thread]) + { + if constexpr (AgentTopKPolicyT::load_algorithm == BLOCK_LOAD_VECTORIZE && AgentTopKPolicyT::vec_size != 0 + && THRUST_NS_QUALIFIER::is_contiguous_iterator_v) + { + LoadDirectBlockedVectorized( + static_cast(threadIdx.x), ::cuda::std::to_address(block_in), thread_data); + } + else + { + block_load_input_t(temp_storage.load_input).Load(block_in, thread_data); + } + } + // Process a range of input data in tiles, calling f(key, index) for each element template _CCCL_DEVICE _CCCL_FORCEINLINE void process_range(InputItT in, const OffsetT num_items, FuncT f) @@ -280,7 +304,7 @@ struct AgentTopK // Ensure that the temporary storage from previous iteration can be reused __syncthreads(); - block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data); + load_tile(in + tile_base, thread_data); for (int j = 0; j < items_per_thread; ++j) { f(thread_data[j], offset + j); @@ -297,7 +321,7 @@ struct AgentTopK if (num_remaining_elements == 0) { - block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data); + load_tile(in + tile_base, thread_data); } else { diff --git a/cub/cub/block/block_load.cuh b/cub/cub/block/block_load.cuh index e42653a89ac5..d2326f6ae7e2 100644 --- a/cub/cub/block/block_load.cuh +++ b/cub/cub/block/block_load.cuh @@ -23,19 +23,21 @@ #include #include #include +#include #include #include #include #include +#include +#include #include #include #include #include #include -#include -#include -#include + +#include CUB_NAMESPACE_BEGIN @@ -182,61 +184,51 @@ _CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectBlocked( #ifndef _CCCL_DOXYGEN_INVOKED // Do not document -//! @brief Internal implementation for load vectorization +namespace detail +{ +//! Whether the storage of ``RandomAccessIterator`` can be loaded with vectors +// TODO(bgruber): remove the check for is_primitive in CCCL 4.0 +template > +inline constexpr bool block_load_can_vectorize_v = + (is_primitive::value || ::cuda::is_trivially_copyable_v) // + &&!::cuda::std::is_volatile_v + && THRUST_NS_QUALIFIER::is_contiguous_iterator_v; + +//! @brief Internal implementation of LoadDirectBlockedVectorized for input aligned to the vector type //! //! @param[in] linear_tid //! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D //! thread blocks) //! //! @param[in] block_src_ptr -//! Input pointer for loading from +//! Input pointer for loading from, aligned to ``cuda::__as_vectorized_storage_t`` //! //! @param[out] dst_items //! Destination to load data into -template +template _CCCL_DEVICE _CCCL_FORCEINLINE void -InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&dst_items)[ItemsPerThread]) +InternalLoadDirectBlockedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst_items)[ItemsPerThread]) { - // Find biggest memory access word that T is a whole multiple of - using device_word_t = typename UnitWord::DeviceWord; - _CCCL_DIAG_PUSH -# if _CCCL_COMPILER(CLANG, >=, 10) - _CCCL_DIAG_SUPPRESS_CLANG("-Wsizeof-array-div") -# endif // _CCCL_COMPILER(CLANG, >=, 10) - // NOLINTNEXTLINE(bugprone-sizeof-expression) - constexpr int total_words = static_cast(sizeof(dst_items) / sizeof(device_word_t)); - _CCCL_DIAG_POP - constexpr int vector_size = (total_words % 4 == 0) ? 4 : (total_words % 2 == 0) ? 2 : 1; - constexpr int vectors_per_thread = total_words / vector_size; - - // Load into an array of vectors in thread-blocked order - using vector_t = typename CubVector::Type; - - // Add the alignment check to ensure the vectorized loading can proceed. - if (::cuda::std::is_sufficiently_aligned(block_src_ptr)) - { - vector_t vec_items[vectors_per_thread]; - // Load into an array of vectors in thread-blocked order - const vector_t* vec_ptr = reinterpret_cast(block_src_ptr) + linear_tid * vectors_per_thread; + // All items of a thread are loaded with the widest vectors that evenly divide them + using vector_t = ::cuda::__as_vectorized_storage_t; + _CCCL_ASSERT(::cuda::std::is_sufficiently_aligned(block_src_ptr), + "block_src_ptr must be aligned to the vector type"); - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < vectors_per_thread; i++) - { - vec_items[i] = ThreadLoad(vec_ptr + i); - } + // Load in thread-blocked order + const vector_t vec_items = ThreadLoad(reinterpret_cast(block_src_ptr) + linear_tid); - // Copy to destination - _CCCL_PRAGMA_UNROLL_FULL() - for (int i = 0; i < ItemsPerThread; i++) - { - dst_items[i] = *(reinterpret_cast(vec_items) + i); - } - } - else + // Copy to destination + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) { - LoadDirectBlocked(linear_tid, block_src_ptr, dst_items); + dst_items[i] = vec_items.template __alias(i); } } +} // namespace detail #endif // _CCCL_DOXYGEN_INVOKED @@ -258,12 +250,23 @@ InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (& //! //! @endrst //! +//! @tparam Modifier +//! **[optional]** Cache load modifier for the vectorized loads (default: ``LOAD_DEFAULT``) +//! +//! @tparam VecSize +//! **[optional]** The number of consecutive items in each vector. Must evenly divide ``ItemsPerThread``. Each thread +//! loads its items with ``ItemsPerThread / VecSize`` vectors, and ``0`` loads them with a single vector +//! (default: ``0``) +//! //! @tparam T //! **[inferred]** The data type to load. //! //! @tparam ItemsPerThread //! **[inferred]** The number of consecutive items partitioned onto each thread. //! +//! @tparam InputT +//! **[inferred]** The value type of the input, which is loaded with vectors and converted to ``T``. +//! //! @param[in] linear_tid //! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + //! linear_tid` for 2D thread blocks) @@ -273,12 +276,75 @@ InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (& //! //! @param[out] dst_items //! destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +LoadDirectBlockedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst_items)[ItemsPerThread]) +{ + constexpr int vec_size = VecSize == 0 ? ItemsPerThread : VecSize; + static_assert(ItemsPerThread % vec_size == 0, "VecSize must evenly divide ItemsPerThread"); + + if constexpr (detail::block_load_can_vectorize_v) + { + // Vector i of a thread is the blocked vector of thread linear_tid * vectors_per_thread + i + constexpr int vectors_per_thread = ItemsPerThread / vec_size; + + // Only sm_100 and newer load vectors wider than 16 bytes with a single access, and only without a cache modifier + using vector16_t = ::cuda::__as_vectorized_storage_t; + +#if __cccl_ptx_isa >= 900ULL + using vector32_t = ::cuda::__as_vectorized_storage_t; + if constexpr (Modifier == LOAD_DEFAULT && alignof(vector16_t) < alignof(vector32_t)) + { + bool is_vector32_aligned = false; + NV_IF_TARGET(NV_PROVIDES_SM_100, + (is_vector32_aligned = ::cuda::std::is_sufficiently_aligned(block_src_ptr);)) + if (is_vector32_aligned) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < vectors_per_thread; i++) + { + detail::InternalLoadDirectBlockedVectorized( + linear_tid * vectors_per_thread + i, + block_src_ptr, + reinterpret_cast(dst_items[i * vec_size])); + } + return; + } + } +#endif // __cccl_ptx_isa >= 900ULL + + // Add the alignment check to ensure the vectorized loading can proceed. + if (::cuda::std::is_sufficiently_aligned(block_src_ptr)) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < vectors_per_thread; i++) + { + detail::InternalLoadDirectBlockedVectorized( + linear_tid * vectors_per_thread + i, block_src_ptr, reinterpret_cast(dst_items[i * vec_size])); + } + return; + } + } + + if constexpr (Modifier != LOAD_DEFAULT) + { + LoadDirectBlocked(linear_tid, CacheModifiedInputIterator(block_src_ptr), dst_items); + } + else + { + LoadDirectBlocked(linear_tid, block_src_ptr, dst_items); + } +} + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document +// Accepts explicitly specified T and ItemsPerThread, which cannot bind to the leading Modifier of the overload above template _CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectBlockedVectorized(int linear_tid, T* block_src_ptr, T (&dst_items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_src_ptr, dst_items); + LoadDirectBlockedVectorized(linear_tid, block_src_ptr, dst_items); } +#endif // _CCCL_DOXYGEN_INVOKED //! @} //! @name Striped arrangement I/O (direct) @@ -1015,18 +1081,13 @@ public: } else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE) { - if constexpr (detail::is_CacheModifiedInputIterator - && ::cuda::std::is_same_v, T>) + if constexpr (detail::is_CacheModifiedInputIterator) { - InternalLoadDirectBlockedVectorized(linear_tid, block_src_it.ptr, dst_items); + LoadDirectBlockedVectorized(linear_tid, block_src_it.ptr, dst_items); } - else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v - && ::cuda::std::is_same_v, T> - // handles contiguous iterators wrapping volatile pointers, like span::iterator - && !::cuda::std::is_volatile_v< - ::cuda::std::remove_reference_t>>) + else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v) { - InternalLoadDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_src_it), dst_items); + LoadDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_src_it), dst_items); } else { diff --git a/cub/cub/device/dispatch/dispatch_topk.cuh b/cub/cub/device/dispatch/dispatch_topk.cuh index 288fa589b5fa..741208cf569c 100644 --- a/cub/cub/device/dispatch/dispatch_topk.cuh +++ b/cub/cub/device/dispatch/dispatch_topk.cuh @@ -84,7 +84,8 @@ __launch_bounds__(int(current_policy().threads_per_block)) policy.items_per_thread, policy.bits_per_pass, policy.load_algorithm, - policy.scan_algorithm>; + policy.scan_algorithm, + policy.vec_size>; using agent_topk_t = AgentTopK().threads_per_block)) policy.items_per_thread, policy.bits_per_pass, policy.load_algorithm, - policy.scan_algorithm>; + policy.scan_algorithm, + policy.vec_size>; using identify_candidates_op_t = NullType; using agent_topk_t = AgentTopK().threads_per_block)) policy.items_per_thread, policy.bits_per_pass, policy.load_algorithm, - policy.scan_algorithm>; + policy.scan_algorithm, + policy.last_filter_vec_size>; using extract_bin_op_t = NullType; using agent_topk_t = AgentTopK 2, int32 -> 4, int16 -> 8. const int items_per_thread = ::cuda::std::max(1, nominal_4b_items_per_thread * 4 / key_size); - return topk_policy{512, items_per_thread, BLOCK_LOAD_VECTORIZE, BLOCK_SCAN_WARP_SCANS, bits_per_pass}; + auto policy = topk_policy{512, items_per_thread, BLOCK_LOAD_VECTORIZE, BLOCK_SCAN_WARP_SCANS, bits_per_pass}; + + // ptxas predicates a single vector of all 8- or 16-bit keys of a thread together with its fallback. Where that + // costs occupancy or spills against vectors of 4 keys with CUDA 13.4, the keys are loaded in vectors of 4 keys + if (key_size <= 2) + { + const bool is_sm90 = cc < ::cuda::compute_capability{10, 0}; + const bool is_sm120 = cc >= ::cuda::compute_capability{12, 0}; + const bool narrow_all = !is_sm90 && offset_size == 4; + const bool narrow_passes = is_sm90 && offset_size == 8 + && (value_size == 0 || value_size == 8 || (value_size == 1 && out_offset_size == 4)); + const bool narrow_last_filter = (is_sm90 || is_sm120) && key_size == 1; + policy.vec_size = narrow_all || narrow_passes ? 4 : 0; + policy.last_filter_vec_size = narrow_all || narrow_last_filter ? 4 : 0; + } + return policy; } // Default tuning used on older architectures. diff --git a/cub/cub/warp/warp_load.cuh b/cub/cub/warp/warp_load.cuh index 3f44400ce1ec..b8fcff1810cb 100644 --- a/cub/cub/warp/warp_load.cuh +++ b/cub/cub/warp/warp_load.cuh @@ -335,19 +335,19 @@ private: _CCCL_DEVICE _CCCL_FORCEINLINE void Load(InputT* block_ptr, InputT (&items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_ptr, items); + LoadDirectBlockedVectorized(linear_tid, block_ptr, items); } _CCCL_DEVICE _CCCL_FORCEINLINE void Load(const InputT* block_ptr, InputT (&items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_ptr, items); + LoadDirectBlockedVectorized(linear_tid, block_ptr, items); } template _CCCL_DEVICE _CCCL_FORCEINLINE void Load(CacheModifiedInputIterator block_itr, InputT (&items)[ItemsPerThread]) { - InternalLoadDirectBlockedVectorized(linear_tid, block_itr.ptr, items); + LoadDirectBlockedVectorized(linear_tid, block_itr.ptr, items); } template diff --git a/cub/test/catch2_test_block_load.cu b/cub/test/catch2_test_block_load.cu index b5ff586dbded..eaf922a3bc88 100644 --- a/cub/test/catch2_test_block_load.cu +++ b/cub/test/catch2_test_block_load.cu @@ -208,10 +208,10 @@ CUB_TEST("Block load works with caching iterators", "[load][block]", CUB_SMALL, } #if IPT == 1 -CUB_TEST("Vectorized block load with const and non-const datatype and different alignment cases", +CUB_TEST("Vectorized block load with const, non-const and cache modified inputs and different alignment cases", "[load][block]", CUB_SMALL, - c2h::type_list) + c2h::type_list>) { using type = int; using input_ptr_type = c2h::get<0, TestType>; @@ -270,5 +270,86 @@ CUB_TEST_CASE("Vectorized block load converts from iterators of a different valu test_vectorized_full_tile_load( d_input, cub::CacheModifiedInputIterator(ptr)); test_vectorized_full_tile_load(d_input, cuda::std::span(ptr, d_input.size()).begin()); + test_vectorized_full_tile_load(d_input, static_cast(ptr)); +} + +template +__global__ void load_direct_blocked_vectorized_kernel(InputIteratorT input, T* output) +{ + T data[ItemsPerThread]; + cub::LoadDirectBlockedVectorized( + static_cast(threadIdx.x), input, data); + for (int i = 0; i < ItemsPerThread; i++) + { + output[threadIdx.x * ItemsPerThread + i] = data[i]; + } +} + +template +void test_load_direct_blocked_vectorized(const c2h::device_vector& d_expected, InputIteratorT input) +{ + REQUIRE(d_expected.size() == ItemsPerThread * ThreadsInBlock); + + c2h::device_vector d_output(d_expected.size(), thrust::no_init); + load_direct_blocked_vectorized_kernel + <<<1, ThreadsInBlock>>>(input, thrust::raw_pointer_cast(d_output.data())); + REQUIRE(cudaSuccess == cudaPeekAtLastError()); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + REQUIRE(d_expected == d_output); +} + +template +__global__ void load_direct_blocked_vectorized_type_first_kernel(T* input, T* output) +{ + T data[ItemsPerThread]; + cub::LoadDirectBlockedVectorized(static_cast(threadIdx.x), input, data); + for (int i = 0; i < ItemsPerThread; i++) + { + output[threadIdx.x * ItemsPerThread + i] = data[i]; + } +} + +CUB_TEST("LoadDirectBlockedVectorized accepts an explicitly specified data type and items per thread", + "[load][block]", + CUB_SMALL, + types) +{ + using type = c2h::get<0, TestType>; + constexpr int items_per_thread = 4; + constexpr int threads_in_block = 64; + + c2h::device_vector d_input(items_per_thread * threads_in_block); + c2h::gen(C2H_SEED(1), d_input); + c2h::device_vector d_output(d_input.size(), thrust::no_init); + load_direct_blocked_vectorized_type_first_kernel + <<<1, threads_in_block>>>(thrust::raw_pointer_cast(d_input.data()), thrust::raw_pointer_cast(d_output.data())); + REQUIRE(cudaSuccess == cudaPeekAtLastError()); + REQUIRE(cudaSuccess == cudaDeviceSynchronize()); + REQUIRE(d_input == d_output); +} + +// Offset 0 takes the vectorized loads. Offset 1 misaligns the vectors of multiple items, which use the scalar fallback. +// An offset of 16 bytes misaligns the vectors wider than 16 bytes, which fall back to 16-byte vectors. +CUB_TEST("LoadDirectBlockedVectorized loads the items of a thread with vectors of VecSize items", + "[load][block]", + CUB_SMALL, + types) +{ + using type = c2h::get<0, TestType>; + constexpr int items_per_thread = 16; + constexpr int threads_in_block = 64; + constexpr int tile_size = items_per_thread * threads_in_block; + + const int offset = GENERATE(0, 1, static_cast(16 / sizeof(type))); + c2h::device_vector d_expected(tile_size); + c2h::gen(C2H_SEED(1), d_expected); + c2h::device_vector d_input(tile_size + offset); + thrust::copy_n(d_expected.begin(), tile_size, d_input.begin() + offset); + const type* input = thrust::raw_pointer_cast(d_input.data()) + offset; + + test_load_direct_blocked_vectorized<0, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<1, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<4, items_per_thread, threads_in_block>(d_expected, input); + test_load_direct_blocked_vectorized<16, items_per_thread, threads_in_block>(d_expected, input); } #endif