/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-hipcub/checkouts/docs-5.5.1/hipcub/include/hipcub/backend/rocprim/block/block_merge_sort.hpp Source File

• home
• docs
• checkouts
• readthedocs.org
• user_builds
• advanced-micro-devices-hipcub
• checkouts
• docs-5.5.1
• hipcub
• include
• hipcub
• backend
• rocprim
• block

block_merge_sort.hpp

1 /******************************************************************************

2 * Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.

3 * Modifications Copyright (c) 2021, Advanced Micro Devices, Inc. All rights reserved.

4 *

5 * Redistribution and use in source and binary forms, with or without

6 * modification, are permitted provided that the following conditions are met:

7 * * Redistributions of source code must retain the above copyright

8 * notice, this list of conditions and the following disclaimer.

9 * * Redistributions in binary form must reproduce the above copyright

10 * notice, this list of conditions and the following disclaimer in the

11 * documentation and/or other materials provided with the distribution.

12 * * Neither the name of the NVIDIA CORPORATION nor the

13 * names of its contributors may be used to endorse or promote products

14 * derived from this software without specific prior written permission.

15 *

16 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND

17 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED

18 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE

19 * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY

20 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES

21 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;

22 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND

23 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT

24 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS

25 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

26 *

27 ******************************************************************************/

28 

29 #ifndef HIPCUB_ROCPRIM_BLOCK_BLOCK_MERGE_SORT_HPP_

30 #define HIPCUB_ROCPRIM_BLOCK_BLOCK_MERGE_SORT_HPP_

31 

32 #include "../thread/thread_sort.hpp"

33 #include "../util_math.hpp"

34 #include "../util_type.hpp"

35 

36 #include <rocprim/detail/various.hpp>

37 #include <rocprim/functional.hpp>

38 

39 BEGIN_HIPCUB_NAMESPACE

40 

41 

42 // Additional details of the Merge-Path Algorithm can be found in:

43 // S. Odeh, O. Green, Z. Mwassi, O. Shmueli, Y. Birk, " Merge Path - Parallel

44 // Merging Made Simple", Multithreaded Architectures and Applications (MTAAP)

45 // Workshop, IEEE 26th International Parallel & Distributed Processing

46 // Symposium (IPDPS), 2012

47 template <typename KeyT,

48  typename KeyIteratorT,

49  typename OffsetT,

50  typename BinaryPred>

51 HIPCUB_DEVICE __forceinline__ OffsetT MergePath(KeyIteratorT keys1,

52  KeyIteratorT keys2,

53  OffsetT keys1_count,

54  OffsetT keys2_count,

55  OffsetT diag,

56  BinaryPred binary_pred)

57 {

58  OffsetT keys1_begin = diag < keys2_count ? 0 : diag - keys2_count;

59  OffsetT keys1_end = (::rocprim::min)(diag, keys1_count);

60 

61  while (keys1_begin < keys1_end)

62  {

63  OffsetT mid = hipcub::MidPoint<OffsetT>(keys1_begin, keys1_end);

64  KeyT key1 = keys1[mid];

65  KeyT key2 = keys2[diag - 1 - mid];

66  bool pred = binary_pred(key2, key1);

67 

68  if (pred)

69  {

70  keys1_end = mid;

71  }

72  else

73  {

74  keys1_begin = mid + 1;

75  }

76  }

77  return keys1_begin;

78 }

79 

80 template <typename KeyT, typename CompareOp, int ITEMS_PER_THREAD>

81 HIPCUB_DEVICE __forceinline__ void SerialMerge(KeyT *keys_shared,

82  int keys1_beg,

83  int keys2_beg,

84  int keys1_count,

85  int keys2_count,

86  KeyT (&output)[ITEMS_PER_THREAD],

87  int (&indices)[ITEMS_PER_THREAD],

88  CompareOp compare_op)

89 {

90  int keys1_end = keys1_beg + keys1_count;

91  int keys2_end = keys2_beg + keys2_count;

92 

93  KeyT key1 = keys_shared[keys1_beg];

94  KeyT key2 = keys_shared[keys2_beg];

95 

96 #pragma unroll

97  for (int item = 0; item < ITEMS_PER_THREAD; ++item)

98  {

99  bool p = (keys2_beg < keys2_end) &&

100  ((keys1_beg >= keys1_end)

101  || compare_op(key2, key1));

102 

103  output[item] = p ? key2 : key1;

104  indices[item] = p ? keys2_beg++ : keys1_beg++;

105 

106  if (p)

107  {

108  key2 = keys_shared[keys2_beg];

109  }

110  else

111  {

112  key1 = keys_shared[keys1_beg];

113  }

114  }

115 }

