|
1 | 1 | /* |
2 | | - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. |
| 2 | + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
3 | 3 | * SPDX-License-Identifier: Apache-2.0 |
4 | 4 | */ |
5 | 5 |
|
6 | 6 | #include "distinct_helpers.hpp" |
7 | 7 |
|
8 | 8 | #include <cudf/detail/algorithms/copy_if.cuh> |
| 9 | +#include <cudf/stream_compaction.hpp> |
| 10 | +#include <cudf/types.hpp> |
| 11 | +#include <cudf/utilities/memory_resource.hpp> |
| 12 | + |
| 13 | +#include <rmm/cuda_stream_view.hpp> |
| 14 | +#include <rmm/exec_policy.hpp> |
| 15 | +#include <rmm/resource_ref.hpp> |
9 | 16 |
|
10 | 17 | #include <cuda/functional> |
11 | 18 | #include <cuda/iterator> |
12 | | -#include <cuda/std/atomic> |
13 | 19 | #include <cuda/std/iterator> |
| 20 | +#include <thrust/uninitialized_fill.h> |
14 | 21 |
|
15 | 22 | namespace cudf::detail { |
16 | 23 |
|
17 | | -template <typename RowEqual> |
18 | | -rmm::device_uvector<size_type> reduce_by_row(distinct_set_t<RowEqual>& set, |
19 | | - size_type num_rows, |
20 | | - duplicate_keep_option keep, |
21 | | - rmm::cuda_stream_view stream, |
22 | | - rmm::device_async_resource_ref mr) |
| 24 | +void initialize_reduction_results(size_type* results, |
| 25 | + size_type num_rows, |
| 26 | + duplicate_keep_option keep, |
| 27 | + rmm::cuda_stream_view stream) |
23 | 28 | { |
24 | | - auto output_indices = rmm::device_uvector<size_type>(num_rows, stream, mr); |
25 | | - |
26 | | - // If we don't care about order, just gather indices of distinct keys taken from set. |
27 | | - if (keep == duplicate_keep_option::KEEP_ANY) { |
28 | | - auto const iter = cuda::counting_iterator<cudf::size_type>{0}; |
29 | | - set.insert_async(iter, iter + num_rows, stream.value()); |
30 | | - auto const output_end = set.retrieve_all(output_indices.begin(), stream.value()); |
31 | | - output_indices.resize(cuda::std::distance(output_indices.begin(), output_end), stream); |
32 | | - return output_indices; |
33 | | - } |
34 | | - |
35 | | - auto reduction_results = rmm::device_uvector<size_type>(num_rows, stream, mr); |
36 | 29 | thrust::uninitialized_fill( |
37 | 30 | rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), |
38 | | - reduction_results.begin(), |
39 | | - reduction_results.end(), |
| 31 | + results, |
| 32 | + results + num_rows, |
40 | 33 | reduction_init_value(keep)); |
| 34 | +} |
41 | 35 |
|
42 | | - auto set_ref = set.ref(cuco::op::insert_and_find); |
43 | | - |
44 | | - thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), |
45 | | - cuda::counting_iterator<cudf::size_type>{0}, |
46 | | - cuda::counting_iterator{num_rows}, |
47 | | - [set_ref, keep, reduction_results = reduction_results.begin()] __device__( |
48 | | - size_type const idx) mutable { |
49 | | - auto const [inserted_idx_ptr, _] = set_ref.insert_and_find(idx); |
50 | | - |
51 | | - auto ref = cuda::atomic_ref<size_type, cuda::thread_scope_device>{ |
52 | | - reduction_results[*inserted_idx_ptr]}; |
53 | | - if (keep == duplicate_keep_option::KEEP_FIRST) { |
54 | | - // Store the smallest index of all rows that are equal. |
55 | | - ref.fetch_min(idx, cuda::memory_order_relaxed); |
56 | | - } else if (keep == duplicate_keep_option::KEEP_LAST) { |
57 | | - // Store the greatest index of all rows that are equal. |
58 | | - ref.fetch_max(idx, cuda::memory_order_relaxed); |
59 | | - } else { |
60 | | - // Count the number of rows in each group of rows that are compared equal. |
61 | | - ref.fetch_add(size_type{1}, cuda::memory_order_relaxed); |
62 | | - } |
63 | | - }); |
64 | | - |
65 | | - auto const map_end = [&] { |
| 36 | +size_type copy_reduction_results(size_type const* results, |
| 37 | + size_type num_rows, |
| 38 | + size_type* output, |
| 39 | + duplicate_keep_option keep, |
| 40 | + rmm::cuda_stream_view stream) |
| 41 | +{ |
| 42 | + auto const output_end = [&] { |
66 | 43 | if (keep == duplicate_keep_option::KEEP_NONE) { |
67 | | - // Reduction results with `KEEP_NONE` are either group sizes of equal rows, or `0`. |
68 | | - // Thus, we only output index of the rows in the groups having group size of `1`. |
| 44 | + // KEEP_NONE stores group sizes; retain only singleton groups. |
69 | 45 | return cudf::detail::copy_if( |
70 | 46 | cuda::counting_iterator<size_type>{0}, |
71 | 47 | cuda::counting_iterator<size_type>{num_rows}, |
72 | | - output_indices.begin(), |
| 48 | + output, |
73 | 49 | cuda::proclaim_return_type<bool>( |
74 | | - [reduction_results = reduction_results.begin()] __device__(auto const idx) { |
75 | | - return reduction_results[idx] == size_type{1}; |
76 | | - }), |
| 50 | + [results] __device__(auto const idx) { return results[idx] == size_type{1}; }), |
77 | 51 | stream); |
78 | 52 | } |
79 | 53 |
|
80 | | - // Reduction results with `KEEP_FIRST` and `KEEP_LAST` are row indices of the first/last row in |
81 | | - // each group of equal rows (which are the desired output indices), or the value given by |
82 | | - // `reduction_init_value()`. |
| 54 | + // KEEP_FIRST and KEEP_LAST store desired row indices or the mode's initial marker. |
83 | 55 | return cudf::detail::copy_if( |
84 | | - reduction_results.begin(), |
85 | | - reduction_results.end(), |
86 | | - output_indices.begin(), |
| 56 | + results, |
| 57 | + results + num_rows, |
| 58 | + output, |
87 | 59 | cuda::proclaim_return_type<bool>([init_value = reduction_init_value(keep)] __device__( |
88 | 60 | auto const idx) { return idx != init_value; }), |
89 | 61 | stream); |
90 | 62 | }(); |
91 | 63 |
|
92 | | - output_indices.resize(cuda::std::distance(output_indices.begin(), map_end), stream); |
93 | | - return output_indices; |
| 64 | + return cuda::std::distance(output, output_end); |
94 | 65 | } |
95 | 66 |
|
96 | | -template rmm::device_uvector<size_type> reduce_by_row( |
97 | | - distinct_set_t<cudf::detail::row::equality::device_row_comparator< |
98 | | - false, |
99 | | - cudf::nullate::DYNAMIC, |
100 | | - cudf::detail::row::equality::nan_equal_physical_equality_comparator>>& set, |
101 | | - size_type num_rows, |
102 | | - duplicate_keep_option keep, |
103 | | - rmm::cuda_stream_view stream, |
104 | | - rmm::device_async_resource_ref mr); |
105 | | - |
106 | | -template rmm::device_uvector<size_type> reduce_by_row( |
107 | | - distinct_set_t<cudf::detail::row::equality::device_row_comparator< |
108 | | - true, |
109 | | - cudf::nullate::DYNAMIC, |
110 | | - cudf::detail::row::equality::nan_equal_physical_equality_comparator>>& set, |
111 | | - size_type num_rows, |
112 | | - duplicate_keep_option keep, |
113 | | - rmm::cuda_stream_view stream, |
114 | | - rmm::device_async_resource_ref mr); |
115 | | - |
116 | | -template rmm::device_uvector<size_type> reduce_by_row( |
117 | | - distinct_set_t<cudf::detail::row::equality::device_row_comparator< |
118 | | - false, |
119 | | - cudf::nullate::DYNAMIC, |
120 | | - cudf::detail::row::equality::physical_equality_comparator>>& set, |
121 | | - size_type num_rows, |
122 | | - duplicate_keep_option keep, |
123 | | - rmm::cuda_stream_view stream, |
124 | | - rmm::device_async_resource_ref mr); |
125 | | - |
126 | | -template rmm::device_uvector<size_type> reduce_by_row( |
127 | | - distinct_set_t<cudf::detail::row::equality::device_row_comparator< |
128 | | - true, |
129 | | - cudf::nullate::DYNAMIC, |
130 | | - cudf::detail::row::equality::physical_equality_comparator>>& set, |
131 | | - size_type num_rows, |
132 | | - duplicate_keep_option keep, |
133 | | - rmm::cuda_stream_view stream, |
134 | | - rmm::device_async_resource_ref mr); |
135 | | - |
136 | 67 | } // namespace cudf::detail |
0 commit comments