Add ability to prefetch in cudf.pandas and change default to managed pool

python/cudf/cudf/pandas/__init__.py

@@ -25,7 +25,11 @@ def install():
     global LOADED
     LOADED = loader is not None
 
-    if (rmm_mode := os.getenv("CUDF_PANDAS_RMM_MODE", None)) is not None:
+    if (
+        rmm_mode := os.getenv("CUDF_PANDAS_RMM_MODE", "managed_pool")
+    ) is not None:
+        if rmm_mode == "disable":
+            return
         # Check if a non-default memory resource is set
         current_mr = rmm.mr.get_current_device_resource()
         if not isinstance(current_mr, rmm.mr.CudaMemoryResource):
@@ -50,16 +54,27 @@ def install():
             mr = rmm.mr.CudaAsyncMemoryResource(initial_pool_size=free_memory)
             rmm.mr.set_current_device_resource(mr)
         elif rmm_mode == "managed":
-            mr = rmm.mr.ManagedMemoryResource()
+            mr = rmm.mr.PrefetchResourceAdaptor(rmm.mr.ManagedMemoryResource())
             rmm.mr.set_current_device_resource(mr)
         elif rmm_mode == "managed_pool":
-            mr = rmm.mr.PoolMemoryResource(
-                rmm.mr.ManagedMemoryResource(),
-                initial_pool_size=free_memory,
+            mr = rmm.mr.PrefetchResourceAdaptor(
+                rmm.mr.PoolMemoryResource(
+                    rmm.mr.ManagedMemoryResource(),
+                    initial_pool_size=free_memory,
+                )
             )
             rmm.mr.set_current_device_resource(mr)
         else:
             raise ValueError(f"Unsupported rmm mode: {rmm_mode}")
+        from cudf._lib import pylibcudf
+
+        for item in {
+            "column_view::get_data",
+            "mutable_column_view::get_data",
+            "gather",
+            "hash_join",
+        }:
+            pylibcudf.experimental.enable_prefetching(item)
 
 
 def pytest_load_initial_conftests(early_config, parser, args):