Use host_mr_options in JNI

abellina · May 22, 2024 · fd0d2b4 · fd0d2b4
1 parent 5b3e7a2
commit fd0d2b4
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 4 deletions.
diff --git a/cpp/include/cudf/io/memory_resource.hpp b/cpp/include/cudf/io/memory_resource.hpp
@@ -47,7 +47,8 @@ struct host_mr_options {
  * @return The previous resource that was in use
  */
 rmm::host_async_resource_ref set_host_memory_resource(
-  rmm::host_async_resource_ref mr, std::optional<host_mr_options> const& default_opts = std::nullopt);
+  rmm::host_async_resource_ref mr,
+  std::optional<host_mr_options> const& default_opts = std::nullopt);
 
 /**
  * @brief Get the rmm resource being used for host memory allocations by

diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp
@@ -256,8 +256,8 @@ rmm::host_async_resource_ref set_host_memory_resource(
   rmm::host_async_resource_ref mr, std::optional<host_mr_options> const& default_opts)
 {
   std::scoped_lock lock{host_mr_mutex()};
-  auto last_mr = host_mr(default_opts);
-  host_mr(std::nullopt)    = mr;
+  auto last_mr          = host_mr(default_opts);
+  host_mr(std::nullopt) = mr;
   return last_mr;
 }
 

diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp
@@ -1036,7 +1036,9 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_setCuioPinnedPoolMemoryResource(J
     // if the regular pinned pool is exhausted
     pinned_fallback_mr.reset(new pinned_fallback_host_memory_resource(pool));
     // set the cuio host mr and store the prior resource in our static variable
-    prior_cuio_host_mr() = cudf::io::set_host_memory_resource(*pinned_fallback_mr);
+    // we pass host_mr_options{0} so we disable any default pinned pool creation from cuDF
+    prior_cuio_host_mr() =
+      cudf::io::set_host_memory_resource(*pinned_fallback_mr, cudf::io::host_mr_options{0});
   }
   CATCH_STD(env, )
 }