MUDA
Loading...
Searching...
No Matches
device_segmented_sort.h
1#pragma once
2#include <muda/cub/device/cub_wrapper.h>
3#include "details/cub_wrapper_macro_def.inl"
4#ifndef __INTELLISENSE__
5#include <cub/device/device_segmented_sort.cuh>
6#endif
7
8namespace muda
9{
10//ref: https://nvlabs.github.io/cub/structcub_1_1_device_segmented_reduce.html
11class DeviceSegmentedSort : public CubWrapper<DeviceSegmentedSort>
12{
13 using Base = CubWrapper<DeviceSegmentedSort>;
14
15 public:
16 using Base::Base;
17
18 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
19 DeviceSegmentedSort& SortKeys(const KeyT* d_keys_in,
20 KeyT* d_keys_out,
21 int num_items,
22 int num_segments,
23 BeginOffsetIteratorT d_begin_offsets,
24 EndOffsetIteratorT d_end_offsets)
25 {
26 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::SortKeys(d_temp_storage,
27 temp_storage_bytes,
28 d_keys_in,
29 d_keys_out,
30 num_items,
31 num_segments,
32 d_begin_offsets,
33 d_end_offsets,
34 _stream,
35 false));
36 }
37
38 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
39 DeviceSegmentedSort& SortKeysDescending(const KeyT* d_keys_in,
40 KeyT* d_keys_out,
41 int num_items,
42 int num_segments,
43 BeginOffsetIteratorT d_begin_offsets,
44 EndOffsetIteratorT d_end_offsets)
45 {
46 MUDA_CUB_WRAPPER_IMPL(
47 cub::DeviceSegmentedSort::SortKeysDescending(d_temp_storage,
48 temp_storage_bytes,
49 d_keys_in,
50 d_keys_out,
51 num_items,
52 num_segments,
53 d_begin_offsets,
54 d_end_offsets,
55 _stream,
56 false));
57 }
58
59
60 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
61 DeviceSegmentedSort& SortKeys(cub::DoubleBuffer<KeyT>& d_keys,
62 int num_items,
63 int num_segments,
64 BeginOffsetIteratorT d_begin_offsets,
65 EndOffsetIteratorT d_end_offsets)
66 {
67 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::SortKeys(
68 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
69 }
70
71
72 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
73 DeviceSegmentedSort& SortKeysDescending(cub::DoubleBuffer<KeyT>& d_keys,
74 int num_items,
75 int num_segments,
76 BeginOffsetIteratorT d_begin_offsets,
77 EndOffsetIteratorT d_end_offsets)
78 {
79 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::SortKeysDescending(
80 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
81 }
82
83
84 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
85 DeviceSegmentedSort& StableSortKeys(const KeyT* d_keys_in,
86 KeyT* d_keys_out,
87 int num_items,
88 int num_segments,
89 BeginOffsetIteratorT d_begin_offsets,
90 EndOffsetIteratorT d_end_offsets)
91 {
92 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::StableSortKeys(d_temp_storage,
93 temp_storage_bytes,
94 d_keys_in,
95 d_keys_out,
96 num_items,
97 num_segments,
98 d_begin_offsets,
99 d_end_offsets,
100 _stream,
101 false));
102 }
103
104
105 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
106 DeviceSegmentedSort& StableSortKeysDescending(const KeyT* d_keys_in,
107 KeyT* d_keys_out,
108 int num_items,
109 int num_segments,
110 BeginOffsetIteratorT d_begin_offsets,
111 EndOffsetIteratorT d_end_offsets)
112 {
113 MUDA_CUB_WRAPPER_IMPL(
114 cub::DeviceSegmentedSort::StableSortKeysDescending(d_temp_storage,
115 temp_storage_bytes,
116 d_keys_in,
117 d_keys_out,
118 num_items,
119 num_segments,
120 d_begin_offsets,
121 d_end_offsets,
122 _stream,
123 false));
124 }
125
126
127 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
128 DeviceSegmentedSort& StableSortKeys(cub::DoubleBuffer<KeyT>& d_keys,
129 int num_items,
130 int num_segments,
131 BeginOffsetIteratorT d_begin_offsets,
132 EndOffsetIteratorT d_end_offsets)
133 {
134 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::StableSortKeys(
135 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
136 }
137
138
139 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
140 DeviceSegmentedSort& StableSortKeysDescending(cub::DoubleBuffer<KeyT>& d_keys,
141 int num_items,
142 int num_segments,
143 BeginOffsetIteratorT d_begin_offsets,
144 EndOffsetIteratorT d_end_offsets)
145 {
146 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::StableSortKeysDescending(
147 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
148 }
149
150
151 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
152 DeviceSegmentedSort& SortPairs(const KeyT* d_keys_in,
153 KeyT* d_keys_out,
154 const ValueT* d_values_in,
155 ValueT* d_values_out,
156 int num_items,
157 int num_segments,
158 BeginOffsetIteratorT d_begin_offsets,
159 EndOffsetIteratorT d_end_offsets)
160 {
161 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::SortPairs(d_temp_storage,
162 temp_storage_bytes,
163 d_keys_in,
164 d_keys_out,
165 d_values_in,
166 d_values_out,
167 num_items,
168 num_segments,
169 d_begin_offsets,
170 d_end_offsets,
171 _stream,
172 false));
173 }
174
175
176 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
177 DeviceSegmentedSort& SortPairsDescending(const KeyT* d_keys_in,
178 KeyT* d_keys_out,
179 const ValueT* d_values_in,
180 ValueT* d_values_out,
181 int num_items,
182 int num_segments,
183 BeginOffsetIteratorT d_begin_offsets,
184 EndOffsetIteratorT d_end_offsets)
185 {
186 MUDA_CUB_WRAPPER_IMPL(
187 cub::DeviceSegmentedSort::SortPairsDescending(d_temp_storage,
188 temp_storage_bytes,
189 d_keys_in,
190 d_keys_out,
191 d_values_in,
192 d_values_out,
193 num_items,
194 num_segments,
195 d_begin_offsets,
196 d_end_offsets,
197 _stream,
198 false));
199 }
200
201 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
202 DeviceSegmentedSort& SortPairs(cub::DoubleBuffer<KeyT>& d_keys,
203 cub::DoubleBuffer<ValueT>& d_values,
204 int num_items,
205 int num_segments,
206 BeginOffsetIteratorT d_begin_offsets,
207 EndOffsetIteratorT d_end_offsets)
208 {
209 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::SortPairs(d_temp_storage,
210 temp_storage_bytes,
211 d_keys,
212 d_values,
213 num_items,
214 num_segments,
215 d_begin_offsets,
216 d_end_offsets,
217 _stream,
218 false));
219 }
220
221
222 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
223 DeviceSegmentedSort& SortPairsDescending(cub::DoubleBuffer<KeyT>& d_keys,
224 cub::DoubleBuffer<ValueT>& d_values,
225 int num_items,
226 int num_segments,
227 BeginOffsetIteratorT d_begin_offsets,
228 EndOffsetIteratorT d_end_offsets)
229 {
230 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::SortPairsDescending(
231 d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
232 }
233
234
235 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
236 DeviceSegmentedSort& StableSortPairs(const KeyT* d_keys_in,
237 KeyT* d_keys_out,
238 const ValueT* d_values_in,
239 ValueT* d_values_out,
240 int num_items,
241 int num_segments,
242 BeginOffsetIteratorT d_begin_offsets,
243 EndOffsetIteratorT d_end_offsets)
244 {
245 MUDA_CUB_WRAPPER_IMPL(
246 cub::DeviceSegmentedSort::StableSortPairs(d_temp_storage,
247 temp_storage_bytes,
248 d_keys_in,
249 d_keys_out,
250 d_values_in,
251 d_values_out,
252 num_items,
253 num_segments,
254 d_begin_offsets,
255 d_end_offsets,
256 _stream,
257 false));
258 }
259
260
261 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
262 DeviceSegmentedSort& StableSortPairsDescending(const KeyT* d_keys_in,
263 KeyT* d_keys_out,
264 const ValueT* d_values_in,
265 ValueT* d_values_out,
266 int num_items,
267 int num_segments,
268 BeginOffsetIteratorT d_begin_offsets,
269 EndOffsetIteratorT d_end_offsets)
270 {
271 MUDA_CUB_WRAPPER_IMPL(
272 cub::DeviceSegmentedSort::StableSortPairsDescending(d_temp_storage,
273 temp_storage_bytes,
274 d_keys_in,
275 d_keys_out,
276 d_values_in,
277 d_values_out,
278 num_items,
279 num_segments,
280 d_begin_offsets,
281 d_end_offsets,
282 _stream,
283 false));
284 }
285
286
287 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
288 DeviceSegmentedSort& StableSortPairs(cub::DoubleBuffer<KeyT>& d_keys,
289 cub::DoubleBuffer<ValueT>& d_values,
290 int num_items,
291 int num_segments,
292 BeginOffsetIteratorT d_begin_offsets,
293 EndOffsetIteratorT d_end_offsets)
294 {
295 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::StableSortPairs(
296 d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
297 }
298
299
300 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
301 DeviceSegmentedSort& StableSortPairsDescending(cub::DoubleBuffer<KeyT>& d_keys,
302 cub::DoubleBuffer<ValueT>& d_values,
303 int num_items,
304 int num_segments,
305 BeginOffsetIteratorT d_begin_offsets,
306 EndOffsetIteratorT d_end_offsets)
307 {
308 MUDA_CUB_WRAPPER_IMPL(cub::DeviceSegmentedSort::StableSortPairsDescending(
309 d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
310 }
311
312 // Origin:
313
314 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
315 DeviceSegmentedSort& SortKeys(void* d_temp_storage,
316 size_t& temp_storage_bytes,
317 const KeyT* d_keys_in,
318 KeyT* d_keys_out,
319 int num_items,
320 int num_segments,
321 BeginOffsetIteratorT d_begin_offsets,
322 EndOffsetIteratorT d_end_offsets)
323 {
324 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
325 cub::DeviceSegmentedSort::SortKeys(d_temp_storage,
326 temp_storage_bytes,
327 d_keys_in,
328 d_keys_out,
329 num_items,
330 num_segments,
331 d_begin_offsets,
332 d_end_offsets,
333 _stream,
334 false));
335 }
336
337 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
338 DeviceSegmentedSort& SortKeysDescending(void* d_temp_storage,
339 size_t& temp_storage_bytes,
340 const KeyT* d_keys_in,
341 KeyT* d_keys_out,
342 int num_items,
343 int num_segments,
344 BeginOffsetIteratorT d_begin_offsets,
345 EndOffsetIteratorT d_end_offsets)
346 {
347 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
348 cub::DeviceSegmentedSort::SortKeysDescending(d_temp_storage,
349 temp_storage_bytes,
350 d_keys_in,
351 d_keys_out,
352 num_items,
353 num_segments,
354 d_begin_offsets,
355 d_end_offsets,
356 _stream,
357 false));
358 }
359
360
361 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
362 DeviceSegmentedSort& SortKeys(void* d_temp_storage,
363 size_t& temp_storage_bytes,
364 cub::DoubleBuffer<KeyT>& d_keys,
365 int num_items,
366 int num_segments,
367 BeginOffsetIteratorT d_begin_offsets,
368 EndOffsetIteratorT d_end_offsets)
369 {
370 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::SortKeys(
371 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
372 }
373
374
375 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
376 DeviceSegmentedSort& SortKeysDescending(void* d_temp_storage,
377 size_t& temp_storage_bytes,
378 cub::DoubleBuffer<KeyT>& d_keys,
379 int num_items,
380 int num_segments,
381 BeginOffsetIteratorT d_begin_offsets,
382 EndOffsetIteratorT d_end_offsets)
383 {
384 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::SortKeysDescending(
385 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
386 }
387
388
389 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
390 DeviceSegmentedSort& StableSortKeys(void* d_temp_storage,
391 size_t& temp_storage_bytes,
392 const KeyT* d_keys_in,
393 KeyT* d_keys_out,
394 int num_items,
395 int num_segments,
396 BeginOffsetIteratorT d_begin_offsets,
397 EndOffsetIteratorT d_end_offsets)
398 {
399 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
400 cub::DeviceSegmentedSort::StableSortKeys(d_temp_storage,
401 temp_storage_bytes,
402 d_keys_in,
403 d_keys_out,
404 num_items,
405 num_segments,
406 d_begin_offsets,
407 d_end_offsets,
408 _stream,
409 false));
410 }
411
412
413 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
414 DeviceSegmentedSort& StableSortKeysDescending(void* d_temp_storage,
415 size_t& temp_storage_bytes,
416 const KeyT* d_keys_in,
417 KeyT* d_keys_out,
418 int num_items,
419 int num_segments,
420 BeginOffsetIteratorT d_begin_offsets,
421 EndOffsetIteratorT d_end_offsets)
422 {
423 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
424 cub::DeviceSegmentedSort::StableSortKeysDescending(d_temp_storage,
425 temp_storage_bytes,
426 d_keys_in,
427 d_keys_out,
428 num_items,
429 num_segments,
430 d_begin_offsets,
431 d_end_offsets,
432 _stream,
433 false));
434 }
435
436
437 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
438 DeviceSegmentedSort& StableSortKeys(void* d_temp_storage,
439 size_t& temp_storage_bytes,
440 cub::DoubleBuffer<KeyT>& d_keys,
441 int num_items,
442 int num_segments,
443 BeginOffsetIteratorT d_begin_offsets,
444 EndOffsetIteratorT d_end_offsets)
445 {
446 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::StableSortKeys(
447 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
448 }
449
450
451 template <typename KeyT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
452 DeviceSegmentedSort& StableSortKeysDescending(void* d_temp_storage,
453 size_t& temp_storage_bytes,
454 cub::DoubleBuffer<KeyT>& d_keys,
455 int num_items,
456 int num_segments,
457 BeginOffsetIteratorT d_begin_offsets,
458 EndOffsetIteratorT d_end_offsets)
459 {
460 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::StableSortKeysDescending(
461 d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
462 }
463
464
465 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
466 DeviceSegmentedSort& SortPairs(void* d_temp_storage,
467 size_t& temp_storage_bytes,
468 const KeyT* d_keys_in,
469 KeyT* d_keys_out,
470 const ValueT* d_values_in,
471 ValueT* d_values_out,
472 int num_items,
473 int num_segments,
474 BeginOffsetIteratorT d_begin_offsets,
475 EndOffsetIteratorT d_end_offsets)
476 {
477 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
478 cub::DeviceSegmentedSort::SortPairs(d_temp_storage,
479 temp_storage_bytes,
480 d_keys_in,
481 d_keys_out,
482 d_values_in,
483 d_values_out,
484 num_items,
485 num_segments,
486 d_begin_offsets,
487 d_end_offsets,
488 _stream,
489 false));
490 }
491
492
493 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
494 DeviceSegmentedSort& SortPairsDescending(void* d_temp_storage,
495 size_t& temp_storage_bytes,
496 const KeyT* d_keys_in,
497 KeyT* d_keys_out,
498 const ValueT* d_values_in,
499 ValueT* d_values_out,
500 int num_items,
501 int num_segments,
502 BeginOffsetIteratorT d_begin_offsets,
503 EndOffsetIteratorT d_end_offsets)
504 {
505 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
506 cub::DeviceSegmentedSort::SortPairsDescending(d_temp_storage,
507 temp_storage_bytes,
508 d_keys_in,
509 d_keys_out,
510 d_values_in,
511 d_values_out,
512 num_items,
513 num_segments,
514 d_begin_offsets,
515 d_end_offsets,
516 _stream,
517 false));
518 }
519
520 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
521 DeviceSegmentedSort& SortPairs(void* d_temp_storage,
522 size_t& temp_storage_bytes,
523 cub::DoubleBuffer<KeyT>& d_keys,
524 cub::DoubleBuffer<ValueT>& d_values,
525 int num_items,
526 int num_segments,
527 BeginOffsetIteratorT d_begin_offsets,
528 EndOffsetIteratorT d_end_offsets)
529 {
530 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::SortPairs(
531 d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
532 }
533
534
535 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
536 DeviceSegmentedSort& SortPairsDescending(void* d_temp_storage,
537 size_t& temp_storage_bytes,
538 cub::DoubleBuffer<KeyT>& d_keys,
539 cub::DoubleBuffer<ValueT>& d_values,
540 int num_items,
541 int num_segments,
542 BeginOffsetIteratorT d_begin_offsets,
543 EndOffsetIteratorT d_end_offsets)
544 {
545 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::SortPairsDescending(
546 d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
547 }
548
549
550 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
551 DeviceSegmentedSort& StableSortPairs(void* d_temp_storage,
552 size_t& temp_storage_bytes,
553 const KeyT* d_keys_in,
554 KeyT* d_keys_out,
555 const ValueT* d_values_in,
556 ValueT* d_values_out,
557 int num_items,
558 int num_segments,
559 BeginOffsetIteratorT d_begin_offsets,
560 EndOffsetIteratorT d_end_offsets)
561 {
562 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
563 cub::DeviceSegmentedSort::StableSortPairs(d_temp_storage,
564 temp_storage_bytes,
565 d_keys_in,
566 d_keys_out,
567 d_values_in,
568 d_values_out,
569 num_items,
570 num_segments,
571 d_begin_offsets,
572 d_end_offsets,
573 _stream,
574 false));
575 }
576
577
578 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
579 DeviceSegmentedSort& StableSortPairsDescending(void* d_temp_storage,
580 size_t& temp_storage_bytes,
581 const KeyT* d_keys_in,
582 KeyT* d_keys_out,
583 const ValueT* d_values_in,
584 ValueT* d_values_out,
585 int num_items,
586 int num_segments,
587 BeginOffsetIteratorT d_begin_offsets,
588 EndOffsetIteratorT d_end_offsets)
589 {
590 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
591 cub::DeviceSegmentedSort::StableSortPairsDescending(d_temp_storage,
592 temp_storage_bytes,
593 d_keys_in,
594 d_keys_out,
595 d_values_in,
596 d_values_out,
597 num_items,
598 num_segments,
599 d_begin_offsets,
600 d_end_offsets,
601 _stream,
602 false));
603 }
604
605
606 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
607 DeviceSegmentedSort& StableSortPairs(void* d_temp_storage,
608 size_t& temp_storage_bytes,
609 cub::DoubleBuffer<KeyT>& d_keys,
610 cub::DoubleBuffer<ValueT>& d_values,
611 int num_items,
612 int num_segments,
613 BeginOffsetIteratorT d_begin_offsets,
614 EndOffsetIteratorT d_end_offsets)
615 {
616 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(cub::DeviceSegmentedSort::StableSortPairs(
617 d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, num_segments, d_begin_offsets, d_end_offsets, _stream, false));
618 }
619
620
621 template <typename KeyT, typename ValueT, typename BeginOffsetIteratorT, typename EndOffsetIteratorT>
622 DeviceSegmentedSort& StableSortPairsDescending(void* d_temp_storage,
623 size_t& temp_storage_bytes,
624 cub::DoubleBuffer<KeyT>& d_keys,
625 cub::DoubleBuffer<ValueT>& d_values,
626 int num_items,
627 int num_segments,
628 BeginOffsetIteratorT d_begin_offsets,
629 EndOffsetIteratorT d_end_offsets)
630 {
631 MUDA_CUB_WRAPPER_FOR_COMPUTE_GRAPH_IMPL(
632 cub::DeviceSegmentedSort::StableSortPairsDescending(d_temp_storage,
633 temp_storage_bytes,
634 d_keys,
635 d_values,
636 num_items,
637 num_segments,
638 d_begin_offsets,
639 d_end_offsets,
640 _stream,
641 false));
642 }
643};
644} // namespace muda
645
646#include "details/cub_wrapper_macro_undef.inl"
muda::CubWrapper
Definition cub_wrapper.h:14
muda::DeviceSegmentedSort
Definition device_segmented_sort.h:12
muda::LaunchCore
Definition launch_base.h:42