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

enable rocm, fix cuda_context.cuh

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

10
src/common/cuda_context.cuh
View File

@ -17,11 +17,21 @@ struct CUDAContext {
/**
* \brief Caching thrust policy.
*/
#if defined(XGBOOST_USE_HIP)
auto CTP() const { return thrust::hip::par(caching_alloc_).on(dh::DefaultStream()); }
#else
auto CTP() const { return thrust::cuda::par(caching_alloc_).on(dh::DefaultStream()); }
#endif
/**
* \brief Thrust policy without caching allocator.
*/
#if defined(XGBOOST_USE_HIP)
auto TP() const { return thrust::hip::par(alloc_).on(dh::DefaultStream()); }
#else
auto TP() const { return thrust::cuda::par(alloc_).on(dh::DefaultStream()); }
#endif
auto Stream() const { return dh::DefaultStream(); }
};
} // namespace xgboost