Add a logging resource adaptor

rapidsai · Jul 31, 2024 · 1e7abc5 · 1e7abc5
1 parent cd5ec3c
commit 1e7abc5
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/cpp/benchmarks/common/cudf_datagen/utils.hpp b/cpp/benchmarks/common/cudf_datagen/utils.hpp
@@ -79,6 +79,22 @@
 #include <utility>
 #include <vector>
 
+class rmm_logger {
+ public:
+  rmm_logger()
+    : existing_mr(rmm::mr::get_current_device_resource()),
+      logging_mr(rmm::mr::make_logging_adaptor(existing_mr))
+  {
+    rmm::mr::set_current_device_resource(&logging_mr);
+  }
+
+  ~rmm_logger() { rmm::mr::set_current_device_resource(existing_mr); }
+
+ private:
+  rmm::mr::device_memory_resource* existing_mr;
+  rmm::mr::loggin_resource_adaptor<rmm::mr::device_memory_resource> logging_mr;
+};
+
 class memory_stats_logger {
  public:
   memory_stats_logger()
@@ -96,7 +112,6 @@ class memory_stats_logger {
   }
 
  private:
-  // TODO change to resource_ref once set_current_device_resource supports it
   rmm::mr::device_memory_resource* existing_mr;
   rmm::mr::statistics_resource_adaptor<rmm::mr::device_memory_resource> statistics_mr;
 };