diff --git a/cpp/include/cudf/detail/sorting.hpp b/cpp/include/cudf/detail/sorting.hpp index bf48a7d29da..8943ea879d2 100644 --- a/cpp/include/cudf/detail/sorting.hpp +++ b/cpp/include/cudf/detail/sorting.hpp @@ -27,7 +27,7 @@ std::unique_ptr sorted_order(table_view const& input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr); + cudf::memory_resources mr); /** * @copydoc cudf::stable_sorted_order @@ -38,7 +38,7 @@ std::unique_ptr stable_sorted_order(table_view const& input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr); + cudf::memory_resources mr); /** * @copydoc cudf::sort_by_key @@ -64,7 +64,7 @@ std::unique_ptr rank(column_view const& input, null_order null_precedence, bool percentage, cuda::stream_ref stream, - rmm::device_async_resource_ref mr); + cudf::memory_resources mr); /** * @copydoc cudf::stable_sort_by_key diff --git a/cpp/include/cudf/sorting.hpp b/cpp/include/cudf/sorting.hpp index 77a80177f80..c8346cd1cd1 100644 --- a/cpp/include/cudf/sorting.hpp +++ b/cpp/include/cudf/sorting.hpp @@ -37,7 +37,7 @@ namespace CUDF_EXPORT cudf { * for each column. Size must be equal to `input.num_columns()` or empty. * If empty, all columns will be sorted in `null_order::BEFORE`. * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned column's device memory + * @param mr Memory resources used for temporary allocations and the returned column * @return A non-nullable column of elements containing the permuted row indices of * `input` if it were sorted */ @@ -46,7 +46,7 @@ std::unique_ptr sorted_order( std::vector const& column_order = {}, std::vector const& null_precedence = {}, cuda::stream_ref stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + cudf::memory_resources mr = cudf::get_current_device_resource_ref()); /** * @brief Computes the row indices that would produce `input` in a stable @@ -61,7 +61,7 @@ std::unique_ptr stable_sorted_order( std::vector const& column_order = {}, std::vector const& null_precedence = {}, cuda::stream_ref stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + cudf::memory_resources mr = cudf::get_current_device_resource_ref()); /** * @brief Checks whether the rows of a `table` are sorted in a lexicographical @@ -216,18 +216,17 @@ std::unique_ptr stable_sort_by_key( * @param null_precedence The desired order of null rows compared to other elements * @param percentage Flag to convert ranks to percentage in range (0,1] * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned column's device memory + * @param mr Memory resources used for temporary allocations and the returned column * @return A column of containing the rank of the each element of the column of `input` */ -std::unique_ptr rank( - column_view const& input, - rank_method method, - order column_order, - null_policy null_handling, - null_order null_precedence, - bool percentage, - cuda::stream_ref stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +std::unique_ptr rank(column_view const& input, + rank_method method, + order column_order, + null_policy null_handling, + null_order null_precedence, + bool percentage, + cuda::stream_ref stream = cudf::get_default_stream(), + cudf::memory_resources mr = cudf::get_current_device_resource_ref()); /** * @brief Returns sorted order after sorting each segment in the table. diff --git a/cpp/src/sort/rank.cu b/cpp/src/sort/rank.cu index 5c7e6b0c11e..0fd74f5a04b 100644 --- a/cpp/src/sort/rank.cu +++ b/cpp/src/sort/rank.cu @@ -57,17 +57,17 @@ struct unique_functor { // Assign rank from 1 to n unique values. Equal values get same rank value. rmm::device_uvector sorted_dense_rank(column_view input_col, column_view sorted_order_view, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { auto const t_input = table_view{{input_col}}; - auto const temp_mr = cudf::get_current_device_resource_ref(); auto const comparator = cudf::detail::row::equality::self_comparator{t_input, stream, temp_mr}; auto const sorted_index_order = cuda::make_permutation_iterator( sorted_order_view.begin(), cuda::counting_iterator{0}); auto const input_size = input_col.size(); - rmm::device_uvector dense_rank_sorted(input_size, stream); + rmm::device_uvector dense_rank_sorted(input_size, stream, temp_mr); auto const comparator_helper = [&](auto const device_comparator) { thrust::transform(rmm::exec_policy_nosync(stream, temp_mr), @@ -120,13 +120,14 @@ void tie_break_ranks_transform(cudf::device_span dense_rank_sor outputIterator rank_iter, TieBreaker tie_breaker, Transformer transformer, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { auto const input_size = sorted_order_view.size(); // algorithm: reduce_by_key(dense_rank, 1, n, reduction_tie_breaker) // reduction_tie_breaker = min, max, min_count - rmm::device_uvector tie_sorted(sorted_order_view.size(), stream); - thrust::reduce_by_key(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + rmm::device_uvector tie_sorted(sorted_order_view.size(), stream, temp_mr); + thrust::reduce_by_key(rmm::exec_policy_nosync(stream, temp_mr), dense_rank_sorted.begin(), dense_rank_sorted.end(), tie_iter, @@ -142,7 +143,7 @@ void tie_break_ranks_transform(cudf::device_span dense_rank_sor [tied_rank = tie_sorted.begin(), transformer] __device__(auto dense_pos) { return transformer(tied_rank[dense_pos - 1]); })); - thrust::scatter(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + thrust::scatter(rmm::exec_policy_nosync(stream, temp_mr), sorted_tied_rank, sorted_tied_rank + input_size, sorted_order_view.begin(), @@ -152,10 +153,11 @@ void tie_break_ranks_transform(cudf::device_span dense_rank_sor template void rank_first(column_view sorted_order_view, mutable_column_view rank_mutable_view, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { // stable sort order ranking (no ties) - thrust::scatter(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + thrust::scatter(rmm::exec_policy_nosync(stream, temp_mr), cuda::counting_iterator{1}, cuda::counting_iterator{rank_mutable_view.size() + 1}, sorted_order_view.begin(), @@ -166,10 +168,11 @@ template void rank_dense(cudf::device_span dense_rank_sorted, column_view sorted_order_view, mutable_column_view rank_mutable_view, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { // All equal values have same rank and rank always increases by 1 between groups - thrust::scatter(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + thrust::scatter(rmm::exec_policy_nosync(stream, temp_mr), dense_rank_sorted.begin(), dense_rank_sorted.end(), sorted_order_view.begin(), @@ -180,7 +183,8 @@ template void rank_min(cudf::device_span group_keys, column_view sorted_order_view, mutable_column_view rank_mutable_view, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { // min of first in the group // All equal values have min of ranks among them. @@ -191,14 +195,16 @@ void rank_min(cudf::device_span group_keys, rank_mutable_view.begin(), cuda::minimum{}, cuda::std::identity{}, - stream); + stream, + temp_mr); } template void rank_max(cudf::device_span group_keys, column_view sorted_order_view, mutable_column_view rank_mutable_view, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { // max of first in the group // All equal values have max of ranks among them. @@ -209,7 +215,8 @@ void rank_max(cudf::device_span group_keys, rank_mutable_view.begin(), cuda::maximum{}, cuda::std::identity{}, - stream); + stream, + temp_mr); } // Returns index, count @@ -221,7 +228,8 @@ struct index_counter { void rank_average(cudf::device_span group_keys, column_view sorted_order_view, mutable_column_view rank_mutable_view, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { // k, k+1, .. k+n-1 // average = (n*k+ n*(n-1)/2)/n @@ -245,7 +253,8 @@ void rank_average(cudf::device_span group_keys, return static_cast(cuda::std::get<0>(minrank_count)) + (static_cast(cuda::std::get<1>(minrank_count)) - 1) / 2.0; }), - stream); + stream, + temp_mr); } } // anonymous namespace @@ -257,77 +266,85 @@ std::unique_ptr rank(column_view const& input, null_order null_precedence, bool percentage, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { + auto const output_mr = mr.get_output_mr(); + auto const temp_mr = mr.get_temporary_mr(); data_type const output_type = (percentage or method == rank_method::AVERAGE) ? data_type(type_id::FLOAT64) : data_type(type_to_id()); - std::unique_ptr rank_column = [&null_handling, &output_type, &input, &stream, &mr] { + std::unique_ptr rank_column = [&null_handling, &output_type, &input, &stream, output_mr] { // na_option=keep assign NA to NA values if (null_handling == null_policy::EXCLUDE) return make_numeric_column(output_type, input.size(), - detail::copy_bitmask(input, stream, mr), + detail::copy_bitmask(input, stream, output_mr), input.null_count(), stream, - mr); + output_mr); else - return make_numeric_column(output_type, input.size(), mask_state::UNALLOCATED, stream, mr); + return make_numeric_column( + output_type, input.size(), mask_state::UNALLOCATED, stream, output_mr); }(); auto rank_mutable_view = rank_column->mutable_view(); std::unique_ptr sorted_order = (method == rank_method::FIRST) ? detail::stable_sorted_order( - table_view{{input}}, {column_order}, {null_precedence}, stream, mr) - : detail::sorted_order(table_view{{input}}, {column_order}, {null_precedence}, stream, mr); + table_view{{input}}, {column_order}, {null_precedence}, stream, {temp_mr, temp_mr}) + : detail::sorted_order( + table_view{{input}}, {column_order}, {null_precedence}, stream, {temp_mr, temp_mr}); column_view sorted_order_view = sorted_order->view(); // dense: All equal values have same rank and rank always increases by 1 between groups // acts as key for min, max, average to denote equal value groups rmm::device_uvector const dense_rank_sorted = - [&method, &input, &sorted_order_view, &stream] { + [&method, &input, &sorted_order_view, &stream, temp_mr] { if (method != rank_method::FIRST) - return sorted_dense_rank(input, sorted_order_view, stream); + return sorted_dense_rank(input, sorted_order_view, stream, temp_mr); else - return rmm::device_uvector(0, stream); + return rmm::device_uvector(0, stream, temp_mr); }(); if (output_type.id() == type_id::FLOAT64) { switch (method) { case rank_method::FIRST: - rank_first(sorted_order_view, rank_mutable_view, stream); + rank_first(sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::DENSE: - rank_dense(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_dense( + dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::MIN: - rank_min(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_min(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::MAX: - rank_max(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_max(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::AVERAGE: - rank_average(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_average(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; default: CUDF_FAIL("Unexpected rank_method for rank()"); } } else { switch (method) { case rank_method::FIRST: - rank_first(sorted_order_view, rank_mutable_view, stream); + rank_first(sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::DENSE: - rank_dense(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_dense( + dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::MIN: - rank_min(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_min( + dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::MAX: - rank_max(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_max( + dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; case rank_method::AVERAGE: - rank_average(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream); + rank_average(dense_rank_sorted, sorted_order_view, rank_mutable_view, stream, temp_mr); break; default: CUDF_FAIL("Unexpected rank_method for rank()"); } @@ -341,7 +358,7 @@ std::unique_ptr rank(column_view const& input, auto drs = dense_rank_sorted.data(); bool const is_dense = (method == rank_method::DENSE); thrust::transform( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + rmm::exec_policy_nosync(stream, temp_mr), rank_iter, rank_iter + input.size(), rank_iter, @@ -360,7 +377,7 @@ std::unique_ptr rank(column_view const& input, null_order null_precedence, bool percentage, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { CUDF_FUNC_RANGE(); return detail::rank( diff --git a/cpp/src/sort/sort.cu b/cpp/src/sort/sort.cu index 36e6beda142..4aa5e433d38 100644 --- a/cpp/src/sort/sort.cu +++ b/cpp/src/sort/sort.cu @@ -23,7 +23,7 @@ std::unique_ptr sorted_order(table_view const& input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { return sorted_order(input, column_order, null_precedence, stream, mr); } @@ -72,7 +72,7 @@ std::unique_ptr sorted_order(table_view const& input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { CUDF_FUNC_RANGE(); return detail::sorted_order(input, column_order, null_precedence, stream, mr); diff --git a/cpp/src/sort/sort.hpp b/cpp/src/sort/sort.hpp index d5ac2a7eced..58dd4c81b2a 100644 --- a/cpp/src/sort/sort.hpp +++ b/cpp/src/sort/sort.hpp @@ -29,7 +29,7 @@ enum class sort_method : bool { STABLE, UNSTABLE }; * @param column_order Ascending or descending sort order * @param null_precedence How null rows are to be ordered * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned column's device memory + * @param mr Memory resources used for temporary allocations and the returned column * @return Sorted indices for the input column. */ template @@ -37,7 +37,7 @@ std::unique_ptr sorted_order(column_view const& input, order column_order, null_order null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr); + cudf::memory_resources mr); } // namespace detail } // namespace cudf diff --git a/cpp/src/sort/sort_column.cu b/cpp/src/sort/sort_column.cu index b6077e1556f..f3a51703df6 100644 --- a/cpp/src/sort/sort_column.cu +++ b/cpp/src/sort/sort_column.cu @@ -17,20 +17,24 @@ namespace detail { /** * @copydoc - * sorted_order(column_view&,order,null_order,cuda::stream_ref,rmm::device_async_resource_ref) + * sorted_order(column_view&,order,null_order,cuda::stream_ref,memory_resources) */ template <> std::unique_ptr sorted_order(column_view const& input, order column_order, null_order null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { - auto sorted_indices = cudf::make_numeric_column( - data_type(type_to_id()), input.size(), mask_state::UNALLOCATED, stream, mr); + auto sorted_indices = cudf::make_numeric_column(data_type(type_to_id()), + input.size(), + mask_state::UNALLOCATED, + stream, + mr.get_output_mr()); mutable_column_view indices_view = sorted_indices->mutable_view(); if (is_radix_sortable(input)) { - sorted_order_radix(input, indices_view, column_order == order::ASCENDING, stream); + sorted_order_radix( + input, indices_view, column_order == order::ASCENDING, stream, mr.get_temporary_mr()); } else { cudf::type_dispatcher(input.type(), column_sorted_order_fn{}, @@ -38,7 +42,8 @@ std::unique_ptr sorted_order(column_view const& i indices_view, column_order == order::ASCENDING, null_precedence, - stream); + stream, + mr.get_temporary_mr()); } return sorted_indices; } diff --git a/cpp/src/sort/sort_column_impl.cuh b/cpp/src/sort/sort_column_impl.cuh index fd48deb56e1..e869e29d0ef 100644 --- a/cpp/src/sort/sort_column_impl.cuh +++ b/cpp/src/sort/sort_column_impl.cuh @@ -74,9 +74,10 @@ struct column_sorted_order_fn { mutable_column_view& indices, bool ascending, null_order null_precedence, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { - auto keys = column_device_view::create(input, stream); + auto keys = column_device_view::create(input, stream, temp_mr); auto comp = simple_comparator{*keys, input.has_nulls(), ascending, null_precedence}; auto in_keys = cuda::counting_iterator{0}; auto out_keys = indices.begin(); @@ -84,13 +85,13 @@ struct column_sorted_order_fn { if constexpr (method == sort_method::STABLE) { cub::DeviceMergeSort::StableSortKeysCopy( nullptr, tmp_bytes, in_keys, out_keys, indices.size(), comp, stream.get()); - auto tmp_stg = rmm::device_buffer(tmp_bytes, stream); + auto tmp_stg = rmm::device_buffer(tmp_bytes, stream, temp_mr); cub::DeviceMergeSort::StableSortKeysCopy( tmp_stg.data(), tmp_bytes, in_keys, out_keys, indices.size(), comp, stream.get()); } else { cub::DeviceMergeSort::SortKeysCopy( nullptr, tmp_bytes, in_keys, out_keys, indices.size(), comp, stream.get()); - auto tmp_stg = rmm::device_buffer(tmp_bytes, stream); + auto tmp_stg = rmm::device_buffer(tmp_bytes, stream, temp_mr); cub::DeviceMergeSort::SortKeysCopy( tmp_stg.data(), tmp_bytes, in_keys, out_keys, indices.size(), comp, stream.get()); } @@ -102,14 +103,20 @@ struct column_sorted_order_fn { mutable_column_view& indices, bool ascending, null_order null_precedence, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { - sorted_order(input, indices, ascending, null_precedence, stream); + sorted_order(input, indices, ascending, null_precedence, stream, temp_mr); } template requires(not cudf::is_relationally_comparable()) - void operator()(column_view const&, mutable_column_view&, bool, null_order, cuda::stream_ref) + void operator()(column_view const&, + mutable_column_view&, + bool, + null_order, + cuda::stream_ref, + rmm::device_async_resource_ref) { CUDF_FAIL("Column type must be relationally comparable"); } @@ -120,24 +127,24 @@ struct column_sorted_order_fn { mutable_column_view& indices, bool ascending, null_order null_precedence, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { auto const keys = dictionary_column_view(input).keys(); // For the keys we do an arg-sort of arg-sort to get the rank and use that as a map // to sort the indices in rank order. // First, get sorted-order of just the keys (slow but expect keys.size <<< indices.size) - auto temp_mr = cudf::get_current_device_resource_ref(); - auto ordered_indices = - cudf::detail::sorted_order(keys, order::ASCENDING, null_precedence, stream, temp_mr); + auto ordered_indices = cudf::detail::sorted_order( + keys, order::ASCENDING, null_precedence, stream, {temp_mr, temp_mr}); // Now, sort the ordered indices to get their ordered positions (very fast integer sort) ordered_indices = cudf::detail::sorted_order( - ordered_indices->view(), order::ASCENDING, null_precedence, stream, temp_mr); + ordered_indices->view(), order::ASCENDING, null_precedence, stream, {temp_mr, temp_mr}); // And use the result as a map over the dictionary indices auto map = ordered_indices->view().template data(); auto itr = cudf::detail::indexalator_factory::make_input_iterator( dictionary_column_view(input).indices()); - auto mapped_indices = rmm::device_uvector(input.size(), stream); - thrust::gather(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + auto mapped_indices = rmm::device_uvector(input.size(), stream, temp_mr); + thrust::gather(rmm::exec_policy_nosync(stream, temp_mr), itr, itr + input.size(), map, @@ -151,9 +158,9 @@ struct column_sorted_order_fn { input.null_count()); // these should be very fast since they are sorting integers if (input.has_nulls()) { - sorted_order(mapped_view, indices, ascending, null_precedence, stream); + sorted_order(mapped_view, indices, ascending, null_precedence, stream, temp_mr); } else { - sorted_order_radix(mapped_view, indices, ascending, stream); + sorted_order_radix(mapped_view, indices, ascending, stream, temp_mr); } } }; diff --git a/cpp/src/sort/sort_impl.cuh b/cpp/src/sort/sort_impl.cuh index 5e719f7f3aa..d9319e88eda 100644 --- a/cpp/src/sort/sort_impl.cuh +++ b/cpp/src/sort/sort_impl.cuh @@ -22,8 +22,8 @@ namespace detail { /** * @copydoc - * sorted_order(table_view&,std::vector,std::vector,rmm::device_async_resource_ref - * ) + * sorted_order(table_view&,std::vector,std::vector,cuda::stream_ref, + * memory_resources) * * @tparam stable Whether to use stable sort * @param stream CUDA stream used for device memory operations and kernel launches @@ -33,11 +33,13 @@ std::unique_ptr sorted_order(table_view input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { + auto const output_mr = mr.get_output_mr(); + auto const temp_mr = mr.get_temporary_mr(); if (input.num_rows() == 0 or input.num_columns() == 0) { return cudf::make_numeric_column( - data_type(type_to_id()), 0, mask_state::UNALLOCATED, stream, mr); + data_type(type_to_id()), 0, mask_state::UNALLOCATED, stream, output_mr); } if (not column_order.empty()) { @@ -58,10 +60,14 @@ std::unique_ptr sorted_order(table_view input, return sorted_order(single_col, col_order, null_prec, stream, mr); } - std::unique_ptr sorted_indices = cudf::make_numeric_column( - data_type(type_to_id()), input.num_rows(), mask_state::UNALLOCATED, stream, mr); + std::unique_ptr sorted_indices = + cudf::make_numeric_column(data_type(type_to_id()), + input.num_rows(), + mask_state::UNALLOCATED, + stream, + output_mr); mutable_column_view mutable_indices_view = sorted_indices->mutable_view(); - thrust::sequence(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + thrust::sequence(rmm::exec_policy_nosync(stream, temp_mr), mutable_indices_view.begin(), mutable_indices_view.end(), 0); @@ -70,18 +76,20 @@ std::unique_ptr sorted_order(table_view input, // Compiling `thrust::*sort*` APIs is expensive. // Thus, we should optimize that by using constexpr condition to only compile what we need. if constexpr (method == sort_method::STABLE) { - thrust::stable_sort(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + thrust::stable_sort(rmm::exec_policy_nosync(stream, temp_mr), mutable_indices_view.begin(), mutable_indices_view.end(), comparator); } else { - thrust::sort(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + thrust::sort(rmm::exec_policy_nosync(stream, temp_mr), mutable_indices_view.begin(), mutable_indices_view.end(), comparator); } }; + // Preprocessing here allocates from the current device resource; the comparator does not + // yet accept a memory resource. auto const comp = cudf::detail::row::lexicographic::self_comparator(input, column_order, null_precedence, stream); if (cudf::detail::has_nested_columns(input)) { diff --git a/cpp/src/sort/sort_radix.hpp b/cpp/src/sort/sort_radix.hpp index 4e077666427..41454b095c2 100644 --- a/cpp/src/sort/sort_radix.hpp +++ b/cpp/src/sort/sort_radix.hpp @@ -45,10 +45,12 @@ std::unique_ptr sort_radix(column_view const& input, * @param indices The indices to return * @param ascending The sort order * @param stream The CUDA stream to use + * @param temp_mr Device memory resource used for temporary allocations */ void sorted_order_radix(column_view const& input, mutable_column_view& indices, bool ascending, - cuda::stream_ref stream); + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr); } // namespace detail } // namespace cudf diff --git a/cpp/src/sort/sorted_order_radix.cu b/cpp/src/sort/sorted_order_radix.cu index d703fa38baa..40182275f30 100644 --- a/cpp/src/sort/sorted_order_radix.cu +++ b/cpp/src/sort/sorted_order_radix.cu @@ -55,22 +55,20 @@ struct float_to_pair_and_seq { * Should not be called if `input.has_nulls()==true` */ struct sorted_order_radix_fn { - column_view const& input; // keys to sort - mutable_column_view& indices; // output of sort - bool ascending; // true for ascending sort - cuda::stream_ref stream; // for allocation and kernel launches + column_view const& input; // keys to sort + mutable_column_view& indices; // output of sort + bool ascending; // true for ascending sort + cuda::stream_ref stream; // for allocation and kernel launches + rmm::device_async_resource_ref temp_mr; // temporary allocations template void radix_sort() { auto d_in = input.begin(); - auto output = rmm::device_uvector(input.size(), stream); + auto output = rmm::device_uvector(input.size(), stream, temp_mr); auto d_out = output.begin(); // not returned - auto seqs = rmm::device_uvector(input.size(), stream); - thrust::sequence(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - seqs.begin(), - seqs.end(), - 0); + auto seqs = rmm::device_uvector(input.size(), stream, temp_mr); + thrust::sequence(rmm::exec_policy_nosync(stream, temp_mr), seqs.begin(), seqs.end(), 0); auto dv_in = seqs.begin(); auto dv_out = indices.begin(); @@ -83,13 +81,13 @@ struct sorted_order_radix_fn { if (ascending) { cub::DeviceRadixSort::SortPairs( nullptr, tmp_bytes, d_in, d_out, dv_in, dv_out, n, 0, end_bit, sv); - auto tmp_stg = rmm::device_buffer(tmp_bytes, stream); + auto tmp_stg = rmm::device_buffer(tmp_bytes, stream, temp_mr); cub::DeviceRadixSort::SortPairs( tmp_stg.data(), tmp_bytes, d_in, d_out, dv_in, dv_out, n, 0, end_bit, sv); } else { cub::DeviceRadixSort::SortPairsDescending( nullptr, tmp_bytes, d_in, d_out, dv_in, dv_out, n, 0, end_bit, sv); - auto tmp_stg = rmm::device_buffer(tmp_bytes, stream); + auto tmp_stg = rmm::device_buffer(tmp_bytes, stream, temp_mr); cub::DeviceRadixSort::SortPairsDescending( tmp_stg.data(), tmp_bytes, d_in, d_out, dv_in, dv_out, n, 0, end_bit, sv); } @@ -99,17 +97,17 @@ struct sorted_order_radix_fn { void operator()() requires(cudf::is_floating_point()) { - auto pair_in = rmm::device_uvector>(input.size(), stream); + auto pair_in = rmm::device_uvector>(input.size(), stream, temp_mr); auto d_in = pair_in.begin(); // pair_out/d_out is not returned to the caller but used as an intermediate - auto pair_out = rmm::device_uvector>(input.size(), stream); + auto pair_out = rmm::device_uvector>(input.size(), stream, temp_mr); auto d_out = pair_out.begin(); - auto vals = rmm::device_uvector(indices.size(), stream); + auto vals = rmm::device_uvector(indices.size(), stream, temp_mr); auto dv_in = vals.begin(); auto dv_out = indices.begin(); auto zip_out = cuda::make_zip_iterator(d_in, dv_in); - 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{0}, cuda::counting_iterator{input.size()}, zip_out, @@ -124,13 +122,13 @@ struct sorted_order_radix_fn { if (ascending) { cub::DeviceRadixSort::SortPairs( nullptr, tmp_bytes, d_in, d_out, dv_in, dv_out, n, decomposer, 0, end_bit, sv); - auto tmp_stg = rmm::device_buffer(tmp_bytes, stream); + auto tmp_stg = rmm::device_buffer(tmp_bytes, stream, temp_mr); cub::DeviceRadixSort::SortPairs( tmp_stg.data(), tmp_bytes, d_in, d_out, dv_in, dv_out, n, decomposer, 0, end_bit, sv); } else { cub::DeviceRadixSort::SortPairsDescending( nullptr, tmp_bytes, d_in, d_out, dv_in, dv_out, n, decomposer, 0, end_bit, sv); - auto tmp_stg = rmm::device_buffer(tmp_bytes, stream); + auto tmp_stg = rmm::device_buffer(tmp_bytes, stream, temp_mr); cub::DeviceRadixSort::SortPairsDescending( tmp_stg.data(), tmp_bytes, d_in, d_out, dv_in, dv_out, n, decomposer, 0, end_bit, sv); } @@ -169,14 +167,16 @@ struct sorted_order_radix_fn { * @param indices The result of the sort * @param ascending Sort order * @param stream CUDA stream used for device memory operations and kernel launches + * @param temp_mr Device memory resource used for temporary allocations */ void sorted_order_radix(column_view const& input, mutable_column_view& indices, bool ascending, - cuda::stream_ref stream) + cuda::stream_ref stream, + rmm::device_async_resource_ref temp_mr) { cudf::type_dispatcher( - input.type(), sorted_order_radix_fn{input, indices, ascending, stream}); + input.type(), sorted_order_radix_fn{input, indices, ascending, stream, temp_mr}); } } // namespace detail } // namespace cudf diff --git a/cpp/src/sort/stable_sort.cu b/cpp/src/sort/stable_sort.cu index 7be0e9a3f72..007f7e592cd 100644 --- a/cpp/src/sort/stable_sort.cu +++ b/cpp/src/sort/stable_sort.cu @@ -23,7 +23,7 @@ std::unique_ptr stable_sorted_order(table_view const& input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { return sorted_order(input, column_order, null_precedence, stream, mr); } @@ -71,7 +71,7 @@ std::unique_ptr stable_sorted_order(table_view const& input, std::vector const& column_order, std::vector const& null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { CUDF_FUNC_RANGE(); return detail::stable_sorted_order(input, column_order, null_precedence, stream, mr); diff --git a/cpp/src/sort/stable_sort_column.cu b/cpp/src/sort/stable_sort_column.cu index 8169b4f046b..6969c4008f9 100644 --- a/cpp/src/sort/stable_sort_column.cu +++ b/cpp/src/sort/stable_sort_column.cu @@ -17,20 +17,24 @@ namespace detail { /** * @copydoc - * stable_sorted_order(column_view&,order,null_order,cuda::stream_ref,rmm::device_async_resource_ref) + * stable_sorted_order(column_view&,order,null_order,cuda::stream_ref,memory_resources) */ template <> std::unique_ptr sorted_order(column_view const& input, order column_order, null_order null_precedence, cuda::stream_ref stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources mr) { - auto sorted_indices = cudf::make_numeric_column( - data_type(type_to_id()), input.size(), mask_state::UNALLOCATED, stream, mr); + auto sorted_indices = cudf::make_numeric_column(data_type(type_to_id()), + input.size(), + mask_state::UNALLOCATED, + stream, + mr.get_output_mr()); mutable_column_view indices_view = sorted_indices->mutable_view(); if (is_radix_sortable(input)) { - sorted_order_radix(input, indices_view, column_order == order::ASCENDING, stream); + sorted_order_radix( + input, indices_view, column_order == order::ASCENDING, stream, mr.get_temporary_mr()); } else { cudf::type_dispatcher(input.type(), column_sorted_order_fn{}, @@ -38,7 +42,8 @@ std::unique_ptr sorted_order(column_view const& inp indices_view, column_order == order::ASCENDING, null_precedence, - stream); + stream, + mr.get_temporary_mr()); } return sorted_indices; } diff --git a/cpp/tests/sort/rank_test.cpp b/cpp/tests/sort/rank_test.cpp index 27a24033e77..127ac988eff 100644 --- a/cpp/tests/sort/rank_test.cpp +++ b/cpp/tests/sort/rank_test.cpp @@ -443,6 +443,35 @@ TEST_F(RankLarge, average_large) CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view(), expected); } +struct RankMemoryResourceTest : public cudf::test::BaseFixtureWithHarness {}; + +TEST_F(RankMemoryResourceTest, Resources) +{ + auto const stream = this->stream(); + auto const mr = this->resources(); + auto& harness = this->harness(); + + cudf::test::fixed_width_column_wrapper input( + {30, 10, 20, 10}, stream, harness.setup_mr()); + cudf::test::fixed_width_column_wrapper expected( + {4, 1, 3, 1}, stream, harness.setup_mr()); + + auto result = cudf::rank(input, + cudf::rank_method::MIN, + cudf::order::ASCENDING, + cudf::null_policy::INCLUDE, + cudf::null_order::BEFORE, + false, + stream, + mr); + harness.synchronize(stream); + harness.expect_output_allocations_live(stream); + harness.expect_temporary_allocation_activity(stream); + harness.expect_temporary_allocations_released(stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL( + expected, result->view(), cudf::test::debug_output_level::FIRST_ERROR, stream, mr); +} + template struct RankListAndStruct : public cudf::test::BaseFixture { void run_all_tests(cudf::rank_method method, diff --git a/cpp/tests/sort/sort_test.cpp b/cpp/tests/sort/sort_test.cpp index e4925f99f01..05ff2bbb068 100644 --- a/cpp/tests/sort/sort_test.cpp +++ b/cpp/tests/sort/sort_test.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -1082,4 +1082,39 @@ TEST_F(SortDouble, InfinityAndNan) CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(results->view(), expected); } +struct SortMemoryResourceTest : public cudf::test::BaseFixtureWithHarness {}; + +TEST_F(SortMemoryResourceTest, SortedOrderResources) +{ + auto const stream = this->stream(); + auto const mr = this->resources(); + auto& harness = this->harness(); + + cudf::test::fixed_width_column_wrapper input( + {5, 1, 4, 2, 3}, stream, harness.setup_mr()); + cudf::test::fixed_width_column_wrapper expected( + {1, 3, 4, 2, 0}, stream, harness.setup_mr()); + + { + auto result = cudf::sorted_order(cudf::table_view{{input}}, {}, {}, stream, mr); + harness.synchronize(stream); + harness.expect_output_allocations_live(stream); + harness.expect_temporary_allocation_activity(stream); + harness.expect_temporary_allocations_released(stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL( + expected, result->view(), cudf::test::debug_output_level::FIRST_ERROR, stream, mr); + } + harness.expect_no_live_allocations(stream); + + { + auto result = cudf::stable_sorted_order(cudf::table_view{{input}}, {}, {}, stream, mr); + harness.synchronize(stream); + harness.expect_output_allocations_live(stream); + harness.expect_temporary_allocation_activity(stream); + harness.expect_temporary_allocations_released(stream); + CUDF_TEST_EXPECT_COLUMNS_EQUAL( + expected, result->view(), cudf::test::debug_output_level::FIRST_ERROR, stream, mr); + } +} + CUDF_TEST_PROGRAM_MAIN()