diff --git a/src/common/numeric.cu b/src/common/numeric.cu index b292edf1a..818de69a0 100644 --- a/src/common/numeric.cu +++ b/src/common/numeric.cu @@ -3,7 +3,12 @@ */ #include +#if defined(XGBOOST_USE_CUDA) #include "device_helpers.cuh" // dh::Reduce, dh::XGBCachingDeviceAllocator +#elif defined(XGBOOST_USE_HIP) +#include "device_helpers.hip.h" // dh::Reduce, dh::XGBCachingDeviceAllocator +#endif + #include "numeric.h" #include "xgboost/context.h" // Context #include "xgboost/host_device_vector.h" // HostDeviceVector @@ -15,8 +20,14 @@ double Reduce(Context const* ctx, HostDeviceVector const& values) { values.SetDevice(ctx->gpu_id); auto const d_values = values.ConstDeviceSpan(); dh::XGBCachingDeviceAllocator alloc; + +#if defined(XGBOOST_USE_CUDA) return dh::Reduce(thrust::cuda::par(alloc), dh::tcbegin(d_values), dh::tcend(d_values), 0.0, thrust::plus{}); +#elif defined(XGBOOST_USE_HIP) + return dh::Reduce(thrust::hip::par(alloc), dh::tcbegin(d_values), dh::tcend(d_values), 0.0, + thrust::plus{}); +#endif } } // namespace cuda_impl } // namespace common diff --git a/src/common/numeric.h b/src/common/numeric.h index 6a1c15fd0..9d255e9af 100644 --- a/src/common/numeric.h +++ b/src/common/numeric.h @@ -97,12 +97,12 @@ void PartialSum(int32_t n_threads, InIt begin, InIt end, T init, OutIt out_it) { namespace cuda_impl { double Reduce(Context const* ctx, HostDeviceVector const& values); -#if !defined(XGBOOST_USE_CUDA) +#if !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP) inline double Reduce(Context const*, HostDeviceVector const&) { AssertGPUSupport(); return 0; } -#endif // !defined(XGBOOST_USE_CUDA) +#endif // !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP) } // namespace cuda_impl /** diff --git a/src/common/numeric.hip b/src/common/numeric.hip index e69de29bb..19c125901 100644 --- a/src/common/numeric.hip +++ b/src/common/numeric.hip @@ -0,0 +1,4 @@ + +#if defined(XGBOOST_USE_HIP) +#include "numeric.cu" +#endif