Add libcudf wrappers around current_device_resource functions. (#16679)

Merge after rapidsai/rmm#1661

Creates and uses CUDF internal wrappers around RMM `current_device_resource` functions.

I've marked this PR as breaking because it breaks the ABI, however the API is compatible.

For reviewers, the most substantial additions are in the new file `<cudf/utilities/memory_resource.hpp>`, and in the `DEVELOPER_GUIDE.md` and `*.rst` docs. The rest are all replacements of an include and all calls to `rmm::get_current_device_resource()` with `cudf::get_current_device_resource_ref()`.

Closes #16676

Authors:
  - Mark Harris (https://github.com/harrism)

Approvers:
  - Nghia Truong (https://github.com/ttnghia)
  - GALI PREM SAGAR (https://github.com/galipremsagar)
  - https://github.com/nvdbaranec
  - David Wendt (https://github.com/davidwendt)

URL: #16679

cpp/benchmarks/common/generate_input.cu

@@ -28,10 +28,10 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/error.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/mr/device/per_device_resource.hpp>
 
 #include <cuda/functional>
 #include <thrust/binary_search.h>
@@ -507,7 +507,7 @@ std::unique_ptr<cudf::column> create_random_column(data_profile const& profile,
                            null_mask.end(),
                            thrust::identity<bool>{},
                            cudf::get_default_stream(),
-                           rmm::mr::get_current_device_resource());
+                           cudf::get_current_device_resource_ref());
 
   return std::make_unique<cudf::column>(
     dtype,
@@ -591,7 +591,7 @@ std::unique_ptr<cudf::column> create_random_utf8_string_column(data_profile cons
                            null_mask.end() - 1,
                            thrust::identity<bool>{},
                            cudf::get_default_stream(),
-                           rmm::mr::get_current_device_resource());
+                           cudf::get_current_device_resource_ref());
   return cudf::make_strings_column(
     num_rows,
     std::make_unique<cudf::column>(std::move(offsets), rmm::device_buffer{}, 0),
@@ -626,7 +626,7 @@ std::unique_ptr<cudf::column> create_random_column<cudf::string_view>(data_profi
                                         cudf::out_of_bounds_policy::DONT_CHECK,
                                         cudf::detail::negative_index_policy::NOT_ALLOWED,
                                         cudf::get_default_stream(),
-                                        rmm::mr::get_current_device_resource());
+                                        cudf::get_current_device_resource_ref());
   return std::move(str_table->release()[0]);
 }
 
@@ -688,7 +688,7 @@ std::unique_ptr<cudf::column> create_random_column<cudf::struct_view>(data_profi
                                         valids.end(),
                                         thrust::identity<bool>{},
                                         cudf::get_default_stream(),
-                                        rmm::mr::get_current_device_resource());
+                                        cudf::get_current_device_resource_ref());
         }
         return std::pair<rmm::device_buffer, cudf::size_type>{};
       }();
@@ -782,7 +782,7 @@ std::unique_ptr<cudf::column> create_random_column<cudf::list_view>(data_profile
                                                           valids.end(),
                                                           thrust::identity<bool>{},
                                                           cudf::get_default_stream(),
-                                                          rmm::mr::get_current_device_resource());
+                                                          cudf::get_current_device_resource_ref());
     list_column                  = cudf::make_lists_column(
       current_num_rows,
       std::move(offsets_column),
@@ -933,7 +933,7 @@ std::pair<rmm::device_buffer, cudf::size_type> create_random_null_mask(
                                   thrust::make_counting_iterator<cudf::size_type>(size),
                                   bool_generator{seed, 1.0 - *null_probability},
                                   cudf::get_default_stream(),
-                                  rmm::mr::get_current_device_resource());
+                                  cudf::get_current_device_resource_ref());
   }
 }
 

cpp/benchmarks/common/tpch_data_generator/random_column_generator.hpp

@@ -17,6 +17,8 @@
 #pragma once
 
 #include <cudf/column/column.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 #include <string>
 
