Skip to content

[MINOR] Route leftover temps through local temp_mr #23730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 5 commits into from
Aug 25, 2026
Merged

[MINOR] Route leftover temps through local temp_mr #23730

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
327bbd9
addressing @igorpeshansky's comments
nirandaperera Aug 19, 2026
27b9887
secondary audit
nirandaperera Aug 19, 2026
586870c
Merge branch 'main' of github.com:NVIDIA/cudf into fix/pr-23665-post-…
nirandaperera Aug 19, 2026
12da70b
Modify left_table_keep_mask initialization
nirandaperera Aug 25, 2026
e27879a
Merge branch 'main' into fix/pr-23665-post-merge-comments
nirandaperera Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 9 additions & 8 deletions cpp/src/dictionary/detail/concatenate.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,9 +103,11 @@ struct compute_children_offsets_fn {
* are used to create the offsets.
*
* @param stream Stream used for allocating the output rmm::device_uvector.
* @param mr Device memory resource used to allocate the returned device vector.
* @return Vector of offsets_pair objects for keys and indices.
*/
rmm::device_uvector<offsets_pair> create_children_offsets(cuda::stream_ref stream)
rmm::device_uvector<offsets_pair> create_children_offsets(cuda::stream_ref stream,
rmm::device_async_resource_ref mr)
{
auto offsets = cudf::detail::make_host_vector<offsets_pair>(columns_ptrs.size(), stream);
thrust::transform_exclusive_scan(
Expand All @@ -121,8 +123,7 @@ struct compute_children_offsets_fn {
[](auto lhs, auto rhs) {
return offsets_pair{lhs.first + rhs.first, lhs.second + rhs.second};
});
return cudf::detail::make_device_uvector(
offsets, stream, cudf::get_current_device_resource_ref());
return cudf::detail::make_device_uvector(offsets, stream, mr);
}

private:
Expand Down Expand Up @@ -175,19 +176,19 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
return keys;
});

auto const temp_mr = cudf::get_current_device_resource_ref();

// TODO: Overload function to accept multiple vectors to do the concatenate at once, with a 2D
// kernel. The keys concatenate below and the indices concatenate further down are two separate
// launches over the same set of input columns and could be fused into a single batched call.
// first, concatenate all the keys
auto all_keys =
cudf::detail::concatenate(keys_views, stream, cudf::get_current_device_resource_ref());
auto all_keys = cudf::detail::concatenate(keys_views, stream, temp_mr);
// compute the unique set of keys to better help map the new indices values
using encode_probe_t = cuco::linear_probing<
1,
cudf::detail::row::hash::device_row_hasher<cudf::hashing::detail::default_hash,
cudf::nullate::NO>>;
auto const tv = cudf::table_view({all_keys->view()});
auto const temp_mr = cudf::get_current_device_resource_ref();
auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream, temp_mr);
auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream, temp_mr);
auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
Expand All @@ -205,7 +206,7 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
auto iota = cuda::counting_iterator<size_type>{0};

