From 05fdca893f94c9423bea5da27794617538b4b32d Mon Sep 17 00:00:00 2001 From: amdsc21 <96135754+amdsc21@users.noreply.github.com> Date: Wed, 8 Mar 2023 06:39:40 +0100 Subject: [PATCH] enable rocm, fix cuda_pinned_allocator.h --- src/common/cuda_pinned_allocator.h | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/common/cuda_pinned_allocator.h b/src/common/cuda_pinned_allocator.h index d11851d99..a5152c8a0 100644 --- a/src/common/cuda_pinned_allocator.h +++ b/src/common/cuda_pinned_allocator.h @@ -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(&result), cnt * sizeof(value_type))); +#else dh::safe_cuda(cudaMallocHost(reinterpret_cast(&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::max)() / sizeof(T); } // NOLINT