[libcudf] Add memory_resources support for gather

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.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 <typename InputItr, typename OutputItr, typename MapIterator>
 void gather_helper(InputItr source_itr,
@@ -112,19 +113,20 @@ 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<MapIterator>::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,
                       source_itr,
                       target_itr,
                       bounds_checker<map_type>{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,15 +161,15 @@ 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 <typename Element, typename MapIterator>
   std::unique_ptr<column> operator()(column_view const& source_column,
                                      MapIterator gather_map_begin,
                                      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<Element> gatherer{};
 
@@ -199,27 +201,29 @@ struct column_gatherer_impl<Element, std::enable_if_t<is_rep_layout_compatible<E
    * @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 <typename MapIterator>
   std::unique_ptr<column> operator()(column_view const& source_column,
                                      MapIterator gather_map_begin,
                                      MapIterator 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 num_rows     = cudf::distance(gather_map_begin, gather_map_end);
-    auto const policy       = cudf::mask_allocation_policy::NEVER;
-    auto destination_column = cudf::allocate_like(source_column, num_rows, policy, stream, mr);
+    auto const num_rows = cudf::distance(gather_map_begin, gather_map_end);
+    auto const policy   = cudf::mask_allocation_policy::NEVER;
+    auto destination_column =
+      cudf::allocate_like(source_column, num_rows, policy, stream, mr.get_output_mr());
 
     gather_helper(source_column.data<Element>(),
                   source_column.size(),
                   destination_column->mutable_view().template begin<Element>(),
                   gather_map_begin,
                   gather_map_end,
                   nullify_out_of_bounds,
-                  stream);
+                  stream,
+                  mr.get_temporary_mr());
 
     return destination_column;
   }
@@ -244,15 +248,15 @@ struct column_gatherer_impl<string_view> {
    * @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 <typename MapItType>
   std::unique_ptr<column> operator()(column_view const& source_column,
                                      MapItType gather_map_begin,
                                      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<true>(
@@ -326,8 +330,10 @@ struct column_gatherer_impl<list_view> {
                                      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
@@ -350,7 +356,7 @@ struct column_gatherer_impl<list_view> {
                                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
@@ -361,7 +367,7 @@ struct column_gatherer_impl<list_view> {
                              std::move(gd.offsets),
                              std::move(child),
                              0,
-                             rmm::device_buffer{0, stream, mr});
+                             rmm::device_buffer{0, stream, output_mr});
   }
 };
 
@@ -380,7 +386,7 @@ struct column_gatherer_impl<dictionary32> {
    * @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 <typename MapItType>
@@ -389,8 +395,11 @@ struct column_gatherer_impl<dictionary32> {
                                      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,20 +410,21 @@ struct column_gatherer_impl<dictionary32> {
     // 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<column>(dictionary.keys(), stream, mr);
+    auto keys_copy = std::make_unique<column>(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(),
       cudf::detail::indexalator_factory::make_output_iterator(new_indices->mutable_view()),
       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 +436,10 @@ struct column_gatherer_impl<struct_view> {
                                      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 +489,9 @@ struct column_gatherer_impl<struct_view> {
       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 +544,13 @@ void gather_bitmask(table_view const& source,
                     std::vector<std::unique_ptr<column>>& 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 +564,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 +574,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<size_type>(
-    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<size_type>(target.size(), stream, temp_mr);
 
   // Dispatch operation enum to get implementation
   auto const impl = [op]() {
@@ -621,7 +634,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 <typename MapIterator>
@@ -630,8 +643,10 @@ std::unique_ptr<table> 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<std::unique_ptr<column>> destination_columns;
 
   // TODO: Could be beneficial to use streams internally here
@@ -661,7 +676,8 @@ std::unique_ptr<table> 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);
       }
     }
   }

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