auto d_indices = rmm::device_uvector<size_type>(all_keys->size(), stream, temp_mr);
auto d_all_keys = column_device_view::create(all_keys->view(), stream);
auto d_all_keys = column_device_view::create(all_keys->view(), stream, temp_mr);
thrust::transform(
policy, iota, iota + all_keys->size(), d_indices.begin(), insert_keys_fn{set_ref, *d_all_keys});
auto keys_indices = rmm::device_uvector<size_type>(all_keys->size(), stream, temp_mr);
Expand Down Expand Up @@ -246,7 +247,7 @@ std::unique_ptr<column> concatenate(host_span<column_view const> columns,
all_indices->type(), all_indices->size(), mask_state::UNALLOCATED, stream, mr);
auto output_view = indices_column->mutable_view();
auto input_view = column_device_view::create(all_indices->view(), stream);
auto children_offsets = child_offsets_fn.create_children_offsets(stream);
auto children_offsets = child_offsets_fn.create_children_offsets(stream, temp_mr);
auto map_fn = map_indices_fn{children_offsets, final_remap, *input_view};
thrust::transform(
policy, iota, iota + all_indices->size(), output_view.begin<size_type>(), map_fn);
Expand Down
17 changes: 8 additions & 9 deletions cpp/src/join/filtered_join/filtered_join.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,10 +130,10 @@ std::unique_ptr<rmm::device_uvector<cudf::size_type>> filtered_join::semi_anti_j
{
cudf::scoped_range range{"filtered_join::semi_anti_join"};

auto const preprocessed_left = [left, stream] {
auto const temp_mr = cudf::get_current_device_resource_ref();
auto const preprocessed_left = [&left, stream, temp_mr] {
cudf::scoped_range range{"filtered_join::semi_anti_join::preprocessed_left"};
return cudf::detail::row::equality::preprocessed_table::create(
left, stream, cudf::get_current_device_resource_ref());
return cudf::detail::row::equality::preprocessed_table::create(left, stream, temp_mr);
}();

auto contains_map = rmm::device_uvector<bool>(left.num_rows(), stream);
Expand All @@ -147,12 +147,11 @@ std::unique_ptr<rmm::device_uvector<cudf::size_type>> filtered_join::semi_anti_j
}

rmm::device_uvector<size_type> gather_map(left.num_rows(), stream, mr);
auto gather_map_end =
thrust::copy_if(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
cuda::counting_iterator<size_type>{0},
cuda::counting_iterator<size_type>{left.num_rows()},
gather_map.begin(),
gather_mask{kind, contains_map_span});
auto gather_map_end = thrust::copy_if(rmm::exec_policy_nosync(stream, temp_mr),
cuda::counting_iterator<size_type>{0},
cuda::counting_iterator<size_type>{left.num_rows()},
gather_map.begin(),
gather_mask{kind, contains_map_span});
gather_map.resize(cuda::std::distance(gather_map.begin(), gather_map_end), stream);
return std::make_unique<rmm::device_uvector<size_type>>(std::move(gather_map));
}
Expand Down
6 changes: 3 additions & 3 deletions cpp/src/join/mark_join.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -741,10 +741,10 @@ std::unique_ptr<rmm::device_uvector<cudf::size_type>> mark_join::semi_anti_join(
{
clear_marks(stream);

auto const preprocessed_right = [right, stream] {
auto const temp_mr = cudf::get_current_device_resource_ref();
auto const preprocessed_right = [&right, stream, temp_mr] {
cudf::scoped_range range{"mark_join::semi_anti_join::preprocessed_right"};
return cudf::detail::row::equality::preprocessed_table::create(
right, stream, cudf::get_current_device_resource_ref());
return cudf::detail::row::equality::preprocessed_table::create(right, stream, temp_mr);
}();

if (is_primitive_row_op_compatible(_left)) {
Expand Down
3 changes: 1 addition & 2 deletions cpp/src/join/mixed_join_semi.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,8 +158,7 @@ std::unique_ptr<rmm::device_uvector<size_type>> mixed_join_semi(
row_set.insert_async(iter, iter + right_num_rows, stream.get());
} else {
cuda::counting_iterator<cudf::size_type> stencil(0);
auto const [row_bitmask, _] =
cudf::detail::bitmask_and(right, stream, cudf::get_current_device_resource_ref());
auto const [row_bitmask, _] = cudf::detail::bitmask_and(right, stream, temp_mr);
row_is_valid pred{static_cast<bitmask_type const*>(row_bitmask.data())};

// insert valid rows
Expand Down
4 changes: 2 additions & 2 deletions cpp/src/sort/rank.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ rmm::device_uvector<size_type> sorted_dense_rank(column_view input_col,
rmm::device_uvector<size_type> dense_rank_sorted(input_size, stream);

@igorpeshansky igorpeshansky Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Optional] Should this also use temp_mr, like you did with d_result in group_nunique? Or would this wait until sorted_dense_rank gets a memory resource parameter?

@nirandaperera nirandaperera Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dense_rank_sorted is the output isnt it? So, it should ideally be set with the output MR, which we will do once the rank.cu is ported to memory_resources

@igorpeshansky igorpeshansky Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's the return value of this (helper) function, but the output (of rank) is rank_column (line 265). This is a const scratch key vector (see the comment in lines 286-287). So using temp_mr here is correct, unless I'm missing something.


auto const comparator_helper = [&](auto const device_comparator) {
thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
cuda::counting_iterator<cudf::size_type>{0},
cuda::counting_iterator{input_size},
dense_rank_sorted.data(),
Expand All @@ -88,7 +88,7 @@ rmm::device_uvector<size_type> sorted_dense_rank(column_view input_col,
comparator_helper(device_comparator);
}

thrust::inclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
thrust::inclusive_scan(rmm::exec_policy_nosync(stream, temp_mr),
dense_rank_sorted.begin(),
dense_rank_sorted.end(),
dense_rank_sorted.data());
Expand Down
6 changes: 3 additions & 3 deletions cpp/tests/utilities/column_utilities.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -882,9 +882,9 @@ bool expect_columns_equal(cudf::column_view const& lhs,
cuda::stream_ref stream,
cudf::memory_resources mr)
{
// TODO: equality row preprocessing (two_table_comparator / preprocessed_table::create) still
// allocates from the current device resource; pass `mr` through once that path accepts
// memory_resources so callers need not disable failing current-resource scopes.
// TODO: check_non_empty_nulls (via has_nonempty_nulls) still allocates temporaries from the
// current device resource. Once it accepts a memory resource, callers can guard comparisons
// with fail_on_current_device_resource_use().
check_non_empty_nulls(lhs, rhs, stream);
auto lhs_indices = generate_all_row_indices(lhs.size(), stream, mr);
auto rhs_indices = generate_all_row_indices(rhs.size(), stream, mr);
Expand Down
7 changes: 3 additions & 4 deletions cpp/tests/utilities_tests/column_wrapper_tests.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -410,8 +410,6 @@ TYPED_TEST(FixedWidthColumnWrapperTest, NullablePairListConstructorAllNullMatch)
this->resources());
cudf::column_view view = col;

// TODO: has_nonempty_nulls (via count_if/transform_reduce) still allocates temporaries from the
// current device resource for strings columns.
Comment on lines -413 to -414

@igorpeshansky igorpeshansky Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So does this mean we can now add fail_on_current_device_resource_use() into this test? Since has_nonempty_nulls does not reach count_if for a fixed-width type?

@nirandaperera nirandaperera Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, fixed width types should work without a problem. Strings and lists were the problematic types.

@igorpeshansky igorpeshansky Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, so let's add the assertion here to give the test some teeth?

CUDF_TEST_EXPECT_COLUMNS_EQUAL(view,
match_view,
cudf::test::debug_output_level::FIRST_ERROR,
Expand Down Expand Up @@ -517,8 +515,9 @@ TYPED_TEST(StringsColumnWrapperTest, NullablePairListConstructorAllNullMatch)
this->resources());
cudf::column_view view = col;

// TODO: has_nonempty_nulls (via count_if/transform_reduce) still allocates temporaries from the
// current device resource for strings columns.
// TODO: check_non_empty_nulls (via has_nonempty_nulls) still allocates temporaries from the

@igorpeshansky igorpeshansky Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your original phrasing pointed at count_if/transform_reduce, which are reachable from has_nonempty_null_rows called from has_nonempty_nulls. That's likely too detailed here, but maybe another TODO is warranted in has_nonempty_null_rows (purge_nonempty_nulls.cu) as well? Per #20780 ("APIs that don't return device memory"), subsequent PRs would add an mr parameter to has_nonempty_nulls (used in cudf::column_device_view::create), but propagating that into count_if→transform_reduce would happen later (aside: transform_reduce is not mentioned in #20780).

The above brings up another minor inaccuracy in my earlier suggested wording: simply accepting a memory resource in has_nonempty_nulls is not sufficient — it also has to be fully propagated through the call chain (into count_if→transform_reduce). So maybe "accepts and fully propagates" is better here (and in column_utilities.cu)…

@igorpeshansky igorpeshansky Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, the actual requests might've gotten lost in the above wall of text. Should we (a) add a TODO in has_nonempty_null_rows and (b) change this TODO to say "accepts and fully propagates"?

// current device resource for string columns. Once it accepts a memory resource, guard this
// comparison with fail_on_current_device_resource_use().
CUDF_TEST_EXPECT_COLUMNS_EQUAL(view,
match_view,
cudf::test::debug_output_level::FIRST_ERROR,
Expand Down
Loading