116 

169 template <typename KeyT,

170  typename ValueT,

171  int NUM_THREADS,

172  int ITEMS_PER_THREAD,

173  typename SynchronizationPolicy>

174 class BlockMergeSortStrategy

175 {

176  static_assert(PowerOfTwo<NUM_THREADS>::VALUE,

177  "NUM_THREADS must be a power of two");

178 

179 private:

180 

181  static constexpr int ITEMS_PER_TILE = ITEMS_PER_THREAD * NUM_THREADS;

182 

183  // Whether or not there are values to be trucked along with keys

184  static constexpr bool KEYS_ONLY = ::rocprim::Equals<ValueT, NullType>::VALUE;

185 

187  union _TempStorage

188  {

189  KeyT keys_shared[ITEMS_PER_TILE + 1];

190  ValueT items_shared[ITEMS_PER_TILE + 1];

191  }; // union TempStorage

192 

194  _TempStorage &temp_storage;

195 

197  HIPCUB_DEVICE __forceinline__ _TempStorage& PrivateStorage()

198  {

199  __shared__ _TempStorage private_storage;

200  return private_storage;

201  }

202 

203  const unsigned int linear_tid;

204 

205 public:

207  struct TempStorage : Uninitialized<_TempStorage> {};

208 

209  BlockMergeSortStrategy() = delete;

210  explicit HIPCUB_DEVICE __forceinline__

211  BlockMergeSortStrategy(unsigned int linear_tid)

212  : temp_storage(PrivateStorage())

213  , linear_tid(linear_tid)

214  {}

215 

216  HIPCUB_DEVICE __forceinline__ BlockMergeSortStrategy(TempStorage &temp_storage,

217  unsigned int linear_tid)

218  : temp_storage(temp_storage.Alias())

219  , linear_tid(linear_tid)

220  {}

221 

222  HIPCUB_DEVICE __forceinline__ unsigned int get_linear_tid() const

223  {

224  return linear_tid;

225  }

226 

249  template <typename CompareOp>

250  HIPCUB_DEVICE __forceinline__ void Sort(KeyT (&keys)[ITEMS_PER_THREAD],

251  CompareOp compare_op)

252  {

253  ValueT items[ITEMS_PER_THREAD];

254  Sort<CompareOp, false>(keys, items, compare_op, ITEMS_PER_TILE, keys[0]);

255  }

256 

291  template <typename CompareOp>

292  HIPCUB_DEVICE __forceinline__ void Sort(KeyT (&keys)[ITEMS_PER_THREAD],

293  CompareOp compare_op,

294  int valid_items,

295  KeyT oob_default)

296  {

297  ValueT items[ITEMS_PER_THREAD];

298  Sort<CompareOp, true>(keys, items, compare_op, valid_items, oob_default);

299  }

300 

325  template <typename CompareOp>

326  HIPCUB_DEVICE __forceinline__ void Sort(KeyT (&keys)[ITEMS_PER_THREAD],

327  ValueT (&items)[ITEMS_PER_THREAD],

328  CompareOp compare_op)

329  {

330  Sort<CompareOp, false>(keys, items, compare_op, ITEMS_PER_TILE, keys[0]);

331  }

332 

373  template <typename CompareOp,

374  bool IS_LAST_TILE = true>

375  HIPCUB_DEVICE __forceinline__ void Sort(KeyT (&keys)[ITEMS_PER_THREAD],

376  ValueT (&items)[ITEMS_PER_THREAD],

377  CompareOp compare_op,

378  int valid_items,

379  KeyT oob_default)

380  {

381  if (IS_LAST_TILE)

382  {

383  // if last tile, find valid max_key

384  // and fill the remaining keys with it

385  //

386  KeyT max_key = oob_default;

387 

388  #pragma unroll

389  for (int item = 1; item < ITEMS_PER_THREAD; ++item)

390  {

391  if (ITEMS_PER_THREAD * static_cast<int>(linear_tid) + item < valid_items)

392  {

393  max_key = compare_op(max_key, keys[item]) ? keys[item] : max_key;

394  }

395  else

396  {

397  keys[item] = max_key;

398  }

399  }

400  }

401 

402  // if first element of thread is in input range, stable sort items

403  //

404  if (!IS_LAST_TILE || ITEMS_PER_THREAD * static_cast<int>(linear_tid) < valid_items)

405  {

406  StableOddEvenSort(keys, items, compare_op);

407  }

408 

409  // each thread has sorted keys

410  // merge sort keys in shared memory

411  //

412  #pragma unroll

413  for (int target_merged_threads_number = 2;

414  target_merged_threads_number <= NUM_THREADS;

415  target_merged_threads_number *= 2)

416  {

417  int merged_threads_number = target_merged_threads_number / 2;

418  int mask = target_merged_threads_number - 1;

419 

420  Sync();

421 

422  // store keys in shmem

423  //

424  #pragma unroll

425  for (int item = 0; item < ITEMS_PER_THREAD; ++item)

426  {

427  int idx = ITEMS_PER_THREAD * linear_tid + item;

428  temp_storage.keys_shared[idx] = keys[item];

429  }

430 

431  Sync();

432 

433  int indices[ITEMS_PER_THREAD];

434 

435  int first_thread_idx_in_thread_group_being_merged = ~mask & linear_tid;

436  int start = ITEMS_PER_THREAD * first_thread_idx_in_thread_group_being_merged;

437  int size = ITEMS_PER_THREAD * merged_threads_number;

438 

439  int thread_idx_in_thread_group_being_merged = mask & linear_tid;

440 

441  int diag =

442  (::rocprim::min)(valid_items,

443  ITEMS_PER_THREAD * thread_idx_in_thread_group_being_merged);

444 

445  int keys1_beg = (::rocprim::min)(valid_items, start);

446  int keys1_end = (::rocprim::min)(valid_items, keys1_beg + size);

447  int keys2_beg = keys1_end;

448  int keys2_end = (::rocprim::min)(valid_items, keys2_beg + size);

449 

450  int keys1_count = keys1_end - keys1_beg;

451  int keys2_count = keys2_end - keys2_beg;

452 

453  int partition_diag = MergePath<KeyT>(&temp_storage.keys_shared[keys1_beg],

454  &temp_storage.keys_shared[keys2_beg],

455  keys1_count,

456  keys2_count,

457  diag,

458  compare_op);

459 

460  int keys1_beg_loc = keys1_beg + partition_diag;

461  int keys1_end_loc = keys1_end;

462  int keys2_beg_loc = keys2_beg + diag - partition_diag;

463  int keys2_end_loc = keys2_end;

464  int keys1_count_loc = keys1_end_loc - keys1_beg_loc;

465  int keys2_count_loc = keys2_end_loc - keys2_beg_loc;

466  SerialMerge(&temp_storage.keys_shared[0],

467  keys1_beg_loc,

468  keys2_beg_loc,

469  keys1_count_loc,

470  keys2_count_loc,

471  keys,

472  indices,

473  compare_op);

474 

475  if (!KEYS_ONLY)

476  {

477  Sync();

478 

479  // store keys in shmem

480  //

481  #pragma unroll

482  for (int item = 0; item < ITEMS_PER_THREAD; ++item)

483  {

484  int idx = ITEMS_PER_THREAD * linear_tid + item;

485  temp_storage.items_shared[idx] = items[item];

486  }

487 

488  Sync();

489 

490  // gather items from shmem

491  //

492  #pragma unroll

493  for (int item = 0; item < ITEMS_PER_THREAD; ++item)

494  {

495  items[item] = temp_storage.items_shared[indices[item]];

496  }

497  }

498  }

499  } // func block_merge_sort

500 

524  template <typename CompareOp>

525  HIPCUB_DEVICE __forceinline__ void StableSort(KeyT (&keys)[ITEMS_PER_THREAD],

526  CompareOp compare_op)

527  {

528  Sort(keys, compare_op);

529  }

530 

557  template <typename CompareOp>

558  HIPCUB_DEVICE __forceinline__ void StableSort(KeyT (&keys)[ITEMS_PER_THREAD],

559  ValueT (&items)[ITEMS_PER_THREAD],

560  CompareOp compare_op)

561  {

562  Sort(keys, items, compare_op);

563  }

564 

601  template <typename CompareOp>

602  HIPCUB_DEVICE __forceinline__ void StableSort(KeyT (&keys)[ITEMS_PER_THREAD],

603  CompareOp compare_op,

604  int valid_items,

605  KeyT oob_default)

606  {

607  Sort(keys, compare_op, valid_items, oob_default);

608  }

609 

651  template <typename CompareOp,

652  bool IS_LAST_TILE = true>

653  HIPCUB_DEVICE __forceinline__ void StableSort(KeyT (&keys)[ITEMS_PER_THREAD],

654  ValueT (&items)[ITEMS_PER_THREAD],

655  CompareOp compare_op,

656  int valid_items,

657  KeyT oob_default)

658  {

659  Sort<CompareOp, IS_LAST_TILE>(keys,

660  items,

661  compare_op,

662  valid_items,

663  oob_default);

664  }

665 

666 private:

667  HIPCUB_DEVICE __forceinline__ void Sync() const

668  {

669  static_cast<const SynchronizationPolicy*>(this)->SyncImplementation();

670  }

671 };

