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

enable rocm, fix cuda_pinned_allocator.h

Browse Source
This commit is contained in:
amdsc21 2023-03-08 06:39:40 +01:00
parent d8cc93f3f2
commit 05fdca893f
1 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
src/common/cuda_pinned_allocator.h
View File

@ -72,11 +72,23 @@ class pinned_allocator {
if (cnt > this->max_size()) { throw std::bad_alloc(); } // end if
pointer result(nullptr);
#if defined(XGBOOST_USE_HIP)
dh::safe_cuda(hipMallocHost(reinterpret_cast<void**>(&result), cnt * sizeof(value_type)));
#else
dh::safe_cuda(cudaMallocHost(reinterpret_cast<void**>(&result), cnt * sizeof(value_type)));
#endif
return result;
}
inline void deallocate(pointer p, size_type) { dh::safe_cuda(cudaFreeHost(p)); } // NOLINT
inline void deallocate(pointer p, size_type) {
#if defined(XGBOOST_USE_HIP)
dh::safe_cuda(hipFreeHost(p));
#else
dh::safe_cuda(cudaFreeHost(p));
#endif
} // NOLINT
inline size_type max_size() const { return (std::numeric_limits<size_type>::max)() / sizeof(T); } // NOLINT