Route valid_if through memory resources

cpp/benchmarks/common/generate_input.cu

@@ -1066,21 +1066,29 @@ std::unique_ptr<cudf::column> create_string_column(cudf::size_type num_rows,
 }
 
 std::pair<rmm::device_buffer, cudf::size_type> create_random_null_mask(
-  cudf::size_type size, std::optional<double> null_probability, unsigned seed)
+  cudf::size_type size,
+  std::optional<double> null_probability,
+  unsigned seed,
+  cuda::stream_ref stream,
+  cudf::memory_resources resources)
 {
   if (not null_probability.has_value()) { return {rmm::device_buffer{}, 0}; }
   CUDF_EXPECTS(*null_probability >= 0.0 and *null_probability <= 1.0,
                "Null probability must be within the range [0.0, 1.0]");
   if (*null_probability == 0.0f) {
-    return {cudf::create_null_mask(size, cudf::mask_state::ALL_VALID), 0};
+    return {
+      cudf::create_null_mask(size, cudf::mask_state::ALL_VALID, stream, resources.get_output_mr()),
+      0};
   } else if (*null_probability == 1.0) {
-    return {cudf::create_null_mask(size, cudf::mask_state::ALL_NULL), size};
+    return {
+      cudf::create_null_mask(size, cudf::mask_state::ALL_NULL, stream, resources.get_output_mr()),
+      size};
   } else {
     return cudf::detail::valid_if(cuda::counting_iterator<cudf::size_type>{0},
                                   cuda::counting_iterator<cudf::size_type>{size},
                                   bool_generator{seed, 1.0 - *null_probability},
-                                  cudf::get_default_stream(),
-                                  cudf::get_current_device_resource_ref());
+                                  stream,
+                                  resources);
   }
 }
 

cpp/benchmarks/common/generate_input.hpp

@@ -6,6 +6,8 @@
 #pragma once
 
 #include <cudf/table/table.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/span.hpp>
 #include <cudf/utilities/traits.hpp>
 
@@ -664,7 +666,13 @@ std::vector<cudf::type_id> mix_dtypes(std::pair<cudf::type_id, cudf::type_id> co
  * @param null_probability probability of a null value
  *  no value implies no null mask, =0 implies all valids, >=1 implies all nulls
  * @param seed Optional, seed for the pseudo-random engine
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param resources Memory resources used for the returned bitmask and temporary allocations
  * @return null mask device buffer with random null mask data and null count
  */
 std::pair<rmm::device_buffer, cudf::size_type> create_random_null_mask(
-  cudf::size_type size, std::optional<double> null_probability = std::nullopt, unsigned seed = 1);
+  cudf::size_type size,
+  std::optional<double> null_probability = std::nullopt,
+  unsigned seed                          = 1,
+  cuda::stream_ref stream                = cudf::get_default_stream(),
+  cudf::memory_resources resources       = cudf::get_current_device_resource_ref());

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

@@ -62,7 +62,7 @@ std::unique_ptr<cudf::column> copy_if_else(StringIterLeft lhs_begin,
       return filter_fn(idx) ? lhs_begin[idx].has_value() : rhs_begin[idx].has_value();
     },
     stream,
-    mr);
+    cudf::memory_resources{mr, mr});
   if (null_count == 0) { null_mask = rmm::device_buffer{}; }
 
   // build vector of strings

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

