Get CUDA memory pool stats after running into OOM.

Log the CUDA memory pool usage, in addition to normal CUDA memory.

Author: trivialfis

src/common/device_vector.cu

@@ -20,6 +20,34 @@ void ThrowOOMError(std::string const &err, std::size_t bytes) {
   ss << "Memory allocation error on worker " << rank << ": " << err << "\n"
      << "- Free memory: " << HumanMemUnit(dh::AvailableMemory(device)) << "\n"
      << "- Requested memory: " << HumanMemUnit(bytes) << std::endl;
+
+  cudaMemPool_t memPool;
+  std::size_t reserved_bytes = 0;
+  std::size_t used_bytes = 0;
+
+  // Get the default memory pool for the current device
+  auto status = cudaDeviceGetDefaultMemPool(&memPool, ::xgboost::curt::CurrentDevice());
+  if (status != cudaSuccess) {
+    ss << "Failed to get default memory pool: " << cudaGetErrorString(status) << "\n";
+    LOG(FATAL) << ss.str();
+  }
+
+  // Get the current total reserved memory size
+  status = cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrReservedMemCurrent, &reserved_bytes);
+  if (status != cudaSuccess) {
+    ss << "Failed to get reserved memory attribute: " << cudaGetErrorString(status) << "\n";
+    LOG(FATAL) << ss.str();
+  }
+  ss << "- Reserved by pool:" << HumanMemUnit(reserved_bytes) << "\n";
+
+  // Get the current total used memory size
+  status = cudaMemPoolGetAttribute(memPool, cudaMemPoolAttrUsedMemCurrent, &used_bytes);
+  if (status != cudaSuccess) {
+    ss << "Failed to get used memory attribute: " << cudaGetErrorString(status) << "\n";
+    LOG(FATAL) << ss.str();
+  }
+
+  ss << "- Used by pool:" << HumanMemUnit(used_bytes) << "\n";
   LOG(FATAL) << ss.str();
 }