codekingpro/portable-devtools
114k
1/******************************************************************************2 * Copyright (c) 2011, Duane Merrill. All rights reserved.3 * Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.4 *5 * Redistribution and use in source and binary forms, with or without6 * modification, are permitted provided that the following conditions are met:7 * * Redistributions of source code must retain the above copyright8 * notice, this list of conditions and the following disclaimer.9 * * Redistributions in binary form must reproduce the above copyright10 * notice, this list of conditions and the following disclaimer in the11 * documentation and/or other materials provided with the distribution.12 * * Neither the name of the NVIDIA CORPORATION nor the13 * names of its contributors may be used to endorse or promote products14 * derived from this software without specific prior written permission.15 *16 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"17 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE18 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE19 * ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY20 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES21 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;22 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND23 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT24 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS25 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.26 *27 ******************************************************************************/28 29/**30 * @file cub::AgentReduce implements a stateful abstraction of CUDA thread31 * blocks for participating in device-wide reduction.32 */33 34#pragma once35 36#include <cub/config.cuh>37 38#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)39# pragma GCC system_header40#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)41# pragma clang system_header42#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)43# pragma system_header44#endif // no system header45 46#include <iterator>47 48#include <cub/block/block_load.cuh>49#include <cub/block/block_reduce.cuh>50#include <cub/detail/type_traits.cuh>51#include <cub/grid/grid_even_share.cuh>52#include <cub/grid/grid_mapping.cuh>53#include <cub/iterator/cache_modified_input_iterator.cuh>54#include <cub/util_type.cuh>55 56CUB_NAMESPACE_BEGIN57 58/******************************************************************************59 * Tuning policy types60 ******************************************************************************/61 62/**63 * Parameterizable tuning policy type for AgentReduce64 * @tparam NOMINAL_BLOCK_THREADS_4B Threads per thread block65 * @tparam NOMINAL_ITEMS_PER_THREAD_4B Items per thread (per tile of input)66 * @tparam ComputeT Dominant compute type67 * @tparam _VECTOR_LOAD_LENGTH Number of items per vectorized load68 * @tparam _BLOCK_ALGORITHM Cooperative block-wide reduction algorithm to use69 * @tparam _LOAD_MODIFIER Cache load modifier for reading input elements70 */71template <int NOMINAL_BLOCK_THREADS_4B,72 int NOMINAL_ITEMS_PER_THREAD_4B,73 typename ComputeT,74 int _VECTOR_LOAD_LENGTH,75 BlockReduceAlgorithm _BLOCK_ALGORITHM,76 CacheLoadModifier _LOAD_MODIFIER,77 typename ScalingType = MemBoundScaling<NOMINAL_BLOCK_THREADS_4B,78 NOMINAL_ITEMS_PER_THREAD_4B,79 ComputeT>>80struct AgentReducePolicy : ScalingType81{82 /// Number of items per vectorized load83 static constexpr int VECTOR_LOAD_LENGTH = _VECTOR_LOAD_LENGTH;84 85 /// Cooperative block-wide reduction algorithm to use86 static constexpr BlockReduceAlgorithm BLOCK_ALGORITHM = _BLOCK_ALGORITHM;87 88 /// Cache load modifier for reading input elements89 static constexpr CacheLoadModifier LOAD_MODIFIER = _LOAD_MODIFIER;90};91 92/******************************************************************************93 * Thread block abstractions94 ******************************************************************************/95 96/**97 * @brief AgentReduce implements a stateful abstraction of CUDA thread blocks98 * for participating in device-wide reduction .99 *100 * Each thread reduces only the values it loads. If `FIRST_TILE`, this partial101 * reduction is stored into `thread_aggregate`. Otherwise it is accumulated102 * into `thread_aggregate`.103 *104 * @tparam AgentReducePolicy105 * Parameterized AgentReducePolicy tuning policy type106 *107 * @tparam InputIteratorT108 * Random-access iterator type for input109 *110 * @tparam OutputIteratorT111 * Random-access iterator type for output112 *113 * @tparam OffsetT114 * Signed integer type for global offsets115 *116 * @tparam ReductionOp117 * Binary reduction operator type having member118 * `auto operator()(T &&a, U &&b)`119 *120 * @tparam AccumT121 * The type of intermediate accumulator (according to P2322R6)122 */123template <typename AgentReducePolicy,124 typename InputIteratorT,125 typename OutputIteratorT,126 typename OffsetT,127 typename ReductionOp,128 typename AccumT>129struct AgentReduce130{131 //---------------------------------------------------------------------132 // Types and constants133 //---------------------------------------------------------------------134 135 /// The input value type136 using InputT = cub::detail::value_t<InputIteratorT>;137 138 /// Vector type of InputT for data movement139 using VectorT =140 typename CubVector<InputT, AgentReducePolicy::VECTOR_LOAD_LENGTH>::Type;141 142 /// Input iterator wrapper type (for applying cache modifier)143 // Wrap the native input pointer with CacheModifiedInputIterator144 // or directly use the supplied input iterator type145 using WrappedInputIteratorT = cub::detail::conditional_t<146 std::is_pointer<InputIteratorT>::value,147 CacheModifiedInputIterator<AgentReducePolicy::LOAD_MODIFIER, InputT, OffsetT>,148 InputIteratorT>;149 150 /// Constants151 static constexpr int BLOCK_THREADS = AgentReducePolicy::BLOCK_THREADS;152 static constexpr int ITEMS_PER_THREAD = AgentReducePolicy::ITEMS_PER_THREAD;153 static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD;154 static constexpr int VECTOR_LOAD_LENGTH =155 CUB_MIN(ITEMS_PER_THREAD, AgentReducePolicy::VECTOR_LOAD_LENGTH);156 157 // Can vectorize according to the policy if the input iterator is a native158 // pointer to a primitive type159 static constexpr bool ATTEMPT_VECTORIZATION = (VECTOR_LOAD_LENGTH > 1) &&160 (ITEMS_PER_THREAD % VECTOR_LOAD_LENGTH == 0) &&161 (std::is_pointer<InputIteratorT>::value) &&162 Traits<InputT>::PRIMITIVE;163 164 static constexpr CacheLoadModifier LOAD_MODIFIER =165 AgentReducePolicy::LOAD_MODIFIER;166 167 static constexpr BlockReduceAlgorithm BLOCK_ALGORITHM =168 AgentReducePolicy::BLOCK_ALGORITHM;169 170 /// Parameterized BlockReduce primitive171 using BlockReduceT =172 BlockReduce<AccumT, BLOCK_THREADS, AgentReducePolicy::BLOCK_ALGORITHM>;173 174 /// Shared memory type required by this thread block175 struct _TempStorage176 {177 typename BlockReduceT::TempStorage reduce;178 };179 180 /// Alias wrapper allowing storage to be unioned181 struct TempStorage : Uninitialized<_TempStorage>182 {};183 184 //---------------------------------------------------------------------185 // Per-thread fields186 //---------------------------------------------------------------------187 188 _TempStorage &temp_storage; ///< Reference to temp_storage189 InputIteratorT d_in; ///< Input data to reduce190 WrappedInputIteratorT d_wrapped_in; ///< Wrapped input data to reduce191 ReductionOp reduction_op; ///< Binary reduction operator192 193 //---------------------------------------------------------------------194 // Utility195 //---------------------------------------------------------------------196 197 // Whether or not the input is aligned with the vector type (specialized for198 // types we can vectorize)199 template <typename Iterator>200 static __device__ __forceinline__ bool201 IsAligned(Iterator d_in, Int2Type<true> /*can_vectorize*/)202 {203 return (size_t(d_in) & (sizeof(VectorT) - 1)) == 0;204 }205 206 // Whether or not the input is aligned with the vector type (specialized for207 // types we cannot vectorize)208 template <typename Iterator>209 static __device__ __forceinline__ bool210 IsAligned(Iterator /*d_in*/, Int2Type<false> /*can_vectorize*/)211 {212 return false;213 }214 215 //---------------------------------------------------------------------216 // Constructor217 //---------------------------------------------------------------------218 219 /**220 * @brief Constructor221 * @param temp_storage Reference to temp_storage222 * @param d_in Input data to reduce223 * @param reduction_op Binary reduction operator224 */225 __device__ __forceinline__ AgentReduce(TempStorage &temp_storage,226 InputIteratorT d_in,227 ReductionOp reduction_op)228 : temp_storage(temp_storage.Alias())229 , d_in(d_in)230 , d_wrapped_in(d_in)231 , reduction_op(reduction_op)232 {}233 234 //---------------------------------------------------------------------235 // Tile consumption236 //---------------------------------------------------------------------237 238 /**239 * @brief Consume a full tile of input (non-vectorized)240 * @param block_offset The offset the tile to consume241 * @param valid_items The number of valid items in the tile242 * @param is_full_tile Whether or not this is a full tile243 * @param can_vectorize Whether or not we can vectorize loads244 */245 template <int IS_FIRST_TILE>246 __device__ __forceinline__ void ConsumeTile(AccumT &thread_aggregate,247 OffsetT block_offset,248 int /*valid_items*/,249 Int2Type<true> /*is_full_tile*/,250 Int2Type<false> /*can_vectorize*/)251 {252 AccumT items[ITEMS_PER_THREAD];253 254 // Load items in striped fashion255 LoadDirectStriped<BLOCK_THREADS>(threadIdx.x,256 d_wrapped_in + block_offset,257 items);258 259 // Reduce items within each thread stripe260 thread_aggregate =261 (IS_FIRST_TILE)262 ? internal::ThreadReduce(items, reduction_op)263 : internal::ThreadReduce(items, reduction_op, thread_aggregate);264 }265 266 /**267 * Consume a full tile of input (vectorized)268 * @param block_offset The offset the tile to consume269 * @param valid_items The number of valid items in the tile270 * @param is_full_tile Whether or not this is a full tile271 * @param can_vectorize Whether or not we can vectorize loads272 */273 template <int IS_FIRST_TILE>274 __device__ __forceinline__ void ConsumeTile(AccumT &thread_aggregate,275 OffsetT block_offset,276 int /*valid_items*/,277 Int2Type<true> /*is_full_tile*/,278 Int2Type<true> /*can_vectorize*/)279 {280 // Alias items as an array of VectorT and load it in striped fashion281 enum282 {283 WORDS = ITEMS_PER_THREAD / VECTOR_LOAD_LENGTH284 };285 286 // Fabricate a vectorized input iterator287 InputT *d_in_unqualified = const_cast<InputT *>(d_in) + block_offset +288 (threadIdx.x * VECTOR_LOAD_LENGTH);289 CacheModifiedInputIterator<AgentReducePolicy::LOAD_MODIFIER, VectorT, OffsetT>290 d_vec_in(reinterpret_cast<VectorT *>(d_in_unqualified));291 292 // Load items as vector items293 InputT input_items[ITEMS_PER_THREAD];294 VectorT *vec_items = reinterpret_cast<VectorT *>(input_items);295#pragma unroll296 for (int i = 0; i < WORDS; ++i)297 {298 vec_items[i] = d_vec_in[BLOCK_THREADS * i];299 }300 301 // Convert from input type to output type302 AccumT items[ITEMS_PER_THREAD];303#pragma unroll304 for (int i = 0; i < ITEMS_PER_THREAD; ++i)305 {306 items[i] = input_items[i];307 }308 309 // Reduce items within each thread stripe310 thread_aggregate =311 (IS_FIRST_TILE)312 ? internal::ThreadReduce(items, reduction_op)313 : internal::ThreadReduce(items, reduction_op, thread_aggregate);314 }315 316 /**317 * Consume a partial tile of input318 * @param block_offset The offset the tile to consume319 * @param valid_items The number of valid items in the tile320 * @param is_full_tile Whether or not this is a full tile321 * @param can_vectorize Whether or not we can vectorize loads322 */323 template <int IS_FIRST_TILE, int CAN_VECTORIZE>324 __device__ __forceinline__ void325 ConsumeTile(AccumT &thread_aggregate,326 OffsetT block_offset,327 int valid_items,328 Int2Type<false> /*is_full_tile*/,329 Int2Type<CAN_VECTORIZE> /*can_vectorize*/)330 {331 // Partial tile332 int thread_offset = threadIdx.x;333 334 // Read first item335 if ((IS_FIRST_TILE) && (thread_offset < valid_items))336 {337 thread_aggregate = d_wrapped_in[block_offset + thread_offset];338 thread_offset += BLOCK_THREADS;339 }340 341 // Continue reading items (block-striped)342 while (thread_offset < valid_items)343 {344 InputT item(d_wrapped_in[block_offset + thread_offset]);345 346 thread_aggregate = reduction_op(thread_aggregate, item);347 thread_offset += BLOCK_THREADS;348 }349 }350 351 //---------------------------------------------------------------352 // Consume a contiguous segment of tiles353 //---------------------------------------------------------------------354 355 /**356 * @brief Reduce a contiguous segment of input tiles357 * @param even_share GridEvenShare descriptor358 * @param can_vectorize Whether or not we can vectorize loads359 */360 template <int CAN_VECTORIZE>361 __device__ __forceinline__ AccumT362 ConsumeRange(GridEvenShare<OffsetT> &even_share,363 Int2Type<CAN_VECTORIZE> can_vectorize)364 {365 AccumT thread_aggregate{};366 367 if (even_share.block_end - even_share.block_offset < TILE_ITEMS)368 {369 // First tile isn't full (not all threads have valid items)370 int valid_items = even_share.block_end - even_share.block_offset;371 ConsumeTile<true>(thread_aggregate,372 even_share.block_offset,373 valid_items,374 Int2Type<false>(),375 can_vectorize);376 return BlockReduceT(temp_storage.reduce)377 .Reduce(thread_aggregate, reduction_op, valid_items);378 }379 380 // Extracting this into a function saves 8% of generated kernel size by allowing to reuse381 // the block reduction below. This also workaround hang in nvcc.382 ConsumeFullTileRange(thread_aggregate, even_share, can_vectorize);383 384 // Compute block-wide reduction (all threads have valid items)385 return BlockReduceT(temp_storage.reduce)386 .Reduce(thread_aggregate, reduction_op);387 }388 389 /**390 * @brief Reduce a contiguous segment of input tiles391 * @param[in] block_offset Threadblock begin offset (inclusive)392 * @param[in] block_end Threadblock end offset (exclusive)393 */394 __device__ __forceinline__ AccumT ConsumeRange(OffsetT block_offset,395 OffsetT block_end)396 {397 GridEvenShare<OffsetT> even_share;398 even_share.template BlockInit<TILE_ITEMS>(block_offset, block_end);399 400 return (IsAligned(d_in + block_offset, Int2Type<ATTEMPT_VECTORIZATION>()))401 ? ConsumeRange(even_share,402 Int2Type < true && ATTEMPT_VECTORIZATION > ())403 : ConsumeRange(even_share,404 Int2Type < false && ATTEMPT_VECTORIZATION > ());405 }406 407 /**408 * Reduce a contiguous segment of input tiles409 * @param[in] even_share GridEvenShare descriptor410 */411 __device__ __forceinline__ AccumT412 ConsumeTiles(GridEvenShare<OffsetT> &even_share)413 {414 // Initialize GRID_MAPPING_STRIP_MINE even-share descriptor for this thread block415 even_share.template BlockInit<TILE_ITEMS, GRID_MAPPING_STRIP_MINE>();416 417 return (IsAligned(d_in, Int2Type<ATTEMPT_VECTORIZATION>()))418 ? ConsumeRange(even_share,419 Int2Type < true && ATTEMPT_VECTORIZATION > ())420 : ConsumeRange(even_share,421 Int2Type < false && ATTEMPT_VECTORIZATION > ());422 }423 424private:425 /**426 * @brief Reduce a contiguous segment of input tiles with more than `TILE_ITEMS` elements427 * @param even_share GridEvenShare descriptor428 * @param can_vectorize Whether or not we can vectorize loads429 */430 template <int CAN_VECTORIZE>431 __device__ __forceinline__ void432 ConsumeFullTileRange(AccumT &thread_aggregate,433 GridEvenShare<OffsetT> &even_share,434 Int2Type<CAN_VECTORIZE> can_vectorize)435 {436 // At least one full block437 ConsumeTile<true>(thread_aggregate,438 even_share.block_offset,439 TILE_ITEMS,440 Int2Type<true>(),441 can_vectorize);442 443 if (even_share.block_end - even_share.block_offset < even_share.block_stride)444 {445 // Exit early to handle offset overflow446 return;447 }448 449 even_share.block_offset += even_share.block_stride;450 451 // Consume subsequent full tiles of input, at least one full tile was processed, so452 // `even_share.block_end >= TILE_ITEMS`453 while (even_share.block_offset <= even_share.block_end - TILE_ITEMS)454 {455 ConsumeTile<false>(thread_aggregate,456 even_share.block_offset,457 TILE_ITEMS,458 Int2Type<true>(),459 can_vectorize);460 461 if (even_share.block_end - even_share.block_offset < even_share.block_stride)462 {463 // Exit early to handle offset overflow464 return;465 }466 467 even_share.block_offset += even_share.block_stride;468 }469 470 // Consume a partially-full tile471 if (even_share.block_offset < even_share.block_end)472 {473 int valid_items = even_share.block_end - even_share.block_offset;474 ConsumeTile<false>(thread_aggregate,475 even_share.block_offset,476 valid_items,477 Int2Type<false>(),478 can_vectorize);479 }480 }481};482 483CUB_NAMESPACE_END484 485