[cudf] 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/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

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))
   {
   }
 

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

@@ -10,6 +10,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <memory>
 #include <vector>
@@ -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<preprocessed_table> 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.

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,
                                      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<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.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<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,
                                      rmm::cuda_stream_view 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);