672 

673 

754 template <typename KeyT,

755  int BLOCK_DIM_X,

756  int ITEMS_PER_THREAD,

757  typename ValueT = NullType,

758  int BLOCK_DIM_Y = 1,

759  int BLOCK_DIM_Z = 1>

760 class BlockMergeSort

761  : public BlockMergeSortStrategy<KeyT,

762  ValueT,

763  BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z,

764  ITEMS_PER_THREAD,

765  BlockMergeSort<KeyT,

766  BLOCK_DIM_X,

767  ITEMS_PER_THREAD,

768  ValueT,

769  BLOCK_DIM_Y,

770  BLOCK_DIM_Z>>

771 {

772 private:

773  // The thread block size in threads

774  static constexpr int BLOCK_THREADS = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z;

775  static constexpr int ITEMS_PER_TILE = ITEMS_PER_THREAD * BLOCK_THREADS;

776 

777  using BlockMergeSortStrategyT =

778  BlockMergeSortStrategy<KeyT,

779  ValueT,

780  BLOCK_THREADS,

781  ITEMS_PER_THREAD,

782  BlockMergeSort>;

783 

784 public:

785  HIPCUB_DEVICE __forceinline__ BlockMergeSort()

786  : BlockMergeSortStrategyT(

787  RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))

788  {}

