From 58a27ba968e1b9021ba9937ad4a0fe55eb3ff2c2 Mon Sep 17 00:00:00 2001
From: Hendrik Groove <hendrik.groove@dw.com>
Date: Sun, 20 Oct 2024 20:59:23 +0200
Subject: [PATCH] more logging

---
 src/common/device_helpers.hip.h | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/common/device_helpers.hip.h b/src/common/device_helpers.hip.h
index 083cd2fc6..300e6a029 100644
--- a/src/common/device_helpers.hip.h
+++ b/src/common/device_helpers.hip.h
@@ -394,13 +394,14 @@ inline void ThrowOOMError(std::string const& err, size_t bytes) {
 template <class T>
 struct XGBDefaultDeviceAllocatorImpl : XGBBaseDeviceAllocator<T> {
   using SuperT = XGBBaseDeviceAllocator<T>;
-  using pointer = thrust::device_ptr<T>;  // NOLINT
+  using pointer = thrust::device_ptr<T>;
+
   template<typename U>
-  struct rebind  // NOLINT
-  {
-    using other = XGBDefaultDeviceAllocatorImpl<U>;  // NOLINT
+  struct rebind {
+    using other = XGBDefaultDeviceAllocatorImpl<U>;
   };
-  pointer allocate(size_t n) {  // NOLINT
+
+  pointer allocate(size_t n) {
     pointer ptr;
     try {
       ptr = SuperT::allocate(n);
@@ -408,17 +409,22 @@ struct XGBDefaultDeviceAllocatorImpl : XGBBaseDeviceAllocator<T> {
     } catch (const std::exception &e) {
       ThrowOOMError(e.what(), n * sizeof(T));
     }
+    std::cerr << "XGBDefaultDeviceAllocatorImpl: Allocated " << n * sizeof(T) 
+              << " bytes at " << ptr.get() << std::endl;
     GlobalMemoryLogger().RegisterAllocation(ptr.get(), n * sizeof(T));
     return ptr;
   }
-  void deallocate(pointer ptr, size_t n) {  // NOLINT
+
+  void deallocate(pointer ptr, size_t n) {
+    std::cerr << "XGBDefaultDeviceAllocatorImpl: Deallocating " << n * sizeof(T) 
+              << " bytes at " << ptr.get() << std::endl;
     GlobalMemoryLogger().RegisterDeallocation(ptr.get(), n * sizeof(T));
     SuperT::deallocate(ptr, n);
   }
-#if defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
-  XGBDefaultDeviceAllocatorImpl()
-    : SuperT(rmm::cuda_stream_default, rmm::mr::get_current_device_resource()) {}
-#endif  // defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
+
+  #if defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
+  XGBDefaultDeviceAllocatorImpl() : SuperT(rmm::cuda_stream_default, rmm::mr::get_current_device_resource()) {}
+  #endif
 };
 
 /**