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

python/cudf/cudf/pandas/__init__.py

@@ -25,14 +25,16 @@ def install():
     global LOADED
     LOADED = loader is not None
 
-    if (rmm_mode := os.getenv("CUDF_PANDAS_RMM_MODE", None)) is not None:
+    rmm_mode = os.getenv("CUDF_PANDAS_RMM_MODE", "managed_pool")
+    if rmm_mode != "disable":
         # 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):
             warnings.warn(
                 f"cudf.pandas detected an already configured memory resource, ignoring 'CUDF_PANDAS_RMM_MODE'={str(rmm_mode)}",
                 UserWarning,
             )
+        enable_prefetching = False
         free_memory, _ = rmm.mr.available_device_memory()
         free_memory = int(round(float(free_memory) * 0.80 / 256) * 256)
 
@@ -50,16 +52,30 @@ 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)
+            enable_prefetching = True
         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)
+            enable_prefetching = True
         else:
             raise ValueError(f"Unsupported rmm mode: {rmm_mode}")
+        from cudf._lib import pylibcudf
+
+        if enable_prefetching:
+            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):