@@ -36,7 +38,7 @@ std::unique_ptr<cudf::column> generate_random_string_column(
   cudf::size_type upper,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a column of random numbers
@@ -61,7 +63,7 @@ std::unique_ptr<cudf::column> generate_random_numeric_column(
   T upper,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a primary key column
@@ -81,7 +83,7 @@ std::unique_ptr<cudf::column> generate_primary_key_column(
   cudf::scalar const& start,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a column where all the rows have the same string value
@@ -101,7 +103,7 @@ std::unique_ptr<cudf::column> generate_repeat_string_column(
   std::string const& value,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a column by randomly choosing from set of strings
@@ -121,7 +123,7 @@ std::unique_ptr<cudf::column> generate_random_string_column_from_set(
   cudf::host_span<const char* const> set,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a column consisting of a repeating sequence of integers
@@ -145,6 +147,6 @@ std::unique_ptr<cudf::column> generate_repeat_sequence_column(
   bool zero_indexed,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 }  // namespace cudf::datagen

cpp/benchmarks/common/tpch_data_generator/table_helpers.cpp

@@ -35,6 +35,9 @@
 #include <cudf/transform.hpp>
 #include <cudf/unary.hpp>
 
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/resource_ref.hpp>
+
 #include <vector>
 
 namespace cudf::datagen {

cpp/benchmarks/common/tpch_data_generator/table_helpers.hpp

@@ -20,6 +20,8 @@
 #include <cudf/column/column_view.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 #include <vector>
 
@@ -37,7 +39,7 @@ std::unique_ptr<cudf::column> add_calendrical_days(
   cudf::column_view const& timestamp_days,
   cudf::column_view const& days,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Perform a left join operation between two tables
@@ -56,7 +58,7 @@ std::unique_ptr<cudf::table> perform_left_join(
   std::vector<cudf::size_type> const& left_on,
   std::vector<cudf::size_type> const& right_on,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `p_retailprice` column of the `part` table
@@ -68,7 +70,7 @@ std::unique_ptr<cudf::table> perform_left_join(
 [[nodiscard]] std::unique_ptr<cudf::column> calculate_p_retailprice(
   cudf::column_view const& p_partkey,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `l_suppkey` column of the `lineitem` table
@@ -84,7 +86,7 @@ std::unique_ptr<cudf::table> perform_left_join(
   cudf::size_type scale_factor,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `ps_suppkey` column of the `partsupp` table
@@ -100,7 +102,7 @@ std::unique_ptr<cudf::table> perform_left_join(
   cudf::size_type scale_factor,
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 /**
  * @brief Calculate the cardinality of the `lineitem` table
  *
@@ -111,7 +113,7 @@ std::unique_ptr<cudf::table> perform_left_join(
 [[nodiscard]] cudf::size_type calculate_l_cardinality(
   cudf::column_view const& o_rep_freqs,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 /**
  * @brief Calculate the charge column for the `lineitem` table
  *
@@ -126,7 +128,7 @@ std::unique_ptr<cudf::table> perform_left_join(
   cudf::column_view const& tax,
   cudf::column_view const& discount,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a column of random addresses according to TPC-H specification clause 4.2.2.7
@@ -138,7 +140,7 @@ std::unique_ptr<cudf::table> perform_left_join(
 [[nodiscard]] std::unique_ptr<cudf::column> generate_address_column(
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate a phone number column according to TPC-H specification clause 4.2.2.9
@@ -150,6 +152,6 @@ std::unique_ptr<cudf::table> perform_left_join(
 [[nodiscard]] std::unique_ptr<cudf::column> generate_phone_column(
   cudf::size_type num_rows,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 }  // namespace cudf::datagen

cpp/benchmarks/common/tpch_data_generator/tpch_data_generator.cpp

@@ -36,6 +36,9 @@
 #include <cudf/transform.hpp>
 #include <cudf/unary.hpp>
 
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/resource_ref.hpp>
+
 #include <array>
 #include <string>
 #include <vector>

cpp/benchmarks/common/tpch_data_generator/tpch_data_generator.hpp

@@ -17,6 +17,8 @@
 #pragma once
 
 #include <cudf/table/table.hpp>
+#include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 namespace CUDF_EXPORT cudf {
 namespace datagen {
@@ -32,7 +34,7 @@ std::tuple<std::unique_ptr<cudf::table>, std::unique_ptr<cudf::table>, std::uniq
 generate_orders_lineitem_part(
   double scale_factor,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `partsupp` table
@@ -44,7 +46,7 @@ generate_orders_lineitem_part(
 std::unique_ptr<cudf::table> generate_partsupp(
   double scale_factor,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `supplier` table
@@ -56,7 +58,7 @@ std::unique_ptr<cudf::table> generate_partsupp(
 std::unique_ptr<cudf::table> generate_supplier(
   double scale_factor,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `customer` table
@@ -68,7 +70,7 @@ std::unique_ptr<cudf::table> generate_supplier(
 std::unique_ptr<cudf::table> generate_customer(
   double scale_factor,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `nation` table
@@ -78,7 +80,7 @@ std::unique_ptr<cudf::table> generate_customer(
  */
 std::unique_ptr<cudf::table> generate_nation(
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**
  * @brief Generate the `region` table
@@ -88,7 +90,7 @@ std::unique_ptr<cudf::table> generate_nation(
  */
 std::unique_ptr<cudf::table> generate_region(
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 }  // namespace datagen
 }  // namespace CUDF_EXPORT cudf

cpp/benchmarks/copying/contiguous_split.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/contiguous_split.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 #include <thrust/iterator/counting_iterator.h>
 
@@ -32,7 +33,7 @@ void contiguous_split(cudf::table_view const& src_table, std::vector<cudf::size_
 
 void chunked_pack(cudf::table_view const& src_table, std::vector<cudf::size_type> const&)
 {
-  auto const mr     = rmm::mr::get_current_device_resource();
+  auto const mr     = cudf::get_current_device_resource_ref();
   auto const stream = cudf::get_default_stream();
   auto user_buffer  = rmm::device_uvector<std::uint8_t>(100L * 1024 * 1024, stream, mr);
   auto chunked_pack = cudf::chunked_pack::create(src_table, user_buffer.size(), mr);

cpp/benchmarks/copying/shift.cu

@@ -20,14 +20,13 @@
 #include <cudf/copying.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/default_stream.hpp>
-
-#include <rmm/resource_ref.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 template <typename T, typename ScalarType = cudf::scalar_type_t<T>>
 std::unique_ptr<cudf::scalar> make_scalar(
   T value                           = 0,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
-  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource())
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref())
 {
   auto s = new ScalarType(value, true, stream, mr);
   return std::unique_ptr<cudf::scalar>(s);

cpp/benchmarks/fixture/benchmark_fixture.hpp

@@ -16,10 +16,11 @@
 
 #pragma once
 
+#include <cudf/utilities/memory_resource.hpp>
+
 #include <rmm/cuda_device.hpp>
 #include <rmm/mr/device/cuda_memory_resource.hpp>
 #include <rmm/mr/device/owning_wrapper.hpp>
-#include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
 #include <rmm/mr/device/statistics_resource_adaptor.hpp>
 
@@ -83,13 +84,13 @@ class benchmark : public ::benchmark::Fixture {
   void SetUp(::benchmark::State const& state) override
   {
     mr = make_pool_instance();
-    rmm::mr::set_current_device_resource(mr.get());  // set default resource to pool
+    cudf::set_current_device_resource(mr.get());  // set default resource to pool
   }
 
   void TearDown(::benchmark::State const& state) override
   {
     // reset default resource to the initial resource
-    rmm::mr::set_current_device_resource(nullptr);
+    cudf::set_current_device_resource(nullptr);
     mr.reset();
   }
 
@@ -106,13 +107,13 @@ class benchmark : public ::benchmark::Fixture {
 class memory_stats_logger {
  public:
   memory_stats_logger()
-    : existing_mr(rmm::mr::get_current_device_resource()),
+    : existing_mr(cudf::get_current_device_resource()),
       statistics_mr(rmm::mr::statistics_resource_adaptor(existing_mr))
   {
-    rmm::mr::set_current_device_resource(&statistics_mr);
+    cudf::set_current_device_resource(&statistics_mr);
   }
 
-  ~memory_stats_logger() { rmm::mr::set_current_device_resource(existing_mr); }
+  ~memory_stats_logger() { cudf::set_current_device_resource(existing_mr); }
 
   [[nodiscard]] size_t peak_memory_usage() const noexcept
   {