@@ -62,7 +62,8 @@ std::unique_ptr<column> make_strings_column(IndexPairIterator begin,
 
   // create null mask
   auto validator = [] __device__(string_index_pair const item) { return item.first != nullptr; };
-  auto new_nulls = cudf::detail::valid_if(begin, end, validator, stream, mr);
+  auto new_nulls =
+    cudf::detail::valid_if(begin, end, validator, stream, cudf::memory_resources{mr, mr});
   auto const null_count = new_nulls.second;
   auto null_mask =
     (null_count > 0) ? std::move(new_nulls.first) : rmm::device_buffer{0, stream, mr};

cpp/src/copying/shift.cu

@@ -54,7 +54,7 @@ std::pair<rmm::device_buffer, size_type> create_null_mask(column_device_view con
                           cuda::counting_iterator<size_type>{size},
                           func_validity,
                           stream,
-                          mr);
+                          cudf::memory_resources{mr, mr});
 }
 
 struct shift_functor {

cpp/src/dictionary/remove_keys.cu

@@ -137,7 +137,7 @@ std::unique_ptr<column> remove_keys_fn(dictionary_column_view const& dictionary_
       return (indices_itr[idx] < max_size);  // new nulls have max values
     },
     stream,
-    mr);
+    cudf::memory_resources{mr, mr});
   rmm::device_buffer new_null_mask =
     (new_nulls.second > 0) ? std::move(new_nulls.first) : rmm::device_buffer{0, stream, mr};
 

cpp/src/dictionary/set_keys.cu

@@ -120,7 +120,7 @@ struct remap_indices_dispatch_fn {
       iota + input.size(),
       [d_indices] __device__(size_type idx) { return d_indices[idx] >= 0; },
       stream,
-      mr);
+      cudf::memory_resources{mr, mr});
 
     return {std::move(indices_column), std::move(null_mask), null_count};
   }

cpp/src/groupby/common/m2_var_std.cu

@@ -140,8 +140,11 @@ std::unique_ptr<column> compute_variance_std(TransformFunc&& transform_fn,
                    out_it + size,
                    transform_fn);
 
-  auto [null_mask, null_count] =
-    cudf::detail::valid_if(validity.begin(), validity.end(), cuda::std::identity{}, stream, mr);
+  auto [null_mask, null_count] = cudf::detail::valid_if(validity.begin(),
+                                                        validity.end(),
+                                                        cuda::std::identity{},
+                                                        stream,
+                                                        cudf::memory_resources{mr, mr});
   if (null_count > 0) { output->set_null_mask(std::move(null_mask), null_count); }
 
   return output;

cpp/src/groupby/hash/hash_compound_agg_finalizer.cu

@@ -126,7 +126,7 @@ void hash_compound_agg_finalizer::operator()<aggregation::MEAN>(aggregation cons
       count_result.end<size_type>(),
       [] __device__(size_type const count) -> bool { return count > 0; },
       stream,
-      mr);
+      cudf::memory_resources{mr, mr});
     if (null_count > 0) { result->set_null_mask(std::move(null_mask), null_count); }
   }
   cache->add_result(col, agg, std::move(result));

cpp/src/groupby/sort/group_bitwise.cu

@@ -73,8 +73,11 @@ struct bitwise_group_reduction_functor {
                    validity.begin(),
                    cuda::std::logical_or{});
 
-      auto [null_mask, null_count] =
-        cudf::detail::valid_if(validity.begin(), validity.end(), cuda::std::identity{}, stream, mr);
+      auto [null_mask, null_count] = cudf::detail::valid_if(validity.begin(),
+                                                            validity.end(),
+                                                            cuda::std::identity{},
+                                                            stream,
+                                                            cudf::memory_resources{mr, mr});
       if (null_count > 0) { result->set_null_mask(std::move(null_mask), null_count); }
     }
     return result;

cpp/src/groupby/sort/group_correlation.cu

@@ -160,8 +160,11 @@ std::unique_ptr<column> group_covariance(column_view const& values_0,
   auto is_null = [ddof, min_periods] __device__(size_type group_size) {
     return not(group_size == 0 or group_size - ddof <= 0 or group_size < min_periods);
   };
-  auto [new_nullmask, null_count] =
-    cudf::detail::valid_if(count.begin<size_type>(), count.end<size_type>(), is_null, stream, mr);
+  auto [new_nullmask, null_count] = cudf::detail::valid_if(count.begin<size_type>(),
+                                                           count.end<size_type>(),
+                                                           is_null,
+                                                           stream,
+                                                           cudf::memory_resources{mr, mr});
   if (null_count != 0) { result->set_null_mask(std::move(new_nullmask), null_count); }
   return result;
 }

cpp/src/groupby/sort/group_single_pass_reduction_util.cuh

@@ -191,8 +191,11 @@ struct group_reduction_functor<
                    validity.begin(),
                    cuda::std::logical_or{});
 
-      auto [null_mask, null_count] =
-        cudf::detail::valid_if(validity.begin(), validity.end(), cuda::std::identity{}, stream, mr);
+      auto [null_mask, null_count] = cudf::detail::valid_if(validity.begin(),
+                                                            validity.end(),
+                                                            cuda::std::identity{},
+                                                            stream,
+                                                            cudf::memory_resources{mr, mr});
       result->set_null_mask(std::move(null_mask), null_count);
     }
     return result;
@@ -245,8 +248,11 @@ struct group_reduction_functor<
                    validity.begin(),
                    cuda::std::logical_or{});
 
-      auto [null_mask, null_count] =
-        cudf::detail::valid_if(validity.begin(), validity.end(), cuda::std::identity{}, stream, mr);
+      auto [null_mask, null_count] = cudf::detail::valid_if(validity.begin(),
+                                                            validity.end(),
+                                                            cuda::std::identity{},
+                                                            stream,
+                                                            cudf::memory_resources{mr, mr});
       result->set_null_mask(std::move(null_mask), null_count);
     }
 

cpp/src/groupby/sort/group_sum_overflow.cu

@@ -93,8 +93,11 @@ struct group_sum_overflow_fn {
         group_valid.begin(),
         cuda::std::equal_to<size_type>{},
         cuda::std::logical_or<bool>{});
-      return cudf::detail::valid_if(
-        group_valid.begin(), group_valid.end(), cuda::std::identity{}, stream, mr);
+      return cudf::detail::valid_if(group_valid.begin(),
+                                    group_valid.end(),
+                                    cuda::std::identity{},
+                                    stream,
+                                    cudf::memory_resources{mr, mr});
     }();
 
     std::vector<std::unique_ptr<column>> children;

cpp/src/labeling/label_bins.cu

@@ -153,7 +153,8 @@ std::unique_ptr<column> label_bins(column_view const& input,
                         left_begin, left_end, right_begin));
   });
 
-  auto mask_and_count = valid_if(output_begin, output_end, filter_null_sentinel(), stream, mr);
+  auto mask_and_count = valid_if(
+    output_begin, output_end, filter_null_sentinel(), stream, cudf::memory_resources{mr, mr});
 
   output->set_null_mask(std::move(mask_and_count.first), mask_and_count.second);
   return output;

cpp/src/lists/combine/concatenate_list_elements.cu

@@ -89,7 +89,7 @@ std::unique_ptr<column> concatenate_lists_ignore_null(column_view const& input,
           [&] __device__(auto const list_idx) { return lists_dv.is_valid(list_idx); });
       },
       stream,
-      mr);
+      cudf::memory_resources{mr, mr});
   }();
 
   return make_lists_column(num_rows,
@@ -212,8 +212,11 @@ std::unique_ptr<column> concatenate_lists_nullifying_rows(column_view const& inp
 
   auto list_entries =
     gather_list_entries(input, offsets_view, num_rows, num_output_entries, stream, mr);
-  auto [null_mask, null_count] = cudf::detail::valid_if(
-    list_validities.begin(), list_validities.end(), cuda::std::identity{}, stream, mr);
+  auto [null_mask, null_count] = cudf::detail::valid_if(list_validities.begin(),
+                                                        list_validities.end(),
+                                                        cuda::std::identity{},
+                                                        stream,
+                                                        cudf::memory_resources{mr, mr});
 
   return make_lists_column(num_rows,
                            std::move(list_offsets),

cpp/src/lists/combine/concatenate_rows.cu

@@ -137,7 +137,7 @@ generate_regrouped_offsets_and_null_mask(table_device_view const& input,
           return null_count != num_columns;
         },
         stream,
-        mr);
+        cudf::memory_resources{mr, mr});
     }
 
     // row is null if -any- input rows are null
