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

---
 src/common/device_helpers.hip.h | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/common/device_helpers.hip.h b/src/common/device_helpers.hip.h
index 79f2f3390..083cd2fc6 100644
--- a/src/common/device_helpers.hip.h
+++ b/src/common/device_helpers.hip.h
@@ -441,12 +441,11 @@ struct XGBCachingDeviceAllocatorImpl : XGBBaseDeviceAllocator<T> {
         std::make_unique<hipcub::CachingDeviceAllocator>(2, 9, 29)};
     return *allocator;
   }
-  pointer allocate(size_t n) {  // NOLINT
+  pointer allocate(size_t n) {
     pointer thrust_ptr;
     if (use_cub_allocator_) {
       T* raw_ptr{nullptr};
-      auto errc =  GetGlobalCachingAllocator().DeviceAllocate(reinterpret_cast<void **>(&raw_ptr),
-                                                              n * sizeof(T));
+      auto errc = GetGlobalCachingAllocator().DeviceAllocate(reinterpret_cast<void **>(&raw_ptr), n * sizeof(T));
       if (errc != hipSuccess) {
         ThrowOOMError("Caching allocator", n * sizeof(T));
       }
@@ -459,10 +458,14 @@ struct XGBCachingDeviceAllocatorImpl : XGBBaseDeviceAllocator<T> {
         ThrowOOMError(e.what(), n * sizeof(T));
       }
     }
+    std::cerr << "XGBCachingDeviceAllocatorImpl: Allocated " << n * sizeof(T) 
+              << " bytes at " << thrust_ptr.get() << std::endl;
     GlobalMemoryLogger().RegisterAllocation(thrust_ptr.get(), n * sizeof(T));
     return thrust_ptr;
   }
-  void deallocate(pointer ptr, size_t n) {  // NOLINT
+  void deallocate(pointer ptr, size_t n) {
+    std::cerr << "XGBCachingDeviceAllocatorImpl: Deallocating " << n * sizeof(T) 
+              << " bytes at " << ptr.get() << std::endl;
     GlobalMemoryLogger().RegisterDeallocation(ptr.get(), n * sizeof(T));
     if (use_cub_allocator_) {
       GetGlobalCachingAllocator().DeviceFree(ptr.get());
@@ -574,6 +577,8 @@ xgboost::common::Span<T> LazyResize(xgboost::Context const *ctx,
 template <typename T>
 void CopyDeviceSpanToVector(std::vector<T> *dst, xgboost::common::Span<T> src) {
   CHECK_EQ(dst->size(), src.size());
+  std::cerr << "CopyDeviceSpanToVector: Copying " << src.size() * sizeof(T) 
+            << " bytes from device to host" << std::endl;
   dh::safe_cuda(hipMemcpyAsync(dst->data(), src.data(), dst->size() * sizeof(T),
                                 hipMemcpyDeviceToHost));
 }
@@ -601,10 +606,11 @@ void CopyToD(HContainer const &h, DContainer *d) {
   d->resize(h.size());
   using HVT = std::remove_cv_t<typename HContainer::value_type>;
   using DVT = std::remove_cv_t<typename DContainer::value_type>;
-  static_assert(std::is_same<HVT, DVT>::value,
-                "Host and device containers must have same value type.");
-  dh::safe_cuda(hipMemcpyAsync(d->data().get(), h.data(), h.size() * sizeof(HVT),
-                                hipMemcpyHostToDevice));
+  static_assert(std::is_same<HVT, DVT>::value, "Host and device containers must have same value type.");
+  std::cerr << "CopyToD: Copying " << h.size() * sizeof(HVT) 
+            << " bytes from host to device" << std::endl;
+  dh::safe_cuda(hipMemcpyAsync(d->data().get(), h.data(), h.size() * sizeof(HVT), 
+                               hipMemcpyHostToDevice));
 }
 
 // Keep track of pinned memory allocation