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

enable rocm, fix threading_utils.cuh

Browse Source
This commit is contained in:
amdsc21 2023-03-08 06:30:52 +01:00
parent 327f1494f1
commit 2eb0b6aae4
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/common/threading_utils.cuh
View File

@ -62,9 +62,17 @@ SegmentedTrapezoidThreads(xgboost::common::Span<U> group_ptr,
dh::InclusiveSum(out_group_threads_ptr.data(), out_group_threads_ptr.data(),
out_group_threads_ptr.size());
size_t total = 0;
#if defined(XGBOOST_USE_HIP)
dh::safe_cuda(hipMemcpy(
&total, out_group_threads_ptr.data() + out_group_threads_ptr.size() - 1,
sizeof(total), hipMemcpyDeviceToHost));
#else
dh::safe_cuda(cudaMemcpy(
&total, out_group_threads_ptr.data() + out_group_threads_ptr.size() - 1,
sizeof(total), cudaMemcpyDeviceToHost));
#endif
return total;
}