cuml.accel: serve managed memory from an RMM pool

docs/source/cuml-accel/usage.rst

@@ -72,7 +72,10 @@ When the platform supports it and RMM has not already been configured,
 ``cuml.accel`` enables `managed memory`_. Managed memory can use host memory in
 addition to GPU memory and migrates data as needed, reducing the risk of GPU
 out-of-memory errors. It does not prevent exhaustion of combined host and
-device memory, and heavy oversubscription can slow execution. Managed memory
+device memory, and heavy oversubscription can slow execution. The managed
+allocations are served from an RMM memory pool (as in ``cudf.pandas``'s
+``managed_pool`` mode), so repeated estimator calls do not pay a CUDA
+allocation and free each; the pool starts small and grows on demand. Managed memory
 is not enabled on WSL 2 or when RMM was configured before ``cuml.accel``.
 If managed-memory oversubscription causes unexpectedly slow execution, use the
 CLI's ``--disable-uvm`` flag to disable it and compare performance.

python/cuml/cuml/accel/core.py

@@ -255,12 +255,25 @@ def install(
                     "skipping enabling managed memory."
                 )
             else:
+                # Pool the managed allocations. Without a pool every
+                # allocation is a cudaMallocManaged and every free a
+                # cudaFree, each costing milliseconds (up to a dozen frees
+                # per predict call), which dominated short estimator calls.
+                # Same construction as cudf.pandas's "managed_pool" mode,
+                # but with a small initial block that grows on demand:
+                # cuml.accel is also installed in every worker process a
+                # user spawns (joblib, multiprocessing), so reserving a
+                # large share of the device up front in each of them is
+                # not acceptable.
                 rmm.mr.set_current_device_resource(
                     rmm.mr.PrefetchResourceAdaptor(
-                        rmm.mr.ManagedMemoryResource()
+                        rmm.mr.PoolMemoryResource(
+                            rmm.mr.ManagedMemoryResource(),
+                            initial_pool_size=256 << 20,
+                        )
                     )
                 )
-                logger.debug("Enabled managed memory.")
+                logger.debug("Enabled pooled managed memory.")
         else:
             logger.debug("Could not enable managed memory on this platform.")