Add default pinned pool that falls back to new pinned allocations

cpp/include/cudf/io/memory_resource.hpp

@@ -41,4 +41,13 @@ rmm::host_async_resource_ref set_host_memory_resource(rmm::host_async_resource_r
  */
 rmm::host_async_resource_ref get_host_memory_resource();
 
+/**
+ * @brief Configure the size of the default host memory resource.
+ *
+ * Must be called before any other function in this header.
+ *
+ * @param size The size of the default host memory resource
+ */
+void config_host_memory_resource(size_t size);
+
 }  // namespace cudf::io

cpp/src/io/utilities/config_utils.cpp

@@ -206,14 +206,16 @@ static_assert(cuda::mr::resource_with<fixed_pinned_pool_memory_resource,
                                       cuda::mr::host_accessible>,
               "");
 
-rmm::host_async_resource_ref default_pinned_mr()
+rmm::host_async_resource_ref make_default_pinned_mr(std::optional<size_t> config_size)
 {
-  static fixed_pinned_pool_memory_resource mr = []() {
-    auto const size = []() -> size_t {
+  static fixed_pinned_pool_memory_resource mr = [config_size]() {
+    auto const size = [&config_size]() -> size_t {
       if (auto const env_val = getenv("LIBCUDF_PINNED_POOL_SIZE"); env_val != nullptr) {
         return std::atol(env_val);
       }
 
+      if (config_size.has_value()) { return *config_size; }
+
       size_t free{}, total{};
       CUDF_CUDA_TRY(cudaMemGetInfo(&free, &total));
       // 0.5% of the total device memory, capped at 100MB
@@ -228,16 +230,22 @@ rmm::host_async_resource_ref default_pinned_mr()
   return mr;
 }
 
+rmm::host_async_resource_ref make_host_mr(std::optional<size_t> size)
+{
+  static rmm::host_async_resource_ref mr_ref = make_default_pinned_mr(size);
+  return mr_ref;
+}
+
 std::mutex& host_mr_mutex()
 {
   static std::mutex map_lock;
   return map_lock;
 }
 
-rmm::host_async_resource_ref host_mr()
+rmm::host_async_resource_ref& host_mr()
 {
-  static rmm::host_async_resource_ref host_mr = default_pinned_mr();
-  return host_mr;
+  static rmm::host_async_resource_ref mr_ref = make_host_mr(std::nullopt);
+  return mr_ref;
 }
 
 }  // namespace
@@ -256,4 +264,10 @@ rmm::host_async_resource_ref get_host_memory_resource()
   return host_mr();
 }
 
+void config_host_memory_resource(size_t size)
+{
+  std::scoped_lock lock{host_mr_mutex()};
+  make_host_mr(size);
+}
+
 }  // namespace cudf::io