diff --git a/cpp/include/cudf/copying.hpp b/cpp/include/cudf/copying.hpp
index f1ab6e30f268..d5ed72451864 100644
--- a/cpp/include/cudf/copying.hpp
+++ b/cpp/include/cudf/copying.hpp
@@ -75,15 +75,14 @@ enum class negative_index_policy : bool {
* better performance. If `policy` is set to `DONT_CHECK` and there are out-of-bounds indices
* in the gather map, the behavior is undefined. Defaults to `DONT_CHECK`.
* @param stream CUDA stream used for device memory operations and kernel launches
- * @param mr Device memory resource used to allocate the returned table's device memory
+ * @param mr Memory resources used for temporary allocations and the returned table
* @return Result of the gather
*/
-std::unique_ptr
gather(
- table_view const& source_table,
- column_view const& gather_map,
- out_of_bounds_policy bounds_policy = out_of_bounds_policy::DONT_CHECK,
- cuda::stream_ref stream = cudf::get_default_stream(),
- rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+std::unique_ptr gather(table_view const& source_table,
+ column_view const& gather_map,
+ out_of_bounds_policy bounds_policy = out_of_bounds_policy::DONT_CHECK,
+ cuda::stream_ref stream = cudf::get_default_stream(),
+ cudf::memory_resources mr = cudf::get_current_device_resource_ref());
/**
* @brief Gathers the specified rows of a set of columns according to a gather map.
@@ -112,16 +111,15 @@ std::unique_ptr gather(
* @param bounds_policy Interpretation of out-of-bounds indices
* @param neg_indices Interpretation of a negative index `i` in the `gather_map`
* @param stream CUDA stream used for device memory operations and kernel launches.
- * @param mr Device memory resource used to allocate the returned table's device memory
+ * @param mr Memory resources used for temporary allocations and the returned table
* @return Result of the gather
*/
-std::unique_ptr gather(
- table_view const& source_table,
- column_view const& gather_map,
- out_of_bounds_policy bounds_policy,
- negative_index_policy neg_indices,
- cuda::stream_ref stream = cudf::get_default_stream(),
- rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+std::unique_ptr gather(table_view const& source_table,
+ column_view const& gather_map,
+ out_of_bounds_policy bounds_policy,
+ negative_index_policy neg_indices,
+ cuda::stream_ref stream = cudf::get_default_stream(),
+ cudf::memory_resources mr = cudf::get_current_device_resource_ref());
/**
* @brief Reverses the rows within a table.
diff --git a/cpp/include/cudf/detail/gather.cuh b/cpp/include/cudf/detail/gather.cuh
index e5bb1f9ff575..8a28c67c0b04 100644
--- a/cpp/include/cudf/detail/gather.cuh
+++ b/cpp/include/cudf/detail/gather.cuh
@@ -104,6 +104,7 @@ struct gather_bitmask_functor {
* @param gather_map_end End of the gather map
* @param nullify_out_of_bounds True if map values are checked against `source_size`
* @param stream CUDA stream used for kernel launches.
+ * @param temp_mr Device memory resource used for temporary allocations
*/
template
void gather_helper(InputItr source_itr,
@@ -112,11 +113,12 @@ void gather_helper(InputItr source_itr,
MapIterator gather_map_begin,
MapIterator gather_map_end,
bool nullify_out_of_bounds,
- rmm::cuda_stream_view stream)
+ rmm::cuda_stream_view stream,
+ rmm::device_async_resource_ref temp_mr = cudf::get_current_device_resource_ref())
{
using map_type = typename std::iterator_traits::value_type;
if (nullify_out_of_bounds) {
- thrust::gather_if(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::gather_if(rmm::exec_policy_nosync(stream, temp_mr),
gather_map_begin,
gather_map_end,
gather_map_begin,
@@ -124,7 +126,7 @@ void gather_helper(InputItr source_itr,
target_itr,
bounds_checker{0, source_size});
} else {
- thrust::gather(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::gather(rmm::exec_policy_nosync(stream, temp_mr),
gather_map_begin,
gather_map_end,
source_itr,
@@ -159,7 +161,7 @@ struct column_gatherer {
* @param gather_map_end End of iterator range of integral values representing the gather map
* @param nullify_out_of_bounds Nullify values in `gather_map` that are out of bounds
* @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
*/
template
std::unique_ptr operator()(column_view const& source_column,
@@ -167,7 +169,7 @@ struct column_gatherer {
MapIterator gather_map_end,
bool nullify_out_of_bounds,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
column_gatherer_impl gatherer{};
@@ -199,7 +201,7 @@ struct column_gatherer_impl
std::unique_ptr operator()(column_view const& source_column,
@@ -207,11 +209,12 @@ struct column_gatherer_impl(),
source_column.size(),
@@ -219,7 +222,8 @@ struct column_gatherer_impl {
* @param gather_map_end End of iterator range of integral values representing the gather map
* @param nullify_out_of_bounds Nullify values in `gather_map` that are out of bounds
* @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
*/
template
std::unique_ptr operator()(column_view const& source_column,
@@ -252,7 +256,7 @@ struct column_gatherer_impl {
MapItType gather_map_end,
bool nullify_out_of_bounds,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
if (true == nullify_out_of_bounds) {
return cudf::strings::detail::gather(
@@ -326,42 +330,49 @@ struct column_gatherer_impl {
MapItRoot gather_map_end,
bool nullify_out_of_bounds,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
+ auto const output_mr = mr.get_output_mr();
+
lists_column_view list(column);
auto gather_map_size = std::distance(gather_map_begin, gather_map_end);
// if the gather map is empty, return an empty column
if (gather_map_size == 0) { return empty_like(column); }
+ // List gather helpers still take a single resource ref on this branch; use the output
+ // resource until the dedicated list/segmented gather MR port lands.
// generate gather_data for the next level (N+1)
- lists::detail::gather_data gd = nullify_out_of_bounds
- ? lists::detail::make_gather_data(
- column, gather_map_begin, gather_map_size, stream, mr)
- : lists::detail::make_gather_data(
- column, gather_map_begin, gather_map_size, stream, mr);
+ lists::detail::gather_data gd =
+ nullify_out_of_bounds
+ ? lists::detail::make_gather_data(
+ column, gather_map_begin, gather_map_size, stream, output_mr)
+ : lists::detail::make_gather_data(
+ column, gather_map_begin, gather_map_size, stream, output_mr);
// the nesting case.
if (list.child().type() == cudf::data_type{type_id::LIST}) {
// gather children
- auto child = lists::detail::gather_list_nested(list.get_sliced_child(stream), gd, stream, mr);
+ auto child =
+ lists::detail::gather_list_nested(list.get_sliced_child(stream), gd, stream, output_mr);
// return the final column
return make_lists_column(gather_map_size,
std::move(gd.offsets),
std::move(child),
0,
- rmm::device_buffer{0, stream, mr});
+ rmm::device_buffer{0, stream, output_mr});
}
// it's a leaf. do a regular gather
- auto child = lists::detail::gather_list_leaf(list.get_sliced_child(stream), gd, stream, mr);
+ auto child =
+ lists::detail::gather_list_leaf(list.get_sliced_child(stream), gd, stream, output_mr);
// assemble final column
return make_lists_column(gather_map_size,
std::move(gd.offsets),
std::move(child),
0,
- rmm::device_buffer{0, stream, mr});
+ rmm::device_buffer{0, stream, output_mr});
}
};
@@ -380,7 +391,7 @@ struct column_gatherer_impl {
* @param gather_map_end End of iterator range of integral values representing the gather map
* @param nullify_out_of_bounds Nullify values in `gather_map` that are out of bounds
* @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 New dictionary column with gathered rows.
*/
template
@@ -389,8 +400,11 @@ struct column_gatherer_impl {
MapItType gather_map_end,
bool nullify_out_of_bounds,
rmm::cuda_stream_view 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();
+
dictionary_column_view dictionary(source_column);
auto output_count = std::distance(gather_map_begin, gather_map_end);
if (output_count == 0) return make_empty_column(type_id::DICTIONARY32);
@@ -401,11 +415,11 @@ struct column_gatherer_impl {
// be relatively smallish.
// Also, there are scenarios where the keys are common with other dictionaries
// and the original intention was to share the keys here.
- auto keys_copy = std::make_unique(dictionary.keys(), stream, mr);
+ auto keys_copy = std::make_unique(dictionary.keys(), stream, output_mr);
// Perform gather on just the indices
column_view indices = dictionary.get_indices_annotated();
- auto new_indices =
- cudf::allocate_like(indices, output_count, cudf::mask_allocation_policy::NEVER, stream, mr);
+ auto new_indices = cudf::allocate_like(
+ indices, output_count, cudf::mask_allocation_policy::NEVER, stream, output_mr);
gather_helper(
cudf::detail::indexalator_factory::make_input_iterator(indices),
indices.size(),
@@ -413,8 +427,9 @@ struct column_gatherer_impl {
gather_map_begin,
gather_map_end,
nullify_out_of_bounds,
- stream);
- return make_dictionary_column(std::move(keys_copy), std::move(new_indices), stream, mr);
+ stream,
+ temp_mr);
+ return make_dictionary_column(std::move(keys_copy), std::move(new_indices), stream, output_mr);
}
};
@@ -426,8 +441,10 @@ struct column_gatherer_impl {
MapItRoot gather_map_end,
bool nullify_out_of_bounds,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
+ auto const output_mr = mr.get_output_mr();
+
auto const gather_map_size = std::distance(gather_map_begin, gather_map_end);
if (gather_map_size == 0) { return empty_like(column); }
@@ -477,9 +494,9 @@ struct column_gatherer_impl {
gather_map_size,
std::move(output_struct_members),
0,
- rmm::device_buffer{0, stream, mr}, // Null mask will be fixed up in cudf::gather().
+ rmm::device_buffer{0, stream, output_mr}, // Null mask will be fixed up in cudf::gather().
stream,
- mr);
+ output_mr);
}
};
@@ -532,10 +549,13 @@ void gather_bitmask(table_view const& source,
std::vector>& target,
gather_bitmask_op op,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
if (target.empty()) { return; }
+ auto const output_mr = mr.get_output_mr();
+ auto const temp_mr = mr.get_temporary_mr();
+
// Validate that all target columns have the same size
auto const target_rows = target.front()->size();
CUDF_EXPECTS(std::all_of(target.begin(),
@@ -549,7 +569,7 @@ void gather_bitmask(table_view const& source,
not target[i]->nullable()) {
auto const state =
op == gather_bitmask_op::PASSTHROUGH ? mask_state::ALL_VALID : mask_state::UNINITIALIZED;
- auto mask = cudf::create_null_mask(target[i]->size(), state, stream, mr);
+ auto mask = cudf::create_null_mask(target[i]->size(), state, stream, output_mr);
target[i]->set_null_mask(std::move(mask), 0);
}
}
@@ -559,12 +579,10 @@ void gather_bitmask(table_view const& source,
std::transform(target.begin(), target.end(), target_masks.begin(), [](auto const& col) {
return col->mutable_view().null_mask();
});
- auto d_target_masks =
- make_device_uvector_async(target_masks, stream, cudf::get_current_device_resource_ref());
+ auto d_target_masks = make_device_uvector_async(target_masks, stream, temp_mr);
- auto const device_source = table_device_view::create(source, stream);
- auto d_valid_counts = make_zeroed_device_uvector_async(
- target.size(), stream, cudf::get_current_device_resource_ref());
+ auto const device_source = table_device_view::create(source, stream, temp_mr);
+ auto d_valid_counts = make_zeroed_device_uvector_async(target.size(), stream, temp_mr);
// Dispatch operation enum to get implementation
auto const impl = [op]() {
@@ -621,7 +639,7 @@ void gather_bitmask(table_view const& source,
* better performance. In case there are out-of-bound indices in the gather map, the behavior
* is undefined. Defaults to `DONT_CHECK`.
* @param[in] stream CUDA stream used for device memory operations and kernel launches.
- * @param[in] mr Device memory resource used to allocate the returned table's device memory
+ * @param[in] mr Memory resources used for temporary allocations and the returned table
* @return cudf::table Result of the gather
*/
template
@@ -630,8 +648,10 @@ std::unique_ptr gather(table_view const& source_table,
MapIterator gather_map_end,
out_of_bounds_policy bounds_policy,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
+ auto const output_mr = mr.get_output_mr();
+
std::vector> destination_columns;
// TODO: Could be beneficial to use streams internally here
@@ -661,7 +681,8 @@ std::unique_ptr gather(table_view const& source_table,
gather_bitmask(source_table, gather_map_begin, destination_columns, op, stream, mr);
} else {
for (size_type i = 0; i < source_table.num_columns(); ++i) {
- set_all_valid_null_masks(source_table.column(i), *destination_columns[i], stream, mr);
+ set_all_valid_null_masks(
+ source_table.column(i), *destination_columns[i], stream, output_mr);
}
}
}
diff --git a/cpp/include/cudf/detail/gather.hpp b/cpp/include/cudf/detail/gather.hpp
index a5fdf224228f..36b87f4e71c6 100644
--- a/cpp/include/cudf/detail/gather.hpp
+++ b/cpp/include/cudf/detail/gather.hpp
@@ -21,21 +21,21 @@ namespace cudf {
namespace detail {
/**
- * @copydoc cudf::gather(table_view const&,column_view const&,table_view
- * const&,cudf::out_of_bounds_policy,cudf::negative_index_policy,cuda::stream_ref,
- * rmm::device_async_resource_ref)
+ * @copydoc cudf::gather(table_view const&,column_view const&,out_of_bounds_policy,
+ * negative_index_policy,cuda::stream_ref,rmm::device_async_resource_ref)
+ *
+ * @param mr Memory resources used for temporary allocations and the returned table
*/
std::unique_ptr gather(table_view const& source_table,
column_view const& gather_map,
out_of_bounds_policy bounds_policy,
negative_index_policy neg_indices,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr);
+ memory_resources mr);
/**
- * @copydoc cudf::detail::gather(table_view const&,column_view const&,table_view
- * const&,cudf::out_of_bounds_policy,cudf::negative_index_policy,cuda::stream_ref,
- * rmm::device_async_resource_ref)
+ * @copydoc cudf::detail::gather(table_view const&,column_view const&,out_of_bounds_policy,
+ * negative_index_policy,cuda::stream_ref,memory_resources)
*
* @throws cudf::logic_error if `gather_map` span size is larger than max of `size_type`.
*/
@@ -44,7 +44,7 @@ std::unique_ptr gather(table_view const& source_table,
out_of_bounds_policy bounds_policy,
negative_index_policy neg_indices,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr);
+ memory_resources mr);
} // namespace detail
} // namespace cudf
diff --git a/cpp/include/cudf/detail/row_operator/equality.cuh b/cpp/include/cudf/detail/row_operator/equality.cuh
index f38acc8b7d78..9a5fd68b6ceb 100644
--- a/cpp/include/cudf/detail/row_operator/equality.cuh
+++ b/cpp/include/cudf/detail/row_operator/equality.cuh
@@ -407,9 +407,12 @@ class self_comparator {
* @param t The table to compare
* @param stream The stream to construct this object on. Not the stream that will be used for
* comparisons using this object.
+ * @param temp_mr Device memory resource used for temporary allocations
*/
- self_comparator(table_view const& t, rmm::cuda_stream_view stream)
- : d_t(preprocessed_table::create(t, stream))
+ self_comparator(table_view const& t,
+ rmm::cuda_stream_view stream,
+ rmm::device_async_resource_ref temp_mr)
+ : d_t(preprocessed_table::create(t, stream, temp_mr))
{
}
@@ -515,10 +518,12 @@ class two_table_comparator {
* @param right The right table to compare.
* @param stream The stream to construct this object on. Not the stream that will be used for
* comparisons using this object.
+ * @param temp_mr Device memory resource used for temporary allocations
*/
two_table_comparator(table_view const& left,
table_view const& right,
- rmm::cuda_stream_view stream);
+ rmm::cuda_stream_view stream,
+ rmm::device_async_resource_ref temp_mr);
/**
* @brief Construct an owning object for performing equality comparisons between two rows from two
diff --git a/cpp/include/cudf/detail/row_operator/hashing.cuh b/cpp/include/cudf/detail/row_operator/hashing.cuh
index e71fc5213483..5bb77b125d7c 100644
--- a/cpp/include/cudf/detail/row_operator/hashing.cuh
+++ b/cpp/include/cudf/detail/row_operator/hashing.cuh
@@ -240,9 +240,12 @@ class row_hasher {
* @param t The table containing rows to hash
* @param stream The stream to construct this object on. Not the stream that will be used for
* comparisons using this object.
+ * @param temp_mr Device memory resource used for temporary allocations
*/
- row_hasher(table_view const& t, rmm::cuda_stream_view stream)
- : d_t(preprocessed_table::create(t, stream))
+ row_hasher(table_view const& t,
+ rmm::cuda_stream_view stream,
+ rmm::device_async_resource_ref temp_mr)
+ : d_t(preprocessed_table::create(t, stream, temp_mr))
{
}
diff --git a/cpp/include/cudf/detail/row_operator/preprocessed_table.cuh b/cpp/include/cudf/detail/row_operator/preprocessed_table.cuh
index 233294201ccd..08949be1f0b7 100644
--- a/cpp/include/cudf/detail/row_operator/preprocessed_table.cuh
+++ b/cpp/include/cudf/detail/row_operator/preprocessed_table.cuh
@@ -10,6 +10,7 @@
#include
#include
#include
+#include
#include
#include
@@ -47,10 +48,12 @@ struct preprocessed_table {
*
* @param table The table to preprocess
* @param stream The cuda stream to use while preprocessing.
+ * @param temp_mr Device memory resource used for temporary allocations
* @return A preprocessed table as shared pointer
*/
static std::shared_ptr create(table_view const& table,
- rmm::cuda_stream_view stream);
+ rmm::cuda_stream_view stream,
+ rmm::device_async_resource_ref temp_mr);
/**
* @brief Implicit conversion operator to a `table_device_view` of the preprocessed table.
diff --git a/cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh b/cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh
index 79f50d9339dc..28a2ce1d2759 100644
--- a/cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh
+++ b/cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh
@@ -249,27 +249,29 @@ static sizes_to_offsets_iterator make_sizes_to_offsets_i
* @param result Output iterator for scan result
* @param initial_offset Initial offset to add to scan
* @param stream CUDA stream used for device memory operations and kernel launches
+ * @param temp_mr Device memory resource used for temporary allocations
* @return The last element of the scan
*/
template
-auto sizes_to_offsets(SizesIterator begin,
- SizesIterator end,
- OffsetsIterator result,
- int64_t initial_offset,
- rmm::cuda_stream_view stream)
+auto sizes_to_offsets(
+ SizesIterator begin,
+ SizesIterator end,
+ OffsetsIterator result,
+ int64_t initial_offset,
+ rmm::cuda_stream_view stream,
+ rmm::device_async_resource_ref temp_mr = cudf::get_current_device_resource_ref())
{
using SizeType = cuda::std::iter_value_t;
static_assert(std::is_integral_v,
"Only numeric types are supported by sizes_to_offsets");
- using LastType = std::conditional_t, int64_t, uint64_t>;
- auto last_element =
- cudf::detail::device_scalar(0, stream, cudf::get_current_device_resource_ref());
+ using LastType = std::conditional_t, int64_t, uint64_t>;
+ auto last_element = cudf::detail::device_scalar(0, stream, temp_mr);
auto output_itr =
make_sizes_to_offsets_iterator(result, result + std::distance(begin, end), last_element.data());
// This function uses the type of the initialization parameter as the accumulator type
// when computing the individual scan output elements.
- thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::exclusive_scan(rmm::exec_policy_nosync(stream, temp_mr),
begin,
end,
output_itr,
@@ -295,19 +297,16 @@ auto sizes_to_offsets(SizesIterator begin,
* @param begin The beginning of the input sequence
* @param end The end of the input sequence
* @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 Offsets column and total elements
*/
template
std::pair, size_type> make_offsets_child_column(
- InputIterator begin,
- InputIterator end,
- rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ InputIterator begin, InputIterator end, rmm::cuda_stream_view stream, cudf::memory_resources mr)
{
auto count = static_cast(std::distance(begin, end));
- auto offsets_column =
- make_numeric_column(data_type{type_id::INT32}, count + 1, mask_state::UNALLOCATED, stream, mr);
+ auto offsets_column = make_numeric_column(
+ data_type{type_id::INT32}, count + 1, mask_state::UNALLOCATED, stream, mr.get_output_mr());
auto offsets_view = offsets_column->mutable_view();
auto d_offsets = offsets_view.template data();
@@ -327,7 +326,7 @@ std::pair, size_type> make_offsets_child_column(
auto input_itr = cudf::detail::make_counting_transform_iterator(0, map_fn);
// Use the sizes-to-offsets iterator to compute the total number of elements
auto const total_elements =
- sizes_to_offsets(input_itr, input_itr + count + 1, d_offsets, 0, stream);
+ sizes_to_offsets(input_itr, input_itr + count + 1, d_offsets, 0, stream, mr.get_temporary_mr());
// the offsets are 32-bit so the total must fit in an int32_t
CUDF_EXPECTS(
total_elements <= static_cast(std::numeric_limits::max()),
diff --git a/cpp/include/cudf/dictionary/detail/encode.hpp b/cpp/include/cudf/dictionary/detail/encode.hpp
index 6a212831661c..df5b394ec88a 100644
--- a/cpp/include/cudf/dictionary/detail/encode.hpp
+++ b/cpp/include/cudf/dictionary/detail/encode.hpp
@@ -37,13 +37,13 @@ namespace dictionary::detail {
* @param column The column to dictionary encode.
* @param indices_type The integer type to use for the indices.
* @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 Returns a dictionary column.
*/
std::unique_ptr encode(column_view const& column,
data_type indices_type,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr);
+ cudf::memory_resources mr);
/**
* @brief Create a column by gathering the keys from the provided
@@ -57,12 +57,12 @@ std::unique_ptr encode(column_view const& column,
*
* @param dictionary_column Existing dictionary column.
* @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 New column with type matching the dictionary_column's keys.
*/
std::unique_ptr decode(dictionary_column_view const& dictionary_column,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr);
+ cudf::memory_resources mr);
/**
* @brief Return minimal integer type for the given number of elements.
diff --git a/cpp/include/cudf/dictionary/encode.hpp b/cpp/include/cudf/dictionary/encode.hpp
index d915927e7c0b..514782812e22 100644
--- a/cpp/include/cudf/dictionary/encode.hpp
+++ b/cpp/include/cudf/dictionary/encode.hpp
@@ -47,14 +47,13 @@ namespace dictionary {
* @param column The column to dictionary encode
* @param indices_type The integer type to use for the indices
* @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 Returns a dictionary column
*/
-std::unique_ptr encode(
- column_view const& column,
- data_type indices_type = data_type{type_id::INT32},
- cuda::stream_ref stream = cudf::get_default_stream(),
- rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+std::unique_ptr encode(column_view const& column,
+ data_type indices_type = data_type{type_id::INT32},
+ cuda::stream_ref stream = cudf::get_default_stream(),
+ cudf::memory_resources mr = cudf::get_current_device_resource_ref());
/**
* @brief Create a column by gathering the keys from the provided
@@ -68,13 +67,12 @@ std::unique_ptr encode(
*
* @param dictionary_column Existing dictionary column
* @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 New column with type matching the dictionary_column's keys
*/
-std::unique_ptr decode(
- dictionary_column_view const& dictionary_column,
- cuda::stream_ref stream = cudf::get_default_stream(),
- rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+std::unique_ptr decode(dictionary_column_view const& dictionary_column,
+ cuda::stream_ref stream = cudf::get_default_stream(),
+ cudf::memory_resources mr = cudf::get_current_device_resource_ref());
/** @} */ // end of group
} // namespace dictionary
diff --git a/cpp/include/cudf/strings/detail/gather.cuh b/cpp/include/cudf/strings/detail/gather.cuh
index 759448ac58a7..3af885f9b4f2 100644
--- a/cpp/include/cudf/strings/detail/gather.cuh
+++ b/cpp/include/cudf/strings/detail/gather.cuh
@@ -215,7 +215,7 @@ CUDF_KERNEL void gather_chars_fn_char_parallel(StringIterator strings_begin,
* @param begin Start of index iterator.
* @param end End of index iterator.
* @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 New strings column containing the gathered strings.
*/
template
@@ -223,13 +223,16 @@ std::unique_ptr gather(strings_column_view const& strings,
MapIterator begin,
MapIterator end,
rmm::cuda_stream_view 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();
+
auto const output_count = std::distance(begin, end);
if (output_count == 0) return make_empty_column(type_id::STRING);
// build offsets column
- auto const d_strings = column_device_view::create(strings.parent(), stream);
+ auto const d_strings = column_device_view::create(strings.parent(), stream, temp_mr);
auto const d_in_offsets = cudf::detail::offsetalator_factory::make_input_iterator(
strings.is_empty() ? make_empty_column(type_id::INT32)->view() : strings.offsets(),
strings.offset());
@@ -252,7 +255,7 @@ std::unique_ptr gather(strings_column_view const& strings,
cudf::prefetch::detail::prefetch(strings.chars_begin(stream), strings.chars_size(stream), stream);
// build output char column
- auto out_chars_data = rmm::device_uvector(out_char_bytes, stream, mr);
+ auto out_chars_data = rmm::device_uvector(out_char_bytes, stream, output_mr);
cudf::prefetch::detail::prefetch(out_chars_data, stream);
auto d_out_chars = out_chars_data.data();
@@ -318,7 +321,7 @@ std::unique_ptr gather(strings_column_view const& strings,
stream.value());
// Allocate temporary storage
- auto d_temp_storage = rmm::device_buffer(temp_storage_bytes, stream, mr);
+ auto d_temp_storage = rmm::device_buffer(temp_storage_bytes, stream, temp_mr);
// Run batched copy algorithm
cub::DeviceMemcpy::Batched(d_temp_storage.data(),
@@ -358,7 +361,7 @@ std::unique_ptr gather(strings_column_view const& strings,
* @param end End of index iterator.
* @param nullify_out_of_bounds If true, indices outside the column's range are nullified.
* @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 New strings column containing the gathered strings.
*/
template
@@ -367,7 +370,7 @@ std::unique_ptr gather(strings_column_view const& strings,
MapIterator end,
bool nullify_out_of_bounds,
rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
if (nullify_out_of_bounds) return gather(strings, begin, end, stream, mr);
return gather(strings, begin, end, stream, mr);
diff --git a/cpp/include/cudf/strings/detail/strings_children.cuh b/cpp/include/cudf/strings/detail/strings_children.cuh
index d15e5fa199f1..59f0334be297 100644
--- a/cpp/include/cudf/strings/detail/strings_children.cuh
+++ b/cpp/include/cudf/strings/detail/strings_children.cuh
@@ -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
*/
#pragma once
@@ -115,23 +115,25 @@ rmm::device_uvector make_chars_buffer(column_view const& offsets,
* @param begin The beginning of the input sequence
* @param end The end of the input sequence
* @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 Offsets column and total elements
*/
template
-std::pair, int64_t> make_offsets_child_column(
- InputIterator begin,
- InputIterator end,
- rmm::cuda_stream_view stream,
- rmm::device_async_resource_ref mr)
+std::pair, int64_t> make_offsets_child_column(InputIterator begin,
+ InputIterator end,
+ rmm::cuda_stream_view stream,
+ cudf::memory_resources mr)
{
+ auto const output_mr = mr.get_output_mr();
+ auto const temp_mr = mr.get_temporary_mr();
+
auto constexpr size_type_max = static_cast(std::numeric_limits::max());
auto const lcount = static_cast(std::distance(begin, end));
CUDF_EXPECTS(
lcount <= size_type_max, "Size of output exceeds the column size limit", std::overflow_error);
auto const strings_count = static_cast(lcount);
auto offsets_column = make_numeric_column(
- data_type{type_id::INT32}, strings_count + 1, mask_state::UNALLOCATED, stream, mr);
+ data_type{type_id::INT32}, strings_count + 1, mask_state::UNALLOCATED, stream, output_mr);
auto d_offsets = offsets_column->mutable_view().template data();
// The number of offsets is strings_count+1 so to build the offsets from the sizes
@@ -141,8 +143,8 @@ std::pair, int64_t> make_offsets_child_column(
auto input_itr =
cudf::detail::make_counting_transform_iterator(0, string_offsets_fn{begin, strings_count});
// Use the sizes-to-offsets iterator to compute the total number of elements
- auto const total_bytes =
- cudf::detail::sizes_to_offsets(input_itr, input_itr + strings_count + 1, d_offsets, 0, stream);
+ auto const total_bytes = cudf::detail::sizes_to_offsets(
+ input_itr, input_itr + strings_count + 1, d_offsets, 0, stream, temp_mr);
auto const threshold = cudf::strings::get_offset64_threshold();
CUDF_EXPECTS(cudf::strings::is_large_strings_enabled() || (total_bytes < threshold),
@@ -151,10 +153,10 @@ std::pair, int64_t> make_offsets_child_column(
if (total_bytes >= cudf::strings::get_offset64_threshold()) {
// recompute as int64 offsets when above the threshold
offsets_column = make_numeric_column(
- data_type{type_id::INT64}, strings_count + 1, mask_state::UNALLOCATED, stream, mr);
+ data_type{type_id::INT64}, strings_count + 1, mask_state::UNALLOCATED, stream, output_mr);
auto d_offsets64 = offsets_column->mutable_view().template data();
cudf::detail::sizes_to_offsets(
- input_itr, input_itr + strings_count + 1, d_offsets64, 0, stream);
+ input_itr, input_itr + strings_count + 1, d_offsets64, 0, stream, temp_mr);
}
return std::pair(std::move(offsets_column), total_bytes);
diff --git a/cpp/include/cudf_test/column_wrapper.hpp b/cpp/include/cudf_test/column_wrapper.hpp
index 62579b196294..81b485fc54a6 100644
--- a/cpp/include/cudf_test/column_wrapper.hpp
+++ b/cpp/include/cudf_test/column_wrapper.hpp
@@ -1075,7 +1075,7 @@ class dictionary_column_wrapper : public detail::column_wrapper {
begin, end, stream, mr.get_temporary_mr()),
cudf::data_type{type_id::INT32},
stream,
- mr.get_output_mr());
+ mr);
}
/**
@@ -1117,7 +1117,7 @@ class dictionary_column_wrapper : public detail::column_wrapper {
begin, end, v, stream, mr.get_temporary_mr()),
cudf::data_type{type_id::INT32},
stream,
- mr.get_output_mr());
+ mr);
}
/**
@@ -1307,7 +1307,7 @@ class dictionary_column_wrapper : public detail::column_wrapper {
cudf::dictionary::encode(strings_column_wrapper(begin, end, stream, mr.get_temporary_mr()),
cudf::data_type{type_id::INT32},
stream,
- mr.get_output_mr());
+ mr);
}
/**
@@ -1352,7 +1352,7 @@ class dictionary_column_wrapper : public detail::column_wrapper {
cudf::dictionary::encode(strings_column_wrapper(begin, end, v, stream, mr.get_temporary_mr()),
cudf::data_type{type_id::INT32},
stream,
- mr.get_output_mr());
+ mr);
}
/**
diff --git a/cpp/src/binaryop/compiled/struct_binary_ops.cuh b/cpp/src/binaryop/compiled/struct_binary_ops.cuh
index 07e313968755..50e7f354204c 100644
--- a/cpp/src/binaryop/compiled/struct_binary_ops.cuh
+++ b/cpp/src/binaryop/compiled/struct_binary_ops.cuh
@@ -144,16 +144,18 @@ void apply_struct_equality_op(mutable_column_view& out,
"Unsupported operator for these types",
cudf::data_type_error);
- auto tlhs = table_view{{lhs}};
- auto trhs = table_view{{rhs}};
- auto table_comparator = cudf::detail::row::equality::two_table_comparator{tlhs, trhs, stream};
+ auto temp_mr = cudf::get_current_device_resource_ref();
+ auto tlhs = table_view{{lhs}};
+ auto trhs = table_view{{rhs}};
+ auto table_comparator =
+ cudf::detail::row::equality::two_table_comparator{tlhs, trhs, stream, temp_mr};
auto outd = column_device_view::create(out, stream);
auto optional_iter =
cudf::detail::make_optional_iterator(*outd, nullate::DYNAMIC{out.has_nulls()});
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(0),
cuda::counting_iterator(out.size()),
out.begin(),
diff --git a/cpp/src/copying/gather.cu b/cpp/src/copying/gather.cu
index 0e84a51261d8..1866a8ede1d8 100644
--- a/cpp/src/copying/gather.cu
+++ b/cpp/src/copying/gather.cu
@@ -28,7 +28,7 @@ std::unique_ptr gather(table_view const& source_table,
out_of_bounds_policy bounds_policy,
negative_index_policy neg_indices,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_EXPECTS(not gather_map.has_nulls(), "gather_map contains nulls", std::invalid_argument);
@@ -55,7 +55,7 @@ std::unique_ptr gather(table_view const& source_table,
out_of_bounds_policy bounds_policy,
negative_index_policy neg_indices,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_EXPECTS(gather_map.size() <= static_cast(std::numeric_limits::max()),
"gather map size exceeds the column size limit",
@@ -74,7 +74,7 @@ std::unique_ptr gather(table_view const& source_table,
column_view const& gather_map,
out_of_bounds_policy bounds_policy,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_FUNC_RANGE();
@@ -89,7 +89,7 @@ std::unique_ptr gather(table_view const& source_table,
out_of_bounds_policy bounds_policy,
negative_index_policy neg_indices,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_FUNC_RANGE();
return detail::gather(source_table, gather_map, bounds_policy, neg_indices, stream, mr);
diff --git a/cpp/src/dictionary/decode.cu b/cpp/src/dictionary/decode.cu
index 1f54b62843ff..d7d2adb185bc 100644
--- a/cpp/src/dictionary/decode.cu
+++ b/cpp/src/dictionary/decode.cu
@@ -36,13 +36,13 @@ struct indices_handler_fn {
*/
std::unique_ptr decode(dictionary_column_view const& source,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
if (source.is_empty()) return make_empty_column(type_id::EMPTY);
// annotated indices include the offset, size and bitmask from it's parent
auto const indices = source.get_indices_annotated();
- auto const d_indices = column_device_view::create(indices, stream);
+ auto const d_indices = column_device_view::create(indices, stream, mr.get_temporary_mr());
auto const d_iterator = cudf::detail::indexalator_factory::make_input_iterator(indices);
auto const indices_begin = cudf::detail::make_counting_transform_iterator(
0, indices_handler_fn{d_iterator, *d_indices, source.keys().size()});
@@ -57,8 +57,8 @@ std::unique_ptr decode(dictionary_column_view const& source,
auto output_column = std::unique_ptr(std::move(table_column.front()));
// apply any nulls to the output column
- output_column->set_null_mask(cudf::detail::copy_bitmask(source.parent(), stream, mr),
- source.null_count());
+ output_column->set_null_mask(
+ cudf::detail::copy_bitmask(source.parent(), stream, mr.get_output_mr()), source.null_count());
return output_column;
}
@@ -67,7 +67,7 @@ std::unique_ptr decode(dictionary_column_view const& source,
std::unique_ptr decode(dictionary_column_view const& source,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_FUNC_RANGE();
return detail::decode(source, stream, mr);
diff --git a/cpp/src/dictionary/detail/concatenate.cu b/cpp/src/dictionary/detail/concatenate.cu
index 1159c63852a9..6e1433432b4b 100644
--- a/cpp/src/dictionary/detail/concatenate.cu
+++ b/cpp/src/dictionary/detail/concatenate.cu
@@ -187,30 +187,29 @@ std::unique_ptr concatenate(host_span columns,
cudf::detail::row::hash::device_row_hasher>;
auto const tv = cudf::table_view({all_keys->view()});
- auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream);
- auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream);
+ 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{};
auto const d_equal =
row_equal.equal_to(cudf::nullate::NO{}, null_equality::EQUAL, comparator);
auto const empty_key = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
auto probe = encode_probe_t{row_hash.device_hasher(cudf::nullate::NO{})};
- auto allocator = rmm::mr::polymorphic_allocator(cudf::get_current_device_resource_ref());
- auto set = cuco::static_set{
+ auto allocator = rmm::mr::polymorphic_allocator(temp_mr);
+ auto set = cuco::static_set{
all_keys->size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
auto set_ref = set.ref(cuco::insert_and_find);
using set_ref_t = decltype(set_ref);
- auto policy = rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref());
+ auto policy = rmm::exec_policy_nosync(stream, temp_mr);
auto iota = cuda::counting_iterator{0};
- auto d_indices = rmm::device_uvector(
- all_keys->size(), stream, cudf::get_current_device_resource_ref());
+ auto d_indices = rmm::device_uvector(all_keys->size(), stream, temp_mr);
auto d_all_keys = column_device_view::create(all_keys->view(), stream);
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(
- all_keys->size(), stream, cudf::get_current_device_resource_ref());
- auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
+ auto keys_indices = rmm::device_uvector(all_keys->size(), stream, temp_mr);
+ auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
// use keys_indices to retrieve the keys (gather)
@@ -223,13 +222,11 @@ std::unique_ptr concatenate(host_span columns,
// build an all_keys_remap: abs position in all_keys to new key index
// use scatter to assign new index values: all_keys_remap[keys_indices[i]] = i
- auto all_keys_remap = rmm::device_uvector(
- all_keys->size(), stream, cudf::get_current_device_resource_ref());
+ auto all_keys_remap = rmm::device_uvector(all_keys->size(), stream, temp_mr);
thrust::scatter(
policy, iota, iota + keys_indices.size(), keys_indices.begin(), all_keys_remap.begin());
// use gather to propagate new indices values to all duplicate positions
- auto final_remap = rmm::device_uvector(
- all_keys->size(), stream, cudf::get_current_device_resource_ref());
+ auto final_remap = rmm::device_uvector(all_keys->size(), stream, temp_mr);
thrust::gather(
policy, d_indices.begin(), d_indices.end(), all_keys_remap.begin(), final_remap.begin());
diff --git a/cpp/src/dictionary/encode.cu b/cpp/src/dictionary/encode.cu
index b3dd4ce3c5af..35ac1b71c6d1 100644
--- a/cpp/src/dictionary/encode.cu
+++ b/cpp/src/dictionary/encode.cu
@@ -56,7 +56,7 @@ struct encode_fn {
std::unique_ptr encode(column_view const& input,
data_type indices_type,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_EXPECTS(is_signed(indices_type) && is_index_type(indices_type),
"indices must be type signed integer",
@@ -68,8 +68,11 @@ std::unique_ptr encode(column_view const& input,
"encoding nested types not supported",
std::invalid_argument);
+ auto const output_mr = mr.get_output_mr();
+ auto const temp_mr = mr.get_temporary_mr();
+
auto indices_column = cudf::make_numeric_column(
- indices_type, input.size(), cudf::mask_state::UNALLOCATED, stream, mr);
+ indices_type, input.size(), cudf::mask_state::UNALLOCATED, stream, output_mr);
if (input.is_empty()) {
return make_dictionary_column(
make_empty_column(input.type()), std::move(indices_column), rmm::device_buffer{}, 0);
@@ -82,13 +85,13 @@ std::unique_ptr encode(column_view const& input,
auto const has_nulls = nullate::DYNAMIC{input.has_nulls()};
auto const tv = cudf::table_view({input});
- auto const row_hash = cudf::detail::row::hash::row_hasher(tv, stream);
- auto const row_equal = cudf::detail::row::equality::self_comparator(tv, stream);
+ 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{};
auto const d_equal = row_equal.equal_to(has_nulls, null_equality::EQUAL, comparator);
auto const empty_key = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
auto probe = encode_probe_t{row_hash.device_hasher(has_nulls)};
- auto allocator = rmm::mr::polymorphic_allocator{};
+ auto allocator = rmm::mr::polymorphic_allocator{temp_mr};
auto set =
cuco::static_set{input.size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
auto set_ref = set.ref(cuco::insert_and_find);
@@ -96,22 +99,20 @@ std::unique_ptr encode(column_view const& input,
// build a static_set of the input values
// and keep track of the indices of the unique values
- auto d_indices = rmm::device_uvector(input.size(), stream);
- auto d_input = column_device_view::create(input, stream);
- thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ auto d_indices = rmm::device_uvector(input.size(), stream, temp_mr);
+ auto d_input = column_device_view::create(input, stream, temp_mr);
+ thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
cuda::counting_iterator{0},
cuda::counting_iterator{input.size()},
d_indices.begin(),
encode_fn{set_ref, *d_input});
- auto keys_indices = rmm::device_uvector(input.size(), stream);
+ auto keys_indices = rmm::device_uvector(input.size(), stream, temp_mr);
auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
// sort the keys_indices so we can use lower-bound on them
- thrust::sort(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
- keys_indices.begin(),
- keys_indices.end());
+ thrust::sort(rmm::exec_policy_nosync(stream, temp_mr), keys_indices.begin(), keys_indices.end());
// use keys_indices to retrieve the keys
auto const oob_policy = cudf::out_of_bounds_policy::DONT_CHECK;
@@ -124,7 +125,7 @@ std::unique_ptr encode(column_view const& input,
// call lower-bound with keys_indices and d_indices to get the output indices_column
auto d_result =
cudf::detail::indexalator_factory::make_output_iterator(indices_column->mutable_view());
- thrust::lower_bound(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::lower_bound(rmm::exec_policy_nosync(stream, temp_mr),
keys_indices.begin(),
keys_indices.end(),
d_indices.begin(),
@@ -134,7 +135,7 @@ std::unique_ptr encode(column_view const& input,
// create column with keys_column and indices_column
return make_dictionary_column(std::move(keys_column),
std::move(indices_column),
- cudf::detail::copy_bitmask(input, stream, mr),
+ cudf::detail::copy_bitmask(input, stream, output_mr),
input.null_count());
}
@@ -155,7 +156,7 @@ data_type get_indices_type_for_size(size_type keys_size)
std::unique_ptr encode(column_view const& input_column,
data_type indices_type,
cuda::stream_ref stream,
- rmm::device_async_resource_ref mr)
+ cudf::memory_resources mr)
{
CUDF_FUNC_RANGE();
return detail::encode(input_column, indices_type, stream, mr);
diff --git a/cpp/src/dictionary/match_keys.cu b/cpp/src/dictionary/match_keys.cu
index f86559ed7181..57dba5f5ebfd 100644
--- a/cpp/src/dictionary/match_keys.cu
+++ b/cpp/src/dictionary/match_keys.cu
@@ -50,13 +50,14 @@ struct unique_keys_dispatch_fn {
auto const has_nulls = nullate::DYNAMIC{false};
auto const keys_tv = table_view({all_keys});
- auto const row_hash = cudf::detail::row::hash::row_hasher(keys_tv, stream);
- auto const row_equal = cudf::detail::row::equality::self_comparator(keys_tv, stream);
+ auto const temp_mr = cudf::get_current_device_resource_ref();
+ auto const row_hash = cudf::detail::row::hash::row_hasher(keys_tv, stream, temp_mr);
+ auto const row_equal = cudf::detail::row::equality::self_comparator(keys_tv, stream, temp_mr);
auto const comparator = cudf::detail::row::equality::nan_equal_physical_equality_comparator{};
auto const d_equal = row_equal.equal_to(has_nulls, null_equality::EQUAL, comparator);
auto const empty_key = cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL};
auto probe = probe_t{row_hash.device_hasher(has_nulls)};
- auto allocator = rmm::mr::polymorphic_allocator{};
+ auto allocator = rmm::mr::polymorphic_allocator{temp_mr};
auto set = cuco::static_set{
all_keys.size(), 0.5, empty_key, d_equal, probe, {}, {}, allocator, stream.get()};
@@ -65,7 +66,7 @@ struct unique_keys_dispatch_fn {
set.insert_async(iter, iter + all_keys.size(), stream.get());
// retrieve the indices of all the unique keys
- auto keys_indices = rmm::device_uvector(all_keys.size(), stream);
+ auto keys_indices = rmm::device_uvector(all_keys.size(), stream, temp_mr);
auto keys_end = set.retrieve_all(keys_indices.begin(), stream.get());
keys_indices.resize(cuda::std::distance(keys_indices.begin(), keys_end), stream);
diff --git a/cpp/src/groupby/hash/groupby.cu b/cpp/src/groupby/hash/groupby.cu
index 8b9a778cba91..814b55300f64 100644
--- a/cpp/src/groupby/hash/groupby.cu
+++ b/cpp/src/groupby/hash/groupby.cu
@@ -42,10 +42,11 @@ std::unique_ptr dispatch_groupby(table_view const& keys,
auto const has_null = nullate::DYNAMIC{cudf::has_nested_nulls(keys)};
auto const skip_rows_with_nulls = keys_have_nulls and include_null_keys == null_policy::EXCLUDE;
- auto preprocessed_keys = cudf::detail::row::hash::preprocessed_table::create(keys, stream);
- auto const comparator = cudf::detail::row::equality::self_comparator{preprocessed_keys};
- auto const row_hash = cudf::detail::row::hash::row_hasher{std::move(preprocessed_keys)};
- auto const d_row_hash = row_hash.device_hasher(has_null);
+ auto preprocessed_keys = cudf::detail::row::hash::preprocessed_table::create(
+ keys, stream, cudf::get_current_device_resource_ref());
+ auto const comparator = cudf::detail::row::equality::self_comparator{preprocessed_keys};
+ auto const row_hash = cudf::detail::row::hash::row_hasher{std::move(preprocessed_keys)};
+ auto const d_row_hash = row_hash.device_hasher(has_null);
if (cudf::detail::has_nested_columns(keys)) {
auto const d_row_equal = comparator.equal_to(has_null, null_keys_are_equal);
diff --git a/cpp/src/groupby/sort/group_nunique.cu b/cpp/src/groupby/sort/group_nunique.cu
index e3a54734c6d9..386e164a2790 100644
--- a/cpp/src/groupby/sort/group_nunique.cu
+++ b/cpp/src/groupby/sort/group_nunique.cu
@@ -78,8 +78,10 @@ std::unique_ptr group_nunique(column_view const& values,
if (num_groups == 0) { return result; }
+ auto temp_mr = cudf::get_current_device_resource_ref();
auto const values_view = table_view{{values}};
- auto const comparator = cudf::detail::row::equality::self_comparator{values_view, stream};
+ auto const comparator =
+ cudf::detail::row::equality::self_comparator{values_view, stream, temp_mr};
auto const d_values_view = column_device_view::create(values, stream);
@@ -92,7 +94,7 @@ std::unique_ptr group_nunique(column_view const& values,
null_handling,
group_offsets.data(),
group_labels.data()};
- 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{values.size()},
d_result.begin(),
diff --git a/cpp/src/groupby/sort/group_rank_scan.cu b/cpp/src/groupby/sort/group_rank_scan.cu
index 6762942ae4f3..a9095df876c3 100644
--- a/cpp/src/groupby/sort/group_rank_scan.cu
+++ b/cpp/src/groupby/sort/group_rank_scan.cu
@@ -94,7 +94,9 @@ std::unique_ptr rank_generator(column_view const& grouped_values,
rmm::device_async_resource_ref mr)
{
auto const grouped_values_view = table_view{{grouped_values}};
- auto const comparator = cudf::detail::row::equality::self_comparator{grouped_values_view, stream};
+ auto const temp_mr = cudf::get_current_device_resource_ref();
+ auto const comparator =
+ cudf::detail::row::equality::self_comparator{grouped_values_view, stream, temp_mr};
auto ranks = make_fixed_width_column(
data_type{type_to_id()}, grouped_values.size(), mask_state::UNALLOCATED, stream, mr);
@@ -104,7 +106,7 @@ std::unique_ptr rank_generator(column_view const& grouped_values,
auto const permuted_equal =
permuted_row_equality_comparator(d_equal, value_order.begin());
- 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(grouped_values.size()),
mutable_ranks.begin(),
@@ -130,14 +132,13 @@ std::unique_ptr rank_generator(column_view const& grouped_values,
cuda::std::reverse_iterator(mutable_ranks.end())};
}
}();
- thrust::inclusive_scan_by_key(
- rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
- group_labels_begin,
- group_labels_begin + group_labels.size(),
- mutable_rank_begin,
- mutable_rank_begin,
- cuda::std::equal_to{},
- scan_op);
+ thrust::inclusive_scan_by_key(rmm::exec_policy_nosync(stream, temp_mr),
+ group_labels_begin,
+ group_labels_begin + group_labels.size(),
+ mutable_rank_begin,
+ mutable_rank_begin,
+ cuda::std::equal_to{},
+ scan_op);
return ranks;
}
} // namespace
diff --git a/cpp/src/groupby/sort/sort_helper_group_offsets.cuh b/cpp/src/groupby/sort/sort_helper_group_offsets.cuh
index dd2e32d54165..c74a8f76fe78 100644
--- a/cpp/src/groupby/sort/sort_helper_group_offsets.cuh
+++ b/cpp/src/groupby/sort/sort_helper_group_offsets.cuh
@@ -38,22 +38,19 @@ size_type compute_group_offsets(table_view const& keys,
rmm::device_uvector& group_offsets,
cuda::stream_ref stream)
{
- auto const comparator = cudf::detail::row::equality::self_comparator{keys, stream};
+ auto const temp_mr = cudf::get_current_device_resource_ref();
+ auto const comparator = cudf::detail::row::equality::self_comparator{keys, stream, temp_mr};
auto const d_key_equal = comparator.equal_to(
cudf::nullate::DYNAMIC{cudf::has_nested_nulls(keys)}, null_equality::EQUAL);
// Using a temporary buffer for intermediate transform results from the iterator containing
// the comparator speeds up compile-time significantly without much degradation in
// runtime performance over using the comparator directly in thrust::unique_copy.
- auto result = rmm::device_uvector(size, stream);
+ auto result = rmm::device_uvector(size, stream, temp_mr);
auto const itr = cuda::counting_iterator{0};
auto const row_eq = permuted_row_equality_comparator(d_key_equal, sorted_order);
auto const ufn = cudf::detail::unique_copy_fn{
itr, duplicate_keep_option::KEEP_FIRST, row_eq, size - 1};
- thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
- itr,
- itr + size,
- result.begin(),
- ufn);
+ thrust::transform(rmm::exec_policy_nosync(stream, temp_mr), itr, itr + size, result.begin(), ufn);
auto const result_end = cudf::detail::copy_if(
itr, itr + size, result.begin(), group_offsets.begin(), cuda::std::identity{}, stream);
return cuda::std::distance(group_offsets.begin(), result_end);
diff --git a/cpp/src/groupby/streaming_groupby/insert.cuh b/cpp/src/groupby/streaming_groupby/insert.cuh
index 90be7c533ae9..9df90ecfcb93 100644
--- a/cpp/src/groupby/streaming_groupby/insert.cuh
+++ b/cpp/src/groupby/streaming_groupby/insert.cuh
@@ -34,7 +34,8 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
auto const has_null = cudf::nullate::DYNAMIC{_has_nullable_keys};
// Preprocess batch for row operators.
- auto preprocessed_batch = cudf::detail::row::hash::preprocessed_table::create(batch_keys, stream);
+ auto preprocessed_batch =
+ cudf::detail::row::hash::preprocessed_table::create(batch_keys, stream, temp_mr);
auto const batch_hasher_obj = cudf::detail::row::hash::row_hasher{preprocessed_batch};
auto const d_batch_hash = batch_hasher_obj.device_hasher(has_null);
@@ -110,7 +111,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_
temp_mr);
auto preprocessed_compacted =
- cudf::detail::row::hash::preprocessed_table::create(compacted->view(), stream);
+ cudf::detail::row::hash::preprocessed_table::create(compacted->view(), stream, temp_mr);
// Store the compacted batch.
auto const new_batch_id = static_cast(_compacted_batches.size());
diff --git a/cpp/src/hash/murmurhash3_x86_32.cu b/cpp/src/hash/murmurhash3_x86_32.cu
index f82d552456d7..4e1f5ced36d5 100644
--- a/cpp/src/hash/murmurhash3_x86_32.cu
+++ b/cpp/src/hash/murmurhash3_x86_32.cu
@@ -15,6 +15,7 @@
#include
#include
#include
+#include
#include
#include
@@ -66,8 +67,8 @@ std::unique_ptr murmurhash3_x86_32(table_view const& input,
rmm::cuda_stream_view stream,
rmm::device_async_resource_ref mr)
{
- auto const preprocessed_input =
- cudf::detail::row::hash::preprocessed_table::create(input, stream);
+ auto const preprocessed_input = cudf::detail::row::hash::preprocessed_table::create(
+ input, stream, cudf::get_current_device_resource_ref());
return murmurhash3_x86_32_impl(
preprocessed_input, input.num_rows(), seed, nullate::DYNAMIC{has_nulls(input)}, stream, mr);
}
diff --git a/cpp/src/hash/xxhash_32.cu b/cpp/src/hash/xxhash_32.cu
index 759a491b193b..0efd6f02ce12 100644
--- a/cpp/src/hash/xxhash_32.cu
+++ b/cpp/src/hash/xxhash_32.cu
@@ -30,9 +30,10 @@ std::unique_ptr xxhash_32(table_view const& input,
if (input.num_rows() == 0) { return output; }
- bool const nullable = has_nulls(input);
- auto const row_hasher = cudf::detail::row::hash::row_hasher(input, stream);
- auto output_view = output->mutable_view();
+ bool const nullable = has_nulls(input);
+ auto const row_hasher =
+ cudf::detail::row::hash::row_hasher(input, stream, cudf::get_current_device_resource_ref());
+ auto output_view = output->mutable_view();
// Compute the hash value for each row
auto const output_begin = output_view.begin();
diff --git a/cpp/src/hash/xxhash_64.cu b/cpp/src/hash/xxhash_64.cu
index fcf7009bd128..69c97724bbd6 100644
--- a/cpp/src/hash/xxhash_64.cu
+++ b/cpp/src/hash/xxhash_64.cu
@@ -32,9 +32,10 @@ std::unique_ptr xxhash_64(table_view const& input,
if (input.num_rows() == 0) { return output; }
- bool const nullable = has_nulls(input);
- auto const row_hasher = cudf::detail::row::hash::row_hasher(input, stream);
- auto output_view = output->mutable_view();
+ bool const nullable = has_nulls(input);
+ auto const row_hasher =
+ cudf::detail::row::hash::row_hasher(input, stream, cudf::get_current_device_resource_ref());
+ auto output_view = output->mutable_view();
// Compute the hash value for each row
auto const output_begin = output_view.begin();
diff --git a/cpp/src/join/distinct_hash_join.cu b/cpp/src/join/distinct_hash_join.cu
index b7d735cfb78d..f18230041ab4 100644
--- a/cpp/src/join/distinct_hash_join.cu
+++ b/cpp/src/join/distinct_hash_join.cu
@@ -157,7 +157,8 @@ distinct_hash_join::distinct_hash_join(cudf::table_view const& right,
: _has_nested_columns{cudf::has_nested_columns(right)},
_nulls_equal{compare_nulls},
_right{right},
- _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)},
+ _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(
+ _right, stream, cudf::get_current_device_resource_ref())},
_hash_table{cuco::extent{static_cast(right.num_rows())},
checked_load_factor(load_factor),
cuco::empty_key{cuco::pair{std::numeric_limits::max(),
@@ -234,7 +235,8 @@ distinct_hash_join::inner_join(cudf::table_view const& left,
auto found_indices = rmm::device_uvector(left_table_num_rows, stream);
auto const found_begin = cuda::make_transform_output_iterator(found_indices.begin(), output_fn{});
- auto preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ auto preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
if (cudf::detail::is_primitive_row_op_compatible(_right)) {
auto const d_hasher =
cudf::detail::row::primitive::row_hasher{nullate::DYNAMIC{has_nulls}, preprocessed_left};
@@ -325,7 +327,8 @@ std::unique_ptr> distinct_hash_join::left_join(
auto const output_begin =
cuda::make_transform_output_iterator(right_indices->begin(), output_fn{});
- auto preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ auto preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
if (cudf::detail::is_primitive_row_op_compatible(_right)) {
auto const d_hasher =
diff --git a/cpp/src/join/filtered_join/filtered_join.cu b/cpp/src/join/filtered_join/filtered_join.cu
index 25f7172ffa2d..c0da2b23b0c3 100644
--- a/cpp/src/join/filtered_join/filtered_join.cu
+++ b/cpp/src/join/filtered_join/filtered_join.cu
@@ -107,7 +107,8 @@ filtered_join::filtered_join(cudf::table_view const& right,
stream.get()},
_right{right},
_nulls_equal{compare_nulls},
- _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)}
+ _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(
+ _right, stream, cudf::get_current_device_resource_ref())}
{
cudf::scoped_range range{"filtered_join::filtered_join"};
if (_right.num_rows() == 0) return;
@@ -129,9 +130,10 @@ std::unique_ptr> filtered_join::semi_anti_j
{
cudf::scoped_range range{"filtered_join::semi_anti_join"};
- auto const preprocessed_left = [&left, stream] {
+ auto const preprocessed_left = [left, stream] {
cudf::scoped_range range{"filtered_join::semi_anti_join::preprocessed_left"};
- return cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ return cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
}();
auto contains_map = rmm::device_uvector(left.num_rows(), stream);
diff --git a/cpp/src/join/hash_join/hash_join.cu b/cpp/src/join/hash_join/hash_join.cu
index 89d990bec73b..9808fa4f2f43 100644
--- a/cpp/src/join/hash_join/hash_join.cu
+++ b/cpp/src/join/hash_join/hash_join.cu
@@ -130,7 +130,8 @@ hash_join::hash_join(cudf::table_view const& right,
rmm::mr::polymorphic_allocator{std::move(mr)},
stream.get()}})},
_right{right},
- _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(_right, stream)}
+ _preprocessed_right{cudf::detail::row::equality::preprocessed_table::create(
+ _right, stream, cudf::get_current_device_resource_ref())}
{
CUDF_FUNC_RANGE();
CUDF_EXPECTS(0 != right.num_columns(), "Hash join right table is empty", std::invalid_argument);
diff --git a/cpp/src/join/hash_join/match_context.cu b/cpp/src/join/hash_join/match_context.cu
index e6812f55b3be..767f8f607b23 100644
--- a/cpp/src/join/hash_join/match_context.cu
+++ b/cpp/src/join/hash_join/match_context.cu
@@ -45,8 +45,8 @@ std::unique_ptr> make_join_match_counts(
"Left table has nulls while right table was not hashed with null check.",
std::invalid_argument);
- auto const preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
auto const left_table_num_rows = left.num_rows();
auto count_matches = [&](auto equality, auto d_hasher) {
diff --git a/cpp/src/join/hash_join/partitioned_join_retrieve.cu b/cpp/src/join/hash_join/partitioned_join_retrieve.cu
index 7c64b9464c55..f0c76746f588 100644
--- a/cpp/src/join/hash_join/partitioned_join_retrieve.cu
+++ b/cpp/src/join/hash_join/partitioned_join_retrieve.cu
@@ -106,8 +106,9 @@ hash_join::partitioned_join_retrieve(join_kind join,
validate_hash_join_probe(_right, left_partition_view, _has_nulls);
+ auto const temp_mr = cudf::get_current_device_resource_ref();
auto const preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left_partition_view, stream);
+ cudf::detail::row::equality::preprocessed_table::create(left_partition_view, stream, temp_mr);
// For FULL_JOIN, probe with LEFT_JOIN semantics (no complement here)
bool const is_outer = (join != join_kind::INNER_JOIN);
@@ -123,8 +124,8 @@ hash_join::partitioned_join_retrieve(join_kind join,
auto retrieve_partition = [&](auto equality, auto d_hasher) {
// Precompute left keys for this partition slice.
- rmm::device_uvector left_keys(n, stream);
- thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ rmm::device_uvector left_keys(n, stream, temp_mr);
+ thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
cuda::counting_iterator(0),
cuda::counting_iterator(partition_size),
left_keys.begin(),
diff --git a/cpp/src/join/hash_join/retrieve_impl.cuh b/cpp/src/join/hash_join/retrieve_impl.cuh
index e19b8e8e0c40..866ae34924ae 100644
--- a/cpp/src/join/hash_join/retrieve_impl.cuh
+++ b/cpp/src/join/hash_join/retrieve_impl.cuh
@@ -173,8 +173,8 @@ hash_join::join_retrieve(cudf::table_view const& left,
}
}
- auto const preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
auto join_indices = cudf::detail::probe_join_hash_table(_right,
left,
diff --git a/cpp/src/join/hash_join/size_impl.cuh b/cpp/src/join/hash_join/size_impl.cuh
index 3b1c44672fcc..ec3eba56a920 100644
--- a/cpp/src/join/hash_join/size_impl.cuh
+++ b/cpp/src/join/hash_join/size_impl.cuh
@@ -80,8 +80,8 @@ std::size_t hash_join::join_size(cudf::table_view const& left,
"Left table has nulls while right table was not hashed with null check.",
std::invalid_argument);
- auto const preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
return cudf::detail::compute_join_output_size(_right,
left,
@@ -109,8 +109,8 @@ std::size_t hash_join::join_size(cudf::table_view const& left,
"Left table has nulls while right table was not hashed with null check.",
std::invalid_argument);
- auto const preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ auto const preprocessed_left = cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref());
return cudf::detail::get_full_join_size(_right,
left,
diff --git a/cpp/src/join/key_remapping.cu b/cpp/src/join/key_remapping.cu
index 8191ea164e08..8057fc35d564 100644
--- a/cpp/src/join/key_remapping.cu
+++ b/cpp/src/join/key_remapping.cu
@@ -380,10 +380,12 @@ class key_remap_table : public key_remap_table_interface {
return std::make_unique>(0, stream, mr);
}
+ auto const temp_mr = cudf::get_current_device_resource_ref();
+
if (this->_right.num_rows() == 0) {
auto result =
std::make_unique>(left_num_rows, stream, mr);
- thrust::fill(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::fill(rmm::exec_policy_nosync(stream, temp_mr),
result->begin(),
result->end(),
cudf::JoinNoMatch);
@@ -395,7 +397,7 @@ class key_remap_table : public key_remap_table_interface {
cuda::make_transform_output_iterator(result->begin(), extract_index{});
auto preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left_keys, stream);
+ cudf::detail::row::equality::preprocessed_table::create(left_keys, stream, temp_mr);
if (cudf::detail::is_primitive_row_op_compatible(_right)) {
auto const d_hasher = cudf::detail::row::primitive::row_hasher{
@@ -500,7 +502,8 @@ std::unique_ptr create_key_remap_table(
if (right.num_rows() == 0 || right.num_columns() == 0) { return nullptr; }
- auto preprocessed_right = cudf::detail::row::equality::preprocessed_table::create(right, stream);
+ auto preprocessed_right = cudf::detail::row::equality::preprocessed_table::create(
+ right, stream, cudf::get_current_device_resource_ref());
if (cudf::detail::is_primitive_row_op_compatible(right)) {
auto const d_hasher = cudf::detail::row::primitive::row_hasher{
diff --git a/cpp/src/join/mark_join.cu b/cpp/src/join/mark_join.cu
index 09d92710763c..aef80b25c7ed 100644
--- a/cpp/src/join/mark_join.cu
+++ b/cpp/src/join/mark_join.cu
@@ -595,7 +595,8 @@ mark_join::mark_join(cudf::table_view const& left,
_left{left},
_nulls_equal{compare_nulls},
_prefilter{prefilter},
- _preprocessed_left{cudf::detail::row::equality::preprocessed_table::create(left, stream)},
+ _preprocessed_left{cudf::detail::row::equality::preprocessed_table::create(
+ left, stream, cudf::get_current_device_resource_ref())},
_bucket_storage{
cuco::extent{compute_mark_join_capacity(left, checked_load_factor(load_factor))},
rmm::mr::polymorphic_allocator{mr},
@@ -740,9 +741,10 @@ std::unique_ptr> mark_join::semi_anti_join(
{
clear_marks(stream);
- auto const preprocessed_right = [&right, stream] {
+ auto const preprocessed_right = [right, stream] {
cudf::scoped_range range{"mark_join::semi_anti_join::preprocessed_right"};
- return cudf::detail::row::equality::preprocessed_table::create(right, stream);
+ return cudf::detail::row::equality::preprocessed_table::create(
+ right, stream, cudf::get_current_device_resource_ref());
}();
if (is_primitive_row_op_compatible(_left)) {
diff --git a/cpp/src/join/mixed_join_semi.cu b/cpp/src/join/mixed_join_semi.cu
index 956bd53423b6..6d02b5cb9dfa 100644
--- a/cpp/src/join/mixed_join_semi.cu
+++ b/cpp/src/join/mixed_join_semi.cu
@@ -105,10 +105,11 @@ std::unique_ptr> mixed_join_semi(
auto left_conditional_view = table_device_view::create(left_conditional, stream);
auto right_conditional_view = table_device_view::create(right_conditional, stream);
+ auto const temp_mr = cudf::get_current_device_resource_ref();
auto const preprocessed_right =
- cudf::detail::row::equality::preprocessed_table::create(right, stream);
+ cudf::detail::row::equality::preprocessed_table::create(right, stream, temp_mr);
auto const preprocessed_left =
- cudf::detail::row::equality::preprocessed_table::create(left, stream);
+ cudf::detail::row::equality::preprocessed_table::create(left, stream, temp_mr);
auto const row_comparator =
cudf::detail::row::equality::two_table_comparator{preprocessed_left, preprocessed_right};
auto const equality_left = row_comparator.equal_to(has_nulls, compare_nulls);
@@ -134,7 +135,7 @@ std::unique_ptr> mixed_join_semi(
auto const equality_right_equality =
row_comparator_right.equal_to(right_nulls, compare_nulls);
auto const preprocessed_right_condtional =
- cudf::detail::row::equality::preprocessed_table::create(right_conditional, stream);
+ cudf::detail::row::equality::preprocessed_table::create(right_conditional, stream, temp_mr);
auto const row_comparator_conditional_right = cudf::detail::row::equality::two_table_comparator{
preprocessed_right_condtional, preprocessed_right_condtional};
auto const equality_right_conditional =
diff --git a/cpp/src/lists/contains.cu b/cpp/src/lists/contains.cu
index 58251a141111..6dc114ac0468 100644
--- a/cpp/src/lists/contains.cu
+++ b/cpp/src/lists/contains.cu
@@ -210,8 +210,8 @@ std::unique_ptr dispatch_index_of(lists_column_view const& lists,
auto const keys_tview = cudf::table_view{{search_keys}};
auto const child_tview = cudf::table_view{{child}};
auto const has_nulls = has_nested_nulls(child_tview) || has_nested_nulls(keys_tview);
- auto const comparator =
- cudf::detail::row::equality::two_table_comparator(child_tview, keys_tview, stream);
+ auto const comparator = cudf::detail::row::equality::two_table_comparator(
+ child_tview, keys_tview, stream, cudf::get_current_device_resource_ref());
if (cudf::is_nested(search_keys.type())) {
auto const d_comp = comparator.equal_to(nullate::DYNAMIC{has_nulls});
index_of(input_it, num_rows, output_it, child, search_keys, find_option, d_comp, stream);
diff --git a/cpp/src/partitioning/partitioning.cu b/cpp/src/partitioning/partitioning.cu
index 3afab1d7c5db..f7feca6edc45 100644
--- a/cpp/src/partitioning/partitioning.cu
+++ b/cpp/src/partitioning/partitioning.cu
@@ -575,8 +575,9 @@ std::pair, std::vector> hash_partition_table(
rmm::device_async_resource_ref mr)
{
auto const num_rows = table_to_hash.num_rows();
+ auto const temp_mr = cudf::get_current_device_resource_ref();
- auto const row_hasher = detail::row::hash::row_hasher(table_to_hash, stream);
+ auto const row_hasher = detail::row::hash::row_hasher(table_to_hash, stream, temp_mr);
auto const hasher =
row_hasher.device_hasher(nullate::DYNAMIC{hash_has_nulls}, seed);
@@ -600,7 +601,7 @@ std::pair, std::vector> hash_partition_table(
std::size_t const grid_size = util::div_rounding_up_safe(num_rows, rows_per_block);
// Allocate array to hold which partition each row belongs to
- auto row_partition_numbers = rmm::device_uvector(num_rows, stream);
+ auto row_partition_numbers = rmm::device_uvector(num_rows, stream, temp_mr);
// Array to hold the size of each partition computed by each block
// i.e., { {block0 partition0 size, block1 partition0 size, ...},
@@ -608,17 +609,18 @@ std::pair, std::vector> hash_partition_table(
// ...
// {block0 partition(num_partitions-1) size, block1
// partition(num_partitions -1) size, ...} }
- auto block_partition_sizes = rmm::device_uvector(grid_size * num_partitions, stream);
+ auto block_partition_sizes =
+ rmm::device_uvector(grid_size * num_partitions, stream, temp_mr);
auto scanned_block_partition_sizes =
- rmm::device_uvector(grid_size * num_partitions, stream);
+ rmm::device_uvector(grid_size * num_partitions, stream, temp_mr);
// Holds the total number of rows in each partition
- auto global_partition_sizes = cudf::detail::make_zeroed_device_uvector_async(
- num_partitions, stream, cudf::get_current_device_resource_ref());
+ auto global_partition_sizes =
+ cudf::detail::make_zeroed_device_uvector_async(num_partitions, stream, temp_mr);
- auto row_partition_offset = cudf::detail::make_zeroed_device_uvector_async(
- num_rows, stream, cudf::get_current_device_resource_ref());
+ auto row_partition_offset =
+ cudf::detail::make_zeroed_device_uvector_async(num_rows, stream, temp_mr);
// If the number of partitions is a power of two, we can compute the partition
// number of each row more efficiently with bitwise operations
@@ -668,7 +670,7 @@ std::pair, std::vector> hash_partition_table(
// Compute exclusive scan of all blocks' partition sizes in-place to determine
// the starting point for each blocks portion of each partition in the output
- thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::exclusive_scan(rmm::exec_policy_nosync(stream, temp_mr),
block_partition_sizes.begin(),
block_partition_sizes.end(),
scanned_block_partition_sizes.data());
@@ -676,7 +678,7 @@ std::pair, std::vector> hash_partition_table(
// Compute exclusive scan of size of each partition to determine offset
// location of each partition in final output.
// TODO This can be done independently on a separate stream
- thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ thrust::exclusive_scan(rmm::exec_policy_nosync(stream, temp_mr),
global_partition_sizes.begin(),
global_partition_sizes.end(),
global_partition_sizes.begin());
diff --git a/cpp/src/reductions/approx_distinct_count.cu b/cpp/src/reductions/approx_distinct_count.cu
index 07acf1d205f6..299d39d764e1 100644
--- a/cpp/src/reductions/approx_distinct_count.cu
+++ b/cpp/src/reductions/approx_distinct_count.cu
@@ -222,8 +222,9 @@ void approx_distinct_count::add(table_view const& input, cuda::stream_re
typename approx_distinct_count::hll_ref_type ref{sketch(), cuda::std::identity{}};
auto const has_nulls = nullate::DYNAMIC{cudf::has_nested_nulls(input)};
+ auto const temp_mr = cudf::get_current_device_resource_ref();
auto const preprocessed_input =
- cudf::detail::row::hash::preprocessed_table::create(input, stream);
+ cudf::detail::row::hash::preprocessed_table::create(input, stream, temp_mr);
auto const row_hasher = cudf::detail::row::hash::row_hasher(preprocessed_input);
auto const hash_key = row_hasher.device_hasher(has_nulls);
@@ -245,9 +246,8 @@ void approx_distinct_count::add(table_view const& input, cuda::stream_re
if (!has_nulls) {
ref.add_async(hash_iter, hash_iter + num_rows, stream);
} else {
- auto const row_bitmask =
- cudf::detail::bitmask_and(input, stream, cudf::get_current_device_resource_ref()).first;
- auto const pred = row_is_valid{static_cast(row_bitmask.data())};
+ auto const row_bitmask = cudf::detail::bitmask_and(input, stream, temp_mr).first;
+ auto const pred = row_is_valid{static_cast(row_bitmask.data())};
ref.add_if_async(hash_iter, hash_iter + num_rows, stencil, pred, stream);
}
} else {
@@ -256,8 +256,7 @@ void approx_distinct_count::add(table_view const& input, cuda::stream_re
auto const pred = check_nans_predicate{*d_table, nullptr};
ref.add_if_async(hash_iter, hash_iter + num_rows, stencil, pred, stream);
} else {
- auto const row_bitmask =
- cudf::detail::bitmask_and(input, stream, cudf::get_current_device_resource_ref()).first;
+ auto const row_bitmask = cudf::detail::bitmask_and(input, stream, temp_mr).first;
auto const bitmask_ptr = static_cast(row_bitmask.data());
auto const pred = check_nans_predicate{*d_table, bitmask_ptr};
ref.add_if_async(hash_iter, hash_iter + num_rows, stencil, pred, stream);
diff --git a/cpp/src/reductions/distinct_count.cu b/cpp/src/reductions/distinct_count.cu
index 0a19155f7e2c..a4a8cf5c9481 100644
--- a/cpp/src/reductions/distinct_count.cu
+++ b/cpp/src/reductions/distinct_count.cu
@@ -135,11 +135,13 @@ cudf::size_type distinct_count(table_view const& keys,
auto const num_rows = keys.num_rows();
if (num_rows == 0) { return 0; } // early exit for empty input
auto const has_nulls = nullate::DYNAMIC{cudf::has_nested_nulls(keys)};
+ auto const temp_mr = cudf::get_current_device_resource_ref();
- auto const preprocessed_input = cudf::detail::row::hash::preprocessed_table::create(keys, stream);
- auto const row_hasher = cudf::detail::row::hash::row_hasher(preprocessed_input);
- auto const hash_key = row_hasher.device_hasher(has_nulls);
- auto const row_comp = cudf::detail::row::equality::self_comparator(preprocessed_input);
+ auto const preprocessed_input =
+ cudf::detail::row::hash::preprocessed_table::create(keys, stream, temp_mr);
+ auto const row_hasher = cudf::detail::row::hash::row_hasher(preprocessed_input);
+ auto const hash_key = row_hasher.device_hasher(has_nulls);
+ auto const row_comp = cudf::detail::row::equality::self_comparator(preprocessed_input);
auto const comparator_helper = [&](auto const row_equal) {
using hasher_type = decltype(hash_key);
@@ -150,7 +152,7 @@ cudf::size_type distinct_count(table_view const& keys,
cuco::linear_probing<1, hasher_type>{hash_key},
{},
{},
- rmm::mr::polymorphic_allocator{},
+ rmm::mr::polymorphic_allocator{temp_mr},
stream.get()};
auto const iter = cuda::counting_iterator{0};
@@ -160,8 +162,7 @@ cudf::size_type distinct_count(table_view const& keys,
cuda::counting_iterator stencil(0);
// We must consider a row if any of its column entries is valid,
// hence OR together the validities of the columns.
- auto const [row_bitmask, null_count] =
- cudf::detail::bitmask_or(keys, stream, cudf::get_current_device_resource_ref());
+ auto const [row_bitmask, null_count] = cudf::detail::bitmask_or(keys, stream, temp_mr);
// Unless all columns have a null mask, row_bitmask will be
// null, and null_count will be zero. Equally, unless there is
diff --git a/cpp/src/reductions/histogram.cu b/cpp/src/reductions/histogram.cu
index b7bcd401463b..380f21e8f65e 100644
--- a/cpp/src/reductions/histogram.cu
+++ b/cpp/src/reductions/histogram.cu
@@ -111,8 +111,9 @@ compute_row_frequencies(table_view const& input,
"Nested types are not yet supported in histogram aggregation.",
std::invalid_argument);
+ auto const temp_mr = cudf::get_current_device_resource_ref();
auto const preprocessed_input =
- cudf::detail::row::hash::preprocessed_table::create(input, stream);
+ cudf::detail::row::hash::preprocessed_table::create(input, stream, temp_mr);
auto const has_nulls = nullate::DYNAMIC{cudf::has_nested_nulls(input)};
auto const row_hasher = cudf::detail::row::hash::row_hasher(preprocessed_input);
@@ -132,11 +133,10 @@ compute_row_frequencies(table_view const& input,
// Construct a vector to store reduced counts and init to zero
rmm::device_uvector reduction_results(num_rows, stream, mr);
- thrust::uninitialized_fill(
- rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
- reduction_results.begin(),
- reduction_results.end(),
- histogram_count_type{0});
+ thrust::uninitialized_fill(rmm::exec_policy_nosync(stream, temp_mr),
+ reduction_results.begin(),
+ reduction_results.end(),
+ histogram_count_type{0});
// Construct a hash set
auto row_set =
@@ -147,7 +147,7 @@ compute_row_frequencies(table_view const& input,
cuco::linear_probing{key_hasher},
{}, // thread scope
{}, // storage
- rmm::mr::polymorphic_allocator{},
+ rmm::mr::polymorphic_allocator{temp_mr},
stream.get()};
// Device-accessible reference to the hash set with `insert_and_find` operator
@@ -156,7 +156,7 @@ compute_row_frequencies(table_view const& input,
// Compute frequencies (aka distinct counts) for the input rows.
// Note that we consider null and NaNs as always equal.
thrust::for_each(
- rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
+ rmm::exec_policy_nosync(stream, temp_mr),
cuda::counting_iterator{0},
cuda::counting_iterator