add context.hip

2023-03-08 07:34:19 +01:00
parent 7a3a9b682a
commit bdcb036592
4 changed files with 11 additions and 5 deletions
--- a/src/common/cuda_context.cuh
+++ b/src/common/cuda_context.cuh
@@ -5,7 +5,11 @@
 #define XGBOOST_COMMON_CUDA_CONTEXT_CUH_
 #include <thrust/execution_policy.h>

+#if defined(XGBOOST_USE_HIP)
+#include "device_helpers.hip.h"
+#elif defined(XGBOOST_USE_CUDA)
 #include "device_helpers.cuh"
+#endif

 namespace xgboost {
 struct CUDAContext {
--- a/src/common/device_helpers.hip.h
+++ b/src/common/device_helpers.hip.h
@@ -23,7 +23,7 @@
 #include <chrono>
 #include <cstddef>  // for size_t
 #include <hipcub/hipcub.hpp>
-#include <cub/util_allocator.cuh>
+#include <hipcub/util_allocator.hpp>
 #include <numeric>
 #include <sstream>
 #include <string>