789 

790  HIPCUB_DEVICE __forceinline__ explicit BlockMergeSort(

791  typename BlockMergeSortStrategyT::TempStorage &temp_storage)

792  : BlockMergeSortStrategyT(

793  temp_storage,

794  RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))

795  {}

796 

797 private:

798  HIPCUB_DEVICE __forceinline__ void SyncImplementation() const

799  {

800  CTA_SYNC();

801  }

802 

803  friend BlockMergeSortStrategyT;

804 };

805 

806 END_HIPCUB_NAMESPACE

807 

808 #endif // HIPCUB_ROCPRIM_BLOCK_BLOCK_MERGE_SORT_HPP_

hipcub::BlockMergeSortStrategy

Generalized merge sort algorithm.

Definition: block_merge_sort.hpp:175

hipcub::BlockMergeSortStrategy::StableSort

__device__ __forceinline__ void StableSort(KeyT(&keys)[ITEMS_PER_THREAD], ValueT(&items)[ITEMS_PER_THREAD], CompareOp compare_op)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:558

hipcub::BlockMergeSortStrategy::StableSort

__device__ __forceinline__ void StableSort(KeyT(&keys)[ITEMS_PER_THREAD], CompareOp compare_op, int valid_items, KeyT oob_default)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:602

hipcub::BlockMergeSortStrategy::StableSort

__device__ __forceinline__ void StableSort(KeyT(&keys)[ITEMS_PER_THREAD], CompareOp compare_op)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:525

hipcub::BlockMergeSortStrategy::Sort

__device__ __forceinline__ void Sort(KeyT(&keys)[ITEMS_PER_THREAD], CompareOp compare_op)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:250

hipcub::BlockMergeSortStrategy::Sort

__device__ __forceinline__ void Sort(KeyT(&keys)[ITEMS_PER_THREAD], ValueT(&items)[ITEMS_PER_THREAD], CompareOp compare_op, int valid_items, KeyT oob_default)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:375

hipcub::BlockMergeSortStrategy::Sort

__device__ __forceinline__ void Sort(KeyT(&keys)[ITEMS_PER_THREAD], ValueT(&items)[ITEMS_PER_THREAD], CompareOp compare_op)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:326

hipcub::BlockMergeSortStrategy::StableSort

__device__ __forceinline__ void StableSort(KeyT(&keys)[ITEMS_PER_THREAD], ValueT(&items)[ITEMS_PER_THREAD], CompareOp compare_op, int valid_items, KeyT oob_default)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:653

hipcub::BlockMergeSortStrategy::Sort

__device__ __forceinline__ void Sort(KeyT(&keys)[ITEMS_PER_THREAD], CompareOp compare_op, int valid_items, KeyT oob_default)

Sorts items partitioned across a CUDA thread block using a merge sorting method.

Definition: block_merge_sort.hpp:292

hipcub::BlockMergeSort

The BlockMergeSort class provides methods for sorting items partitioned across a CUDA thread block us...

Definition: block_merge_sort.hpp:771

hipcub::BlockMergeSortStrategy::TempStorage

\smemstorage{BlockMergeSort}

Definition: block_merge_sort.hpp:207

hipcub::PowerOfTwo

Definition: util_type.hpp:78

hipcub::Uninitialized

A storage-backing wrapper that allows types with non-trivial constructors to be aliased in unions.

Definition: util_type.hpp:363