@@ -146,7 +146,7 @@ generate_regrouped_offsets_and_null_mask(table_device_view const& input,
       row_null_counts.begin() + input.num_rows(),
       [] __device__(size_type null_count) { return null_count == 0; },
       stream,
-      mr);
+      cudf::memory_resources{mr, mr});
   }();
 
   return {std::move(offsets), std::move(null_mask), null_count};
@@ -248,7 +248,7 @@ std::unique_ptr<column> concatenate_rows(table_view const& input,
               return row_null_counts[row_index] != num_columns;
             }),
           stream,
-          cudf::get_current_device_resource_ref());
+          cudf::memory_resources{mr, mr});
       }
       // NULLIFY_OUTPUT_ROW.  Output row is nullfied if any input row is null
       return cudf::detail::valid_if(
@@ -261,7 +261,7 @@ std::unique_ptr<column> concatenate_rows(table_view const& input,
             return row_null_counts[row_index] == 0;
           }),
         stream,
-        cudf::get_current_device_resource_ref());
+        cudf::memory_resources{mr, mr});
     }();
     concat->set_null_mask(std::move(null_mask), null_count);
   }

cpp/src/lists/contains.cu

@@ -226,7 +226,7 @@ std::unique_ptr<column> dispatch_index_of(lists_column_view const& lists,
       output_it + num_rows,
       [] __device__(auto const idx) { return idx != NULL_SENTINEL; },
       stream,
-      mr);
+      cudf::memory_resources{mr, mr});
     out_positions->set_null_mask(std::move(null_mask), null_count);
   }
   return out_positions;

cpp/src/lists/copying/scatter_helper.cu

@@ -63,7 +63,7 @@ std::pair<rmm::device_buffer, size_type> construct_child_nullmask(
                                 cuda::counting_iterator<size_type>{num_child_rows},
                                 is_valid_predicate,
                                 stream,
-                                mr);
+                                cudf::memory_resources{mr, mr});
 }
 
 /**

cpp/src/lists/explode.cu

@@ -82,7 +82,7 @@ std::unique_ptr<table> build_table(
                                                explode_col_gather_map->end(),
                                                [] __device__(auto i) { return i != InvalidIndex; },
                                                stream,
-                                               mr)
+                                               cudf::memory_resources{mr, mr})
                                            : std::pair<rmm::device_buffer, size_type>{
                                                rmm::device_buffer(0, stream), size_type{0}};
 

cpp/src/lists/interleave_columns.cu

@@ -273,8 +273,11 @@ struct interleave_list_entries_impl<T, std::enable_if_t<cudf::is_fixed_width<T>(
       });
 
     if (data_has_null_mask) {
-      auto [null_mask, null_count] = cudf::detail::valid_if(
-        validities.begin(), validities.end(), cuda::std::identity{}, stream, mr);
+      auto [null_mask, null_count] = cudf::detail::valid_if(validities.begin(),
+                                                            validities.end(),
+                                                            cuda::std::identity{},
+                                                            stream,
+                                                            cudf::memory_resources{mr, mr});
       if (null_count > 0) { output->set_null_mask(std::move(null_mask), null_count); }
     }
 
@@ -367,8 +370,11 @@ std::unique_ptr<column> interleave_columns(table_view const& input,
       num_output_lists, std::move(list_offsets), std::move(list_entries), 0, rmm::device_buffer{});
   }
 
-  auto [null_mask, null_count] = cudf::detail::valid_if(
-    list_validities.begin(), list_validities.end(), cuda::std::identity{}, stream, mr);
+  auto [null_mask, null_count] = cudf::detail::valid_if(list_validities.begin(),
+                                                        list_validities.end(),
+                                                        cuda::std::identity{},
+                                                        stream,
+                                                        cudf::memory_resources{mr, mr});
   return make_lists_column(num_output_lists,
                            std::move(list_offsets),
                            std::move(list_entries),

cpp/src/quantiles/quantile.cu

@@ -117,7 +117,7 @@ struct quantile_functor {
           return select_quantile_validity(sorted_validity, size, q, interp);
         },
         stream,
-        mr);
+        cudf::memory_resources{mr, mr});
 
       output->set_null_mask(std::move(mask), null_count);
     }

cpp/src/quantiles/tdigest/tdigest.cu

@@ -220,7 +220,7 @@ std::unique_ptr<column> compute_approx_percentiles(tdigest_column_view const& in
                    return percentiles.is_valid(i % percentiles.size());
                  },
                  stream,
-                 mr)
+                 cudf::memory_resources{mr, mr})
              : std::pair<rmm::device_buffer, size_type>{rmm::device_buffer{}, 0};
   }();
 
@@ -385,8 +385,11 @@ std::unique_ptr<column> percentile_approx(tdigest_column_view const& input,
     if (null_count == 0) {
       return std::pair<rmm::device_buffer, size_type>{rmm::device_buffer{}, null_count};
     }
-    return cudf::detail::valid_if(
-      tdigest_is_empty, tdigest_is_empty + tdv.size(), cuda::std::logical_not{}, stream, mr);
+    return cudf::detail::valid_if(tdigest_is_empty,
+                                  tdigest_is_empty + tdv.size(),
+                                  cuda::std::logical_not{},
+                                  stream,
+                                  cudf::memory_resources{mr, mr});
   }();
 
   return cudf::make_lists_column(input.size(),

cpp/src/reshape/interleave_columns.cu

@@ -119,7 +119,7 @@ struct interleave_columns_impl<T, std::enable_if_t<std::is_same_v<T, cudf::struc
                                     cuda::counting_iterator<size_type>{output_size},
                                     validity_fn,
                                     stream,
-                                    mr);
+                                    cudf::memory_resources{mr, mr});
     };
 
     // Only create null mask if at least one input structs column is nullable.
@@ -223,7 +223,8 @@ struct interleave_columns_impl<T, std::enable_if_t<cudf::is_fixed_width<T>()>> {
                          func_value,
                          func_validity);
 
-    auto [mask, null_count] = valid_if(index_begin, index_end, func_validity, stream, mr);
+    auto [mask, null_count] =
+      valid_if(index_begin, index_end, func_validity, stream, cudf::memory_resources{mr, mr});
 
     output->set_null_mask(std::move(mask), null_count);
 

cpp/src/rolling/detail/rolling_collect_list.cuh

@@ -198,7 +198,7 @@ std::unique_ptr<column> rolling_collect_list(column_view const& input,
       return (preceding_begin[i] + following_begin[i]) >= min_periods;
     },
     stream,
-    mr);
+    cudf::memory_resources{mr, mr});
 
   return make_lists_column(input.size(),
                            std::move(offsets),

cpp/src/strings/combine/concatenate.cu

@@ -146,7 +146,7 @@ std::unique_ptr<column> concatenate(table_view const& strings_columns,
         thrust::seq, d_table.begin(), d_table.end(), [idx](auto col) { return col.is_null(idx); });
     },
     stream,
-    mr);
+    cudf::memory_resources{mr, mr});
 
   return make_strings_column(
     strings_count, std::move(offsets_column), chars.release(), null_count, std::move(null_mask));
@@ -239,7 +239,7 @@ std::unique_ptr<column> concatenate(table_view const& strings_columns,
         thrust::seq, d_table.begin(), d_table.end(), [idx](auto col) { return col.is_null(idx); });
     },
     stream,
-    mr);
+    cudf::memory_resources{mr, mr});
 
   return make_strings_column(
     strings_count, std::move(offsets_column), chars.release(), null_count, std::move(null_mask));