Add capability to use different allocators in cuML Python benchmarks

python/cuml/cuml/benchmark/run_benchmarks.py

@@ -18,6 +18,7 @@
 import json
 
 import numpy as np
+import rmm
 
 from cuml.benchmark import algorithms, datagen, runners
 
@@ -211,8 +212,34 @@ def extract_param_overrides(params_to_sweep):
         default="fp32",
         help="Precision of the dataset to benchmark with",
     )
+    parser.add_argument(
+        "--rmm-allocator",
+        choices=["cuda", "managed", "prefetched"],
+        default="cuda",
+        help="RMM memory resource to use (default: cuda)",
+    )
     args = parser.parse_args()
 
+    # Setup RMM allocator based on command line option
+    match args.rmm_allocator:
+        case "cuda":
+            dev_resource = rmm.mr.CudaMemoryResource()
+            rmm.mr.set_current_device_resource(dev_resource)
+            print("Using CUDA Memory Resource...")
+        case "managed":
+            managed_resource = rmm.mr.ManagedMemoryResource()
+            rmm.mr.set_current_device_resource(managed_resource)
+            print("Using Managed Memory Resource...")
+        case "prefetched":
+            upstream_mr = rmm.mr.ManagedMemoryResource()
+            prefetch_mr = rmm.mr.PrefetchResourceAdaptor(upstream_mr)
+            rmm.mr.set_current_device_resource(prefetch_mr)
+            print("Using Prefetched Managed Memory Resource...")
+        case _:
+            raise ValueError(
+                f"Unknown RMM allocator type: {args.rmm_allocator}"
+            )
+
     args.dtype = PrecisionMap[args.dtype]
 
     if args.print_algorithms: