Route dictionary encode temp allocations through memory_resources

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<table> 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<table> 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<table> 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<table> 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<table> 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.

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<table> 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<table> 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

cpp/include/cudf/detail/sizes_to_offsets_iterator.cuh

@@ -249,27 +249,29 @@ static sizes_to_offsets_iterator<ScanIterator, LastType> 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 <typename SizesIterator, typename OffsetsIterator>
-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<SizesIterator>;
   static_assert(std::is_integral_v<SizeType>,
                 "Only numeric types are supported by sizes_to_offsets");
 
-  using LastType = std::conditional_t<std::is_signed_v<SizeType>, int64_t, uint64_t>;
-  auto last_element =
-    cudf::detail::device_scalar<LastType>(0, stream, cudf::get_current_device_resource_ref());
+  using LastType    = std::conditional_t<std::is_signed_v<SizeType>, int64_t, uint64_t>;
+  auto last_element = cudf::detail::device_scalar<LastType>(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 <typename InputIterator>
 std::pair<std::unique_ptr<column>, 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<size_type>(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<int32_t>();
 
@@ -327,7 +326,7 @@ std::pair<std::unique_ptr<column>, 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<decltype(total_elements)>(std::numeric_limits<int32_t>::max()),

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<column> 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<column> 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<column> 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.

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<column> 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<column> 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<column> 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<column> 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<column> 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

cpp/include/cudf/strings/detail/gather.cuh

@@ -215,21 +215,24 @@ 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 <bool NullifyOutOfBounds, typename MapIterator>
 std::unique_ptr<cudf::column> gather(strings_column_view const& strings,
                                      MapIterator begin,
                                      MapIterator end,
                                      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();
+
   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<cudf::column> 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<char>(out_char_bytes, stream, mr);
+  auto out_chars_data = rmm::device_uvector<char>(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<cudf::column> gather(strings_column_view const& strings,
                                  stream.get());
 
       // 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<cudf::column> 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 <typename MapIterator>
@@ -367,7 +370,7 @@ std::unique_ptr<cudf::column> gather(strings_column_view const& strings,
                                      MapIterator end,
                                      bool nullify_out_of_bounds,
                                      cuda::stream_ref stream,
-                                     rmm::device_async_resource_ref mr)
+                                     cudf::memory_resources mr)
 {
   if (nullify_out_of_bounds) return gather<true>(strings, begin, end, stream, mr);
   return gather<false>(strings, begin, end, stream, mr);

cpp/include/cudf/strings/detail/strings_children.cuh

@@ -115,23 +115,25 @@ rmm::device_uvector<char> 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 <typename InputIterator>
-std::pair<std::unique_ptr<column>, int64_t> make_offsets_child_column(
-  InputIterator begin,
-  InputIterator end,
-  cuda::stream_ref stream,
-  rmm::device_async_resource_ref mr)
+std::pair<std::unique_ptr<column>, int64_t> make_offsets_child_column(InputIterator begin,
+                                                                      InputIterator end,
+                                                                      cuda::stream_ref 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<int64_t>(std::numeric_limits<size_type>::max());
   auto const lcount            = static_cast<int64_t>(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<size_type>(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<int32_t>();
 
   // The number of offsets is strings_count+1 so to build the offsets from the sizes
@@ -141,8 +143,8 @@ std::pair<std::unique_ptr<column>, 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<std::unique_ptr<column>, 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<int64_t>();
     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);