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

Use caching allocator from RMM, when RMM is enabled (#6131)

Browse Source
This commit is contained in:
Jiaming Yuan
2020-09-18 12:51:49 +08:00
committed by GitHub
parent 6bc9b9dc4f
commit 5384ed85c8
2 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/common/device_helpers.cuh
View File

@@ -406,10 +406,14 @@ struct XGBDefaultDeviceAllocatorImpl : XGBBaseDeviceAllocator<T> {
}
#if defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
XGBDefaultDeviceAllocatorImpl()
: SuperT(rmm::mr::get_current_device_resource(), cudaStream_t{0}) {}
: SuperT(rmm::mr::get_current_device_resource(), cudaStream_t{nullptr}) {}
#endif // defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
};
#if defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
template <typename T>
using XGBCachingDeviceAllocatorImpl = XGBDefaultDeviceAllocatorImpl<T>;
#else
/**
* \brief Caching memory allocator, uses cub::CachingDeviceAllocator as a back-end and logs
* allocations if verbose. Does not initialise memory on construction.
@@ -448,6 +452,7 @@ struct XGBCachingDeviceAllocatorImpl : thrust::device_malloc_allocator<T> {
// no-op
}
};
#endif // defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
} // namespace detail
// Declare xgboost allocators