Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

more logging

Browse Source
This commit is contained in:
Hendrik Groove 2024-10-20 20:59:23 +02:00
parent c964dd62b4
commit 58a27ba968
1 changed files with 16 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
src/common/device_helpers.hip.h
View File

@ -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
};
/**