| /****************************************************************************** |
| * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. |
| * |
| * Redistribution and use in source and binary forms, with or without |
| * modification, are permitted provided that the following conditions are met: |
| * * Redistributions of source code must retain the above copyright |
| * notice, this list of conditions and the following disclaimer. |
| * * Redistributions in binary form must reproduce the above copyright |
| * notice, this list of conditions and the following disclaimer in the |
| * documentation and/or other materials provided with the distribution. |
| * * Neither the name of the NVIDIA CORPORATION nor the |
| * names of its contributors may be used to endorse or promote products |
| * derived from this software without specific prior written permission. |
| * |
| * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND |
| * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED |
| * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE |
| * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY |
| * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES |
| * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; |
| * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND |
| * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS |
| * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| * |
| ******************************************************************************/ |
| |
| #pragma once |
| |
| #include "../config.cuh" |
| #include "../util_type.cuh" |
| #include "../util_namespace.cuh" |
| #include "../block/block_load.cuh" |
| #include "../block/block_store.cuh" |
| #include "../block/block_merge_sort.cuh" |
| |
| #include <thrust/system/cuda/detail/core/util.h> |
| |
| CUB_NAMESPACE_BEGIN |
| |
| |
| template < |
| int _BLOCK_THREADS, |
| int _ITEMS_PER_THREAD = 1, |
| cub::BlockLoadAlgorithm _LOAD_ALGORITHM = cub::BLOCK_LOAD_DIRECT, |
| cub::CacheLoadModifier _LOAD_MODIFIER = cub::LOAD_LDG, |
| cub::BlockStoreAlgorithm _STORE_ALGORITHM = cub::BLOCK_STORE_DIRECT> |
| struct AgentMergeSortPolicy |
| { |
| static constexpr int BLOCK_THREADS = _BLOCK_THREADS; |
| static constexpr int ITEMS_PER_THREAD = _ITEMS_PER_THREAD; |
| static constexpr int ITEMS_PER_TILE = BLOCK_THREADS * ITEMS_PER_THREAD; |
| |
| static constexpr cub::BlockLoadAlgorithm LOAD_ALGORITHM = _LOAD_ALGORITHM; |
| static constexpr cub::CacheLoadModifier LOAD_MODIFIER = _LOAD_MODIFIER; |
| static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = _STORE_ALGORITHM; |
| }; |
| |
| /// \brief This agent is responsible for the initial in-tile sorting. |
| template <typename Policy, |
| typename KeyInputIteratorT, |
| typename ValueInputIteratorT, |
| typename KeyIteratorT, |
| typename ValueIteratorT, |
| typename OffsetT, |
| typename CompareOpT, |
| typename KeyT, |
| typename ValueT> |
| struct AgentBlockSort |
| { |
| //--------------------------------------------------------------------- |
| // Types and constants |
| //--------------------------------------------------------------------- |
| |
| static constexpr bool KEYS_ONLY = std::is_same<ValueT, NullType>::value; |
| |
| using BlockMergeSortT = |
| BlockMergeSort<KeyT, Policy::BLOCK_THREADS, Policy::ITEMS_PER_THREAD, ValueT>; |
| |
| using KeysLoadIt = typename THRUST_NS_QUALIFIER::cuda_cub::core::LoadIterator<Policy, KeyInputIteratorT>::type; |
| using ItemsLoadIt = typename THRUST_NS_QUALIFIER::cuda_cub::core::LoadIterator<Policy, ValueInputIteratorT>::type; |
| |
| using BlockLoadKeys = typename cub::BlockLoadType<Policy, KeysLoadIt>::type; |
| using BlockLoadItems = typename cub::BlockLoadType<Policy, ItemsLoadIt>::type; |
| |
| using BlockStoreKeysIt = typename cub::BlockStoreType<Policy, KeyIteratorT>::type; |
| using BlockStoreItemsIt = typename cub::BlockStoreType<Policy, ValueIteratorT>::type; |
| using BlockStoreKeysRaw = typename cub::BlockStoreType<Policy, KeyT *>::type; |
| using BlockStoreItemsRaw = typename cub::BlockStoreType<Policy, ValueT *>::type; |
| |
| union _TempStorage |
| { |
| typename BlockLoadKeys::TempStorage load_keys; |
| typename BlockLoadItems::TempStorage load_items; |
| typename BlockStoreKeysIt::TempStorage store_keys_it; |
| typename BlockStoreItemsIt::TempStorage store_items_it; |
| typename BlockStoreKeysRaw::TempStorage store_keys_raw; |
| typename BlockStoreItemsRaw::TempStorage store_items_raw; |
| typename BlockMergeSortT::TempStorage block_merge; |
| }; |
| |
| /// Alias wrapper allowing storage to be unioned |
| struct TempStorage : Uninitialized<_TempStorage> {}; |
| |
| static constexpr int BLOCK_THREADS = Policy::BLOCK_THREADS; |
| static constexpr int ITEMS_PER_THREAD = Policy::ITEMS_PER_THREAD; |
| static constexpr int ITEMS_PER_TILE = Policy::ITEMS_PER_TILE; |
| static constexpr int SHARED_MEMORY_SIZE = |
| static_cast<int>(sizeof(TempStorage)); |
| |
| //--------------------------------------------------------------------- |
| // Per thread data |
| //--------------------------------------------------------------------- |
| |
| bool ping; |
| _TempStorage &storage; |
| KeysLoadIt keys_in; |
| ItemsLoadIt items_in; |
| OffsetT keys_count; |
| KeyIteratorT keys_out_it; |
| ValueIteratorT items_out_it; |
| KeyT *keys_out_raw; |
| ValueT *items_out_raw; |
| CompareOpT compare_op; |
| |
| __device__ __forceinline__ AgentBlockSort(bool ping_, |
| TempStorage &storage_, |
| KeysLoadIt keys_in_, |
| ItemsLoadIt items_in_, |
| OffsetT keys_count_, |
| KeyIteratorT keys_out_it_, |
| ValueIteratorT items_out_it_, |
| KeyT *keys_out_raw_, |
| ValueT *items_out_raw_, |
| CompareOpT compare_op_) |
| : ping(ping_) |
| , storage(storage_.Alias()) |
| , keys_in(keys_in_) |
| , items_in(items_in_) |
| , keys_count(keys_count_) |
| , keys_out_it(keys_out_it_) |
| , items_out_it(items_out_it_) |
| , keys_out_raw(keys_out_raw_) |
| , items_out_raw(items_out_raw_) |
| , compare_op(compare_op_) |
| { |
| } |
| |
| __device__ __forceinline__ void Process() |
| { |
| auto tile_idx = static_cast<OffsetT>(blockIdx.x); |
| auto num_tiles = static_cast<OffsetT>(gridDim.x); |
| auto tile_base = tile_idx * ITEMS_PER_TILE; |
| int items_in_tile = (cub::min)(keys_count - tile_base, int{ITEMS_PER_TILE}); |
| |
| if (tile_idx < num_tiles - 1) |
| { |
| consume_tile<false>(tile_base, ITEMS_PER_TILE); |
| } |
| else |
| { |
| consume_tile<true>(tile_base, items_in_tile); |
| } |
| } |
| |
| template <bool IS_LAST_TILE> |
| __device__ __forceinline__ void consume_tile(OffsetT tile_base, |
| int num_remaining) |
| { |
| ValueT items_local[ITEMS_PER_THREAD]; |
| if (!KEYS_ONLY) |
| { |
| if (IS_LAST_TILE) |
| { |
| BlockLoadItems(storage.load_items) |
| .Load(items_in + tile_base, |
| items_local, |
| num_remaining, |
| *(items_in + tile_base)); |
| } |
| else |
| { |
| BlockLoadItems(storage.load_items).Load(items_in + tile_base, items_local); |
| } |
| |
| CTA_SYNC(); |
| } |
| |
| KeyT keys_local[ITEMS_PER_THREAD]; |
| if (IS_LAST_TILE) |
| { |
| BlockLoadKeys(storage.load_keys) |
| .Load(keys_in + tile_base, |
| keys_local, |
| num_remaining, |
| *(keys_in + tile_base)); |
| } |
| else |
| { |
| BlockLoadKeys(storage.load_keys) |
| .Load(keys_in + tile_base, keys_local); |
| } |
| |
| CTA_SYNC(); |
| |
| if (IS_LAST_TILE) |
| { |
| BlockMergeSortT(storage.block_merge) |
| .Sort(keys_local, items_local, compare_op, num_remaining, keys_local[0]); |
| } |
| else |
| { |
| BlockMergeSortT(storage.block_merge).Sort(keys_local, items_local, compare_op); |
| } |
| |
| CTA_SYNC(); |
| |
| if (ping) |
| { |
| if (IS_LAST_TILE) |
| { |
| BlockStoreKeysIt(storage.store_keys_it) |
| .Store(keys_out_it + tile_base, keys_local, num_remaining); |
| } |
| else |
| { |
| BlockStoreKeysIt(storage.store_keys_it) |
| .Store(keys_out_it + tile_base, keys_local); |
| } |
| |
| if (!KEYS_ONLY) |
| { |
| CTA_SYNC(); |
| |
| if (IS_LAST_TILE) |
| { |
| BlockStoreItemsIt(storage.store_items_it) |
| .Store(items_out_it + tile_base, items_local, num_remaining); |
| } |
| else |
| { |
| BlockStoreItemsIt(storage.store_items_it) |
| .Store(items_out_it + tile_base, items_local); |
| } |
| } |
| } |
| else |
| { |
| if (IS_LAST_TILE) |
| { |
| BlockStoreKeysRaw(storage.store_keys_raw) |
| .Store(keys_out_raw + tile_base, keys_local, num_remaining); |
| } |
| else |
| { |
| BlockStoreKeysRaw(storage.store_keys_raw) |
| .Store(keys_out_raw + tile_base, keys_local); |
| } |
| |
| if (!KEYS_ONLY) |
| { |
| CTA_SYNC(); |
| |
| if (IS_LAST_TILE) |
| { |
| BlockStoreItemsRaw(storage.store_items_raw) |
| .Store(items_out_raw + tile_base, items_local, num_remaining); |
| } |
| else |
| { |
| BlockStoreItemsRaw(storage.store_items_raw) |
| .Store(items_out_raw + tile_base, items_local); |
| } |
| } |
| } |
| } |
| }; |
| |
| /** |
| * \brief This agent is responsible for partitioning a merge path into equal segments |
| * |
| * There are two sorted arrays to be merged into one array. If the first array |
| * is partitioned between parallel workers by slicing it into ranges of equal |
| * size, there could be a significant workload imbalance. The imbalance is |
| * caused by the fact that the distribution of elements from the second array |
| * is unknown beforehand. Instead, the MergePath is partitioned between workers. |
| * This approach guarantees an equal amount of work being assigned to each worker. |
| * |
| * This approach is outlined in the paper: |
| * Odeh et al, "Merge Path - Parallel Merging Made Simple" |
| * doi:10.1109/IPDPSW.2012.202 |
| */ |
| template < |
| typename KeyIteratorT, |
| typename OffsetT, |
| typename CompareOpT, |
| typename KeyT> |
| struct AgentPartition |
| { |
| bool ping; |
| KeyIteratorT keys_ping; |
| KeyT *keys_pong; |
| OffsetT keys_count; |
| OffsetT partition_idx; |
| OffsetT *merge_partitions; |
| CompareOpT compare_op; |
| OffsetT target_merged_tiles_number; |
| int items_per_tile; |
|
|
| __device__ __forceinline__ AgentPartition(bool ping, |
| KeyIteratorT keys_ping, |
| KeyT *keys_pong, |
| OffsetT keys_count, |
| OffsetT partition_idx, |
| OffsetT *merge_partitions, |
| CompareOpT compare_op, |
| OffsetT target_merged_tiles_number, |
| int items_per_tile) |
| : ping(ping) |
| , keys_ping(keys_ping) |
| , keys_pong(keys_pong) |
| , keys_count(keys_count) |
| , partition_idx(partition_idx) |
| , merge_partitions(merge_partitions) |
| , compare_op(compare_op) |
| , target_merged_tiles_number(target_merged_tiles_number) |
| , items_per_tile(items_per_tile) |
| {} |
| |
| __device__ __forceinline__ void Process() |
| { |
| OffsetT merged_tiles_number = target_merged_tiles_number / 2; |
| |
| // target_merged_tiles_number is a power of two. |
| OffsetT mask = target_merged_tiles_number - 1; |
| |
| // The first tile number in the tiles group being merged, equal to: |
| // target_merged_tiles_number * (partition_idx / target_merged_tiles_number) |
| OffsetT list = ~mask & partition_idx; |
| OffsetT start = items_per_tile * list; |
| OffsetT size = items_per_tile * merged_tiles_number; |
| |
| // Tile number within the tile group being merged, equal to: |
| // partition_idx / target_merged_tiles_number |
| OffsetT local_tile_idx = mask & partition_idx; |
| |
| OffsetT keys1_beg = (cub::min)(keys_count, start); |
| OffsetT keys1_end = (cub::min)(keys_count, start + size); |
| OffsetT keys2_beg = keys1_end; |
| OffsetT keys2_end = (cub::min)(keys_count, keys2_beg + size); |
| |
| OffsetT partition_at = (cub::min)(keys2_end - keys1_beg, |
| items_per_tile * local_tile_idx); |
| |
| OffsetT partition_diag = ping ? MergePath<KeyT>(keys_ping + keys1_beg, |
| keys_ping + keys2_beg, |
| keys1_end - keys1_beg, |
| keys2_end - keys2_beg, |
| partition_at, |
| compare_op) |
| : MergePath<KeyT>(keys_pong + keys1_beg, |
| keys_pong + keys2_beg, |
| keys1_end - keys1_beg, |
| keys2_end - keys2_beg, |
| partition_at, |
| compare_op); |
| |
| merge_partitions[partition_idx] = keys1_beg + partition_diag; |
| } |
| }; |
| |
| /// \brief The agent is responsible for merging N consecutive sorted arrays into N/2 sorted arrays. |
| template < |
| typename Policy, |
| typename KeyIteratorT, |
| typename ValueIteratorT, |
| typename OffsetT, |
| typename CompareOpT, |
| typename KeyT, |
| typename ValueT> |
| struct AgentMerge |
| { |
|
|
| //--------------------------------------------------------------------- |
| // Types and constants |
| //--------------------------------------------------------------------- |
| using KeysLoadPingIt = typename THRUST_NS_QUALIFIER::cuda_cub::core::LoadIterator<Policy, KeyIteratorT>::type; |
| using ItemsLoadPingIt = typename THRUST_NS_QUALIFIER::cuda_cub::core::LoadIterator<Policy, ValueIteratorT>::type; |
| using KeysLoadPongIt = typename THRUST_NS_QUALIFIER::cuda_cub::core::LoadIterator<Policy, KeyT *>::type; |
| using ItemsLoadPongIt = typename THRUST_NS_QUALIFIER::cuda_cub::core::LoadIterator<Policy, ValueT *>::type; |
|
|
| using KeysOutputPongIt = KeyIteratorT; |
| using ItemsOutputPongIt = ValueIteratorT; |
| using KeysOutputPingIt = KeyT*; |
| using ItemsOutputPingIt = ValueT*; |
|
|
| using BlockStoreKeysPong = typename BlockStoreType<Policy, KeysOutputPongIt>::type; |
| using BlockStoreItemsPong = typename BlockStoreType<Policy, ItemsOutputPongIt>::type; |
| using BlockStoreKeysPing = typename BlockStoreType<Policy, KeysOutputPingIt>::type; |
| using BlockStoreItemsPing = typename BlockStoreType<Policy, ItemsOutputPingIt>::type; |
|
|
| /// Parameterized BlockReduce primitive |
|
|
| union _TempStorage |
| { |
| typename BlockStoreKeysPing::TempStorage store_keys_ping; |
| typename BlockStoreItemsPing::TempStorage store_items_ping; |
| typename BlockStoreKeysPong::TempStorage store_keys_pong; |
| typename BlockStoreItemsPong::TempStorage store_items_pong; |
| |
| KeyT keys_shared[Policy::ITEMS_PER_TILE + 1]; |
| ValueT items_shared[Policy::ITEMS_PER_TILE + 1]; |
| }; |
| |
| /// Alias wrapper allowing storage to be unioned |
| struct TempStorage : Uninitialized<_TempStorage> {}; |
|
|
| static constexpr bool KEYS_ONLY = std::is_same<ValueT, NullType>::value; |
| static constexpr int BLOCK_THREADS = Policy::BLOCK_THREADS; |
| static constexpr int ITEMS_PER_THREAD = Policy::ITEMS_PER_THREAD; |
| static constexpr int ITEMS_PER_TILE = Policy::ITEMS_PER_TILE; |
| static constexpr int SHARED_MEMORY_SIZE = |
| static_cast<int>(sizeof(TempStorage)); |
| |
| //--------------------------------------------------------------------- |
| // Per thread data |
| //--------------------------------------------------------------------- |
|
|
| bool ping; |
| _TempStorage& storage; |
| |
| KeysLoadPingIt keys_in_ping; |
| ItemsLoadPingIt items_in_ping; |
| KeysLoadPongIt keys_in_pong; |
| ItemsLoadPongIt items_in_pong; |
| |
| OffsetT keys_count; |
|
|
| KeysOutputPongIt keys_out_pong; |
| ItemsOutputPongIt items_out_pong; |
| KeysOutputPingIt keys_out_ping; |
| ItemsOutputPingIt items_out_ping; |
|
|
| CompareOpT compare_op; |
| OffsetT *merge_partitions; |
| OffsetT target_merged_tiles_number; |
| |
| //--------------------------------------------------------------------- |
| // Utility functions |
| //--------------------------------------------------------------------- |
| |
| /** |
| * \brief Concatenates up to ITEMS_PER_THREAD elements from input{1,2} into output array |
| * |
| * Reads data in a coalesced fashion [BLOCK_THREADS * item + tid] and |
| * stores the result in output[item]. |
| */ |
| template <bool IS_FULL_TILE, class T, class It1, class It2> |
| __device__ __forceinline__ void |
| gmem_to_reg(T (&output)[ITEMS_PER_THREAD], |
| It1 input1, |
| It2 input2, |
| int count1, |
| int count2) |
| { |
| if (IS_FULL_TILE) |
| { |
| #pragma unroll |
| for (int item = 0; item < ITEMS_PER_THREAD; ++item) |
| { |
| int idx = BLOCK_THREADS * item + threadIdx.x; |
| output[item] = (idx < count1) ? input1[idx] : input2[idx - count1]; |
| } |
| } |
| else |
| { |
| #pragma unroll |
| for (int item = 0; item < ITEMS_PER_THREAD; ++item) |
| { |
| int idx = BLOCK_THREADS * item + threadIdx.x; |
| if (idx < count1 + count2) |
| { |
| output[item] = (idx < count1) ? input1[idx] : input2[idx - count1]; |
| } |
| } |
| } |
| } |
| |
| /// \brief Stores data in a coalesced fashion in[item] -> out[BLOCK_THREADS * item + tid] |
| template <class T, class It> |
| __device__ __forceinline__ void |
| reg_to_shared(It output, |
| T (&input)[ITEMS_PER_THREAD]) |
| { |
| #pragma unroll |
| for (int item = 0; item < ITEMS_PER_THREAD; ++item) |
| { |
| int idx = BLOCK_THREADS * item + threadIdx.x; |
| output[idx] = input[item]; |
| } |
| } |
| |
| template <bool IS_FULL_TILE> |
| __device__ __forceinline__ void |
| consume_tile(int tid, OffsetT tile_idx, OffsetT tile_base, int count) |
| { |
| OffsetT partition_beg = merge_partitions[tile_idx + 0]; |
| OffsetT partition_end = merge_partitions[tile_idx + 1]; |
| |
| // target_merged_tiles_number is a power of two. |
| OffsetT merged_tiles_number = target_merged_tiles_number / 2; |
| |
| OffsetT mask = target_merged_tiles_number - 1; |
| |
| // The first tile number in the tiles group being merged, equal to: |
| // target_merged_tiles_number * (tile_idx / target_merged_tiles_number) |
| OffsetT list = ~mask & tile_idx; |
| OffsetT start = ITEMS_PER_TILE * list; |
| OffsetT size = ITEMS_PER_TILE * merged_tiles_number; |
| |
| OffsetT diag = ITEMS_PER_TILE * tile_idx - start; |
| |
| OffsetT keys1_beg = partition_beg; |
| OffsetT keys1_end = partition_end; |
| OffsetT keys2_beg = (cub::min)(keys_count, 2 * start + size + diag - partition_beg); |
| OffsetT keys2_end = (cub::min)(keys_count, 2 * start + size + diag + ITEMS_PER_TILE - partition_end); |
| |
| // Check if it's the last tile in the tile group being merged |
| if (mask == (mask & tile_idx)) |
| { |
| keys1_end = (cub::min)(keys_count, start + size); |
| keys2_end = (cub::min)(keys_count, start + size * 2); |
| } |
| |
| // number of keys per tile |
| // |
| int num_keys1 = static_cast<int>(keys1_end - keys1_beg); |
| int num_keys2 = static_cast<int>(keys2_end - keys2_beg); |
| |
| // load keys1 & keys2 |
| KeyT keys_local[ITEMS_PER_THREAD]; |
| if (ping) |
| { |
| gmem_to_reg<IS_FULL_TILE>(keys_local, |
| keys_in_ping + keys1_beg, |
| keys_in_ping + keys2_beg, |
| num_keys1, |
| num_keys2); |
| } |
| else |
| { |
| gmem_to_reg<IS_FULL_TILE>(keys_local, |
| keys_in_pong + keys1_beg, |
| keys_in_pong + keys2_beg, |
| num_keys1, |
| num_keys2); |
| } |
| reg_to_shared(&storage.keys_shared[0], keys_local); |
| |
| // preload items into registers already |
| // |
| ValueT items_local[ITEMS_PER_THREAD]; |
| if (!KEYS_ONLY) |
| { |
| if (ping) |
| { |
| gmem_to_reg<IS_FULL_TILE>(items_local, |
| items_in_ping + keys1_beg, |
| items_in_ping + keys2_beg, |
| num_keys1, |
| num_keys2); |
| } |
| else |
| { |
| gmem_to_reg<IS_FULL_TILE>(items_local, |
| items_in_pong + keys1_beg, |
| items_in_pong + keys2_beg, |
| num_keys1, |
| num_keys2); |
| } |
| } |
| |
| CTA_SYNC(); |
| |
| // use binary search in shared memory |
| // to find merge path for each of thread |
| // we can use int type here, because the number of |
| // items in shared memory is limited |
| // |
| int diag0_local = (cub::min)(num_keys1 + num_keys2, ITEMS_PER_THREAD * tid); |
| |
| int keys1_beg_local = MergePath<KeyT>(&storage.keys_shared[0], |
| &storage.keys_shared[num_keys1], |
| num_keys1, |
| num_keys2, |
| diag0_local, |
| compare_op); |
| int keys1_end_local = num_keys1; |
| int keys2_beg_local = diag0_local - keys1_beg_local; |
| int keys2_end_local = num_keys2; |
| |
| int num_keys1_local = keys1_end_local - keys1_beg_local; |
| int num_keys2_local = keys2_end_local - keys2_beg_local; |
| |
| // perform serial merge |
| // |
| int indices[ITEMS_PER_THREAD]; |
| |
| SerialMerge(&storage.keys_shared[0], |
| keys1_beg_local, |
| keys2_beg_local + num_keys1, |
| num_keys1_local, |
| num_keys2_local, |
| keys_local, |
| indices, |
| compare_op); |
| |
| CTA_SYNC(); |
| |
| // write keys |
| // |
| if (ping) |
| { |
| if (IS_FULL_TILE) |
| { |
| BlockStoreKeysPing(storage.store_keys_ping) |
| .Store(keys_out_ping + tile_base, keys_local); |
| } |
| else |
| { |
| BlockStoreKeysPing(storage.store_keys_ping) |
| .Store(keys_out_ping + tile_base, keys_local, num_keys1 + num_keys2); |
| } |
| } |
| else |
| { |
| if (IS_FULL_TILE) |
| { |
| BlockStoreKeysPong(storage.store_keys_pong) |
| .Store(keys_out_pong + tile_base, keys_local); |
| } |
| else |
| { |
| BlockStoreKeysPong(storage.store_keys_pong) |
| .Store(keys_out_pong + tile_base, keys_local, num_keys1 + num_keys2); |
| } |
| } |
| |
| // if items are provided, merge them |
| if (!KEYS_ONLY) |
| { |
| CTA_SYNC(); |
| |
| reg_to_shared(&storage.items_shared[0], items_local); |
| |
| CTA_SYNC(); |
| |
| // gather items from shared mem |
| // |
| #pragma unroll |
| for (int item = 0; item < ITEMS_PER_THREAD; ++item) |
| { |
| items_local[item] = storage.items_shared[indices[item]]; |
| } |
| |
| CTA_SYNC(); |
| |
| // write from reg to gmem |
| // |
| if (ping) |
| { |
| if (IS_FULL_TILE) |
| { |
| BlockStoreItemsPing(storage.store_items_ping) |
| .Store(items_out_ping + tile_base, items_local); |
| } |
| else |
| { |
| BlockStoreItemsPing(storage.store_items_ping) |
| .Store(items_out_ping + tile_base, items_local, count); |
| } |
| } |
| else |
| { |
| if (IS_FULL_TILE) |
| { |
| BlockStoreItemsPong(storage.store_items_pong) |
| .Store(items_out_pong + tile_base, items_local); |
| } |
| else |
| { |
| BlockStoreItemsPong(storage.store_items_pong) |
| .Store(items_out_pong + tile_base, items_local, count); |
| } |
| } |
| } |
| } |
| |
| __device__ __forceinline__ AgentMerge(bool ping_, |
| TempStorage &storage_, |
| KeysLoadPingIt keys_in_ping_, |
| ItemsLoadPingIt items_in_ping_, |
| KeysLoadPongIt keys_in_pong_, |
| ItemsLoadPongIt items_in_pong_, |
| OffsetT keys_count_, |
| KeysOutputPingIt keys_out_ping_, |
| ItemsOutputPingIt items_out_ping_, |
| KeysOutputPongIt keys_out_pong_, |
| ItemsOutputPongIt items_out_pong_, |
| CompareOpT compare_op_, |
| OffsetT *merge_partitions_, |
| OffsetT target_merged_tiles_number_) |
| : ping(ping_) |
| , storage(storage_.Alias()) |
| , keys_in_ping(keys_in_ping_) |
| , items_in_ping(items_in_ping_) |
| , keys_in_pong(keys_in_pong_) |
| , items_in_pong(items_in_pong_) |
| , keys_count(keys_count_) |
| , keys_out_pong(keys_out_pong_) |
| , items_out_pong(items_out_pong_) |
| , keys_out_ping(keys_out_ping_) |
| , items_out_ping(items_out_ping_) |
| , compare_op(compare_op_) |
| , merge_partitions(merge_partitions_) |
| , target_merged_tiles_number(target_merged_tiles_number_) |
| {} |
| |
| __device__ __forceinline__ void Process() |
| { |
| int tile_idx = static_cast<int>(blockIdx.x); |
| int num_tiles = static_cast<int>(gridDim.x); |
| OffsetT tile_base = OffsetT(tile_idx) * ITEMS_PER_TILE; |
| int tid = static_cast<int>(threadIdx.x); |
| int items_in_tile = static_cast<int>( |
| (cub::min)(static_cast<OffsetT>(ITEMS_PER_TILE), keys_count - tile_base)); |
| |
| if (tile_idx < num_tiles - 1) |
| { |
| consume_tile<true>(tid, tile_idx, tile_base, ITEMS_PER_TILE); |
| } |
| else |
| { |
| consume_tile<false>(tid, tile_idx, tile_base, items_in_tile); |
| } |
| } |
| }; |
| |
|
|
| CUB_NAMESPACE_END |
|
|
