Skip to content

Commit 94286a2

Browse files
nirandapereranirandaperera
authored
[cudf] Thread memory resources through equality/hash row preprocessing (#23665)
A part of #20780. Equality and hash `preprocessed_table` paths still allocated temps from the current device resource, so harness-based tests could not prove explicit MR routing. This change requires `cudf::memory_resources` on those APIs and updates callers accordingly. - Equality/hash `preprocessed_table::create`, `self_comparator`, `two_table_comparator`, and `row_hasher` take required `mr` and use `mr.get_temporary_mr()` for preprocessing temps - Call sites across join/groupby/hash/search/reductions/etc. pass an explicit `mr` - Switch to `BaseFixtureWithHarness` and convert row-operator tests to use harness `stream()`/`resources()`, with TODOs where lexicographic still falls back to the current resource - If a callsite introduces a `auto temp_mr = get_current_resource_ref()`, it replaces all `get_current_resouce_ref()` calls with `temp_mr` in that function scope. ## Latent bug fix - In https://github.com/NVIDIA/cudf/blob/main/cpp/src/stream_compaction/unique.cu#L53 a transient column was allocated on `mr`, which should have been allocated on a `temp_mr`. Fixed Authors: - Niranda Perera (https://github.com/nirandaperera) Approvers: - Bradley Dice (https://github.com/bdice) - Vyas Ramasubramani (https://github.com/vyasr) URL: #23665
1 parent 32f2516 commit 94286a2

49 files changed

Lines changed: 555 additions & 373 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎cpp/include/cudf/detail/row_operator/equality.cuh‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -407,9 +407,12 @@ class self_comparator {
407407
* @param t The table to compare
408408
* @param stream The stream to construct this object on. Not the stream that will be used for
409409
* comparisons using this object.
410+
* @param temp_mr Device memory resource used for temporary allocations
410411
*/
411-
self_comparator(table_view const& t, rmm::cuda_stream_view stream)
412-
: d_t(preprocessed_table::create(t, stream))
412+
self_comparator(table_view const& t,
413+
rmm::cuda_stream_view stream,
414+
rmm::device_async_resource_ref temp_mr)
415+
: d_t(preprocessed_table::create(t, stream, temp_mr))
413416
{
414417
}
415418

@@ -515,10 +518,12 @@ class two_table_comparator {
515518
* @param right The right table to compare.
516519
* @param stream The stream to construct this object on. Not the stream that will be used for
517520
* comparisons using this object.
521+
* @param temp_mr Device memory resource used for temporary allocations
518522
*/
519523
two_table_comparator(table_view const& left,
520524
table_view const& right,
521-
rmm::cuda_stream_view stream);
525+
rmm::cuda_stream_view stream,
526+
rmm::device_async_resource_ref temp_mr);
522527

523528
/**
524529
* @brief Construct an owning object for performing equality comparisons between two rows from two

‎cpp/include/cudf/detail/row_operator/hashing.cuh‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -240,9 +240,12 @@ class row_hasher {
240240
* @param t The table containing rows to hash
241241
* @param stream The stream to construct this object on. Not the stream that will be used for
242242
* comparisons using this object.
243+
* @param temp_mr Device memory resource used for temporary allocations
243244
*/
244-
row_hasher(table_view const& t, rmm::cuda_stream_view stream)
245-
: d_t(preprocessed_table::create(t, stream))
245+
row_hasher(table_view const& t,
246+
rmm::cuda_stream_view stream,
247+
rmm::device_async_resource_ref temp_mr)
248+
: d_t(preprocessed_table::create(t, stream, temp_mr))
246249
{
247250
}
248251

‎cpp/include/cudf/detail/row_operator/preprocessed_table.cuh‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#include <rmm/cuda_stream_view.hpp>
1111
#include <rmm/device_buffer.hpp>
1212
#include <rmm/device_uvector.hpp>
13+
#include <rmm/resource_ref.hpp>
1314

1415
#include <memory>
1516
#include <vector>
@@ -47,10 +48,12 @@ struct preprocessed_table {
4748
*
4849
* @param table The table to preprocess
4950
* @param stream The cuda stream to use while preprocessing.
51+
* @param temp_mr Device memory resource used for temporary allocations
5052
* @return A preprocessed table as shared pointer
5153
*/
5254
static std::shared_ptr<preprocessed_table> create(table_view const& table,
53-
rmm::cuda_stream_view stream);
55+
rmm::cuda_stream_view stream,
56+
rmm::device_async_resource_ref temp_mr);
5457

5558
/**
5659
* @brief Implicit conversion operator to a `table_device_view` of the preprocessed table.

‎cpp/src/binaryop/compiled/struct_binary_ops.cuh‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -144,16 +144,18 @@ void apply_struct_equality_op(mutable_column_view& out,
144144
"Unsupported operator for these types",
145145
cudf::data_type_error);
146146

147-
auto tlhs = table_view{{lhs}};
148-
auto trhs = table_view{{rhs}};
149-
auto table_comparator = cudf::detail::row::equality::two_table_comparator{tlhs, trhs, stream};
147+
auto temp_mr = cudf::get_current_device_resource_ref();
148+
auto tlhs = table_view{{lhs}};
149+
auto trhs = table_view{{rhs}};
150+
auto table_comparator =
151+
cudf::detail::row::equality::two_table_comparator{tlhs, trhs, stream, temp_mr};
150152

151153
auto outd = column_device_view::create(out, stream);
152154
auto optional_iter =
153155
cudf::detail::make_optional_iterator<bool>(*outd, nullate::DYNAMIC{out.has_nulls()});
154156

155157
auto const comparator_helper = [&](auto const device_comparator) {
156-
thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
158+
thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
157159
cuda::counting_iterator<size_type>(0),
158160
cuda::counting_iterator<size_type>(out.size()),
159161
out.begin<bool>(),

‎cpp/src/dictionary/detail/concatenate.cu‎

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -187,30 +187,29 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
187187
cudf::detail::row::hash::device_row_hasher<cudf::hashing::detail::default_hash,
188188
cudf::nullate::NO>>;
189189
auto const tv = cudf::table_view({all_keys->view()});
190-
auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream);
191-
auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream);
190+
auto const temp_mr = cudf::get_current_device_resource_ref();
191+
auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream, temp_mr);
192+
auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream, temp_mr);
192193
auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
193194
auto const d_equal =
194195
row_equal.equal_to<false>(cudf::nullate::NO{}, null_equality::EQUAL, comparator);
195196
auto const empty_key = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
196197
auto probe = encode_probe_t{row_hash.device_hasher(cudf::nullate::NO{})};
197-
auto allocator = rmm::mr::polymorphic_allocator<char>(cudf::get_current_device_resource_ref());
198-
auto set = cuco::static_set{
198+
auto allocator = rmm::mr::polymorphic_allocator<char>(temp_mr);
199+
auto set = cuco::static_set{
199200
all_keys->size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
200201
auto set_ref = set.ref(cuco::insert_and_find);
201202
using set_ref_t = decltype(set_ref);
202203

203-
auto policy = rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref());
204+
auto policy = rmm::exec_policy_nosync(stream, temp_mr);
204205
auto iota = cuda::counting_iterator<size_type>{0};
205206

206-
auto d_indices = rmm::device_uvector<size_type>(
207-
all_keys->size(), stream, cudf::get_current_device_resource_ref());
207+
auto d_indices = rmm::device_uvector<size_type>(all_keys->size(), stream, temp_mr);
208208
auto d_all_keys = column_device_view::create(all_keys->view(), stream);
209209
thrust::transform(
210210
policy, iota, iota + all_keys->size(), d_indices.begin(), insert_keys_fn{set_ref, *d_all_keys});
211-
auto keys_indices = rmm::device_uvector<size_type>(
212-
all_keys->size(), stream, cudf::get_current_device_resource_ref());
213-
auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
211+
auto keys_indices = rmm::device_uvector<size_type>(all_keys->size(), stream, temp_mr);
212+
auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
214213
keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
215214

216215
// use keys_indices to retrieve the keys (gather)
@@ -223,13 +222,11 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
223222

224223
// build an all_keys_remap: abs position in all_keys to new key index
225224
// use scatter to assign new index values: all_keys_remap[keys_indices[i]] = i
226-
auto all_keys_remap = rmm::device_uvector<size_type>(
227-
all_keys->size(), stream, cudf::get_current_device_resource_ref());
225+
auto all_keys_remap = rmm::device_uvector<size_type>(all_keys->size(), stream, temp_mr);
228226
thrust::scatter(
229227
policy, iota, iota + keys_indices.size(), keys_indices.begin(), all_keys_remap.begin());
230228
// use gather to propagate new indices values to all duplicate positions
231-
auto final_remap = rmm::device_uvector<size_type>(
232-
all_keys->size(), stream, cudf::get_current_device_resource_ref());
229+
auto final_remap = rmm::device_uvector<size_type>(all_keys->size(), stream, temp_mr);
233230
thrust::gather(
234231
policy, d_indices.begin(), d_indices.end(), all_keys_remap.begin(), final_remap.begin());
235232

‎cpp/src/dictionary/encode.cu‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -82,36 +82,35 @@ std::unique_ptr<column> encode(column_view const& input,
8282

8383
auto const has_nulls = nullate::DYNAMIC{input.has_nulls()};
8484
auto const tv = cudf::table_view({input});
85-
auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream);
86-
auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream);
85+
auto const temp_mr = cudf::get_current_device_resource_ref();
86+
auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream, temp_mr);
87+
auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream, temp_mr);
8788
auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
8889
auto const d_equal = row_equal.equal_to<false>(has_nulls, null_equality::EQUAL, comparator);
8990
auto const empty_key = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
9091
auto probe = encode_probe_t{row_hash.device_hasher(has_nulls)};
91-
auto allocator = rmm::mr::polymorphic_allocator<char>{};
92+
auto allocator = rmm::mr::polymorphic_allocator<char>{temp_mr};
9293
auto set =
9394
cuco::static_set{input.size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
9495
auto set_ref = set.ref(cuco::insert_and_find);
9596
using set_ref_t = decltype(set_ref);
9697

9798
// build a static_set of the input values
9899
// and keep track of the indices of the unique values
99-
auto d_indices = rmm::device_uvector<size_type>(input.size(), stream);
100-
auto d_input = column_device_view::create(input, stream);
101-
thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
100+
auto d_indices = rmm::device_uvector<size_type>(input.size(), stream, temp_mr);
101+
auto d_input = column_device_view::create(input, stream, temp_mr);
102+
thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
102103
cuda::counting_iterator<size_type>{0},
103104
cuda::counting_iterator<size_type>{input.size()},
104105
d_indices.begin(),
105106
encode_fn{set_ref, *d_input});
106107

107-
auto keys_indices = rmm::device_uvector<size_type>(input.size(), stream);
108+
auto keys_indices = rmm::device_uvector<size_type>(input.size(), stream, temp_mr);
108109
auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
109110
keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
110111

111112
// sort the keys_indices so we can use lower-bound on them
112-
thrust::sort(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
113-
keys_indices.begin(),
114-
keys_indices.end());
113+
thrust::sort(rmm::exec_policy_nosync(stream, temp_mr), keys_indices.begin(), keys_indices.end());
115114

116115
// use keys_indices to retrieve the keys
117116
auto const oob_policy = cudf::out_of_bounds_policy::DONT_CHECK;
@@ -124,7 +123,7 @@ std::unique_ptr<column> encode(column_view const& input,
124123
// call lower-bound with keys_indices and d_indices to get the output indices_column
125124
auto d_result =
126125
cudf::detail::indexalator_factory::make_output_iterator(indices_column->mutable_view());
127-
thrust::lower_bound(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
126+
thrust::lower_bound(rmm::exec_policy_nosync(stream, temp_mr),
128127
keys_indices.begin(),
129128
keys_indices.end(),
130129
d_indices.begin(),

‎cpp/src/dictionary/match_keys.cu‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,14 @@ struct unique_keys_dispatch_fn {
5050

5151
auto const has_nulls = nullate::DYNAMIC{false};
5252
auto const keys_tv = table_view({all_keys});
53-
auto const row_hash = cudf::detail::row::hash::row_hasher(keys_tv, stream);
54-
auto const row_equal = cudf::detail::row::equality::self_comparator(keys_tv, stream);
53+
auto const temp_mr = cudf::get_current_device_resource_ref();
54+
auto const row_hash = cudf::detail::row::hash::row_hasher(keys_tv, stream, temp_mr);
55+
auto const row_equal = cudf::detail::row::equality::self_comparator(keys_tv, stream, temp_mr);
5556
auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
5657
auto const d_equal = row_equal.equal_to<false>(has_nulls, null_equality::EQUAL, comparator);
5758
auto const empty_key = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
5859
auto probe = probe_t{row_hash.device_hasher(has_nulls)};
59-
auto allocator = rmm::mr::polymorphic_allocator<char>{};
60+
auto allocator = rmm::mr::polymorphic_allocator<char>{temp_mr};
6061
auto set = cuco::static_set{
6162
all_keys.size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
6263

@@ -65,7 +66,7 @@ struct unique_keys_dispatch_fn {
6566
set.insert_async(iter, iter + all_keys.size(), stream.get());
6667

6768
// retrieve the indices of all the unique keys
68-
auto keys_indices = rmm::device_uvector<size_type>(all_keys.size(), stream);
69+
auto keys_indices = rmm::device_uvector<size_type>(all_keys.size(), stream, temp_mr);
6970
auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
7071
keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
7172

‎cpp/src/groupby/hash/groupby.cu‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,11 @@ std::unique_ptr<table> dispatch_groupby(table_view const& keys,
4242
auto const has_null = nullate::DYNAMIC{cudf::has_nested_nulls(keys)};
4343
auto const skip_rows_with_nulls = keys_have_nulls and include_null_keys == null_policy::EXCLUDE;
4444

45-
auto preprocessed_keys = cudf::detail::row::hash::preprocessed_table::create(keys, stream);
46-
auto const comparator = cudf::detail::row::equality::self_comparator{preprocessed_keys};
47-
auto const row_hash = cudf::detail::row::hash::row_hasher{std::move(preprocessed_keys)};
48-
auto const d_row_hash = row_hash.device_hasher(has_null);
45+
auto preprocessed_keys = cudf::detail::row::hash::preprocessed_table::create(
46+
keys, stream, cudf::get_current_device_resource_ref());
47+
auto const comparator = cudf::detail::row::equality::self_comparator{preprocessed_keys};
48+
auto const row_hash = cudf::detail::row::hash::row_hasher{std::move(preprocessed_keys)};
49+
auto const d_row_hash = row_hash.device_hasher(has_null);
4950

5051
if (cudf::detail::has_nested_columns(keys)) {
5152
auto const d_row_equal = comparator.equal_to<true>(has_null, null_keys_are_equal);

‎cpp/src/groupby/sort/group_nunique.cu‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,10 @@ std::unique_ptr<column> group_nunique(column_view const& values,
7777

7878
if (num_groups == 0) { return result; }
7979

80+
auto temp_mr = cudf::get_current_device_resource_ref();
8081
auto const values_view = table_view{{values}};
81-
auto const comparator = cudf::detail::row::equality::self_comparator{values_view, stream};
82+
auto const comparator =
83+
cudf::detail::row::equality::self_comparator{values_view, stream, temp_mr};
8284

8385
auto const d_values_view = column_device_view::create(values, stream);
8486

@@ -91,7 +93,7 @@ std::unique_ptr<column> group_nunique(column_view const& values,
9193
null_handling,
9294
group_offsets.data(),
9395
group_labels.data()};
94-
thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
96+
thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
9597
cuda::counting_iterator<size_type>{0},
9698
cuda::counting_iterator<size_type>{values.size()},
9799
d_result.begin(),

‎cpp/src/groupby/sort/group_rank_scan.cu‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,9 @@ std::unique_ptr<column> rank_generator(column_view const& grouped_values,
9494
rmm::device_async_resource_ref mr)
9595
{
9696
auto const grouped_values_view = table_view{{grouped_values}};
97-
auto const comparator = cudf::detail::row::equality::self_comparator{grouped_values_view, stream};
97+
auto const temp_mr = cudf::get_current_device_resource_ref();
98+
auto const comparator =
99+
cudf::detail::row::equality::self_comparator{grouped_values_view, stream, temp_mr};
98100

99101
auto ranks = make_fixed_width_column(
100102
data_type{type_to_id<size_type>()}, grouped_values.size(), mask_state::UNALLOCATED, stream, mr);
@@ -104,7 +106,7 @@ std::unique_ptr<column> rank_generator(column_view const& grouped_values,
104106
auto const permuted_equal =
105107
permuted_row_equality_comparator(d_equal, value_order.begin<size_type>());
106108

107-
thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
109+
thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
108110
cuda::counting_iterator<size_type>(0),
109111
cuda::counting_iterator<size_type>(grouped_values.size()),
110112
mutable_ranks.begin<size_type>(),
@@ -130,14 +132,13 @@ std::unique_ptr<column> rank_generator(column_view const& grouped_values,
130132
cuda::std::reverse_iterator(mutable_ranks.end<size_type>())};
131133
}
132134
}();
133-
thrust::inclusive_scan_by_key(
134-
rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
135-
group_labels_begin,
136-
group_labels_begin + group_labels.size(),
137-
mutable_rank_begin,
138-
mutable_rank_begin,
139-
cuda::std::equal_to{},
140-
scan_op);
135+
thrust::inclusive_scan_by_key(rmm::exec_policy_nosync(stream, temp_mr),
136+
group_labels_begin,
137+
group_labels_begin + group_labels.size(),
138+
mutable_rank_begin,
139+
mutable_rank_begin,
140+
cuda::std::equal_to{},
141+
scan_op);
141142
return ranks;
142143
}
143144
} // namespace

0 commit comments

Comments
 (0)