From 1c58ff61d172769d6fe13e3d725f79777bb12853 Mon Sep 17 00:00:00 2001
From: amdsc21 <96135754+amdsc21@users.noreply.github.com>
Date: Fri, 10 Mar 2023 00:46:29 +0100
Subject: [PATCH] finish fit_stump.cu

---
 src/tree/constraints.cu |  2 +-
 src/tree/fit_stump.cc   |  4 ++--
 src/tree/fit_stump.cu   | 12 ++++++++++++
 src/tree/fit_stump.hip  |  4 ++++
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/tree/constraints.cu b/src/tree/constraints.cu
index 1065b9689..c5993dd1d 100644
--- a/src/tree/constraints.cu
+++ b/src/tree/constraints.cu
@@ -15,7 +15,7 @@
 #include "constraints.cuh"
 #include "param.h"
 
-#if defined(XGBOOST_USE_hip.CUDA)
+#if defined(XGBOOST_USE_CUDA)
 #include "../common/device_helpers.cuh"
 #elif defined(XGBOOST_USE_HIP)
 #include "../common/device_helpers.hip.h"
diff --git a/src/tree/fit_stump.cc b/src/tree/fit_stump.cc
index 82efff2c7..d8c08da12 100644
--- a/src/tree/fit_stump.cc
+++ b/src/tree/fit_stump.cc
@@ -56,12 +56,12 @@ namespace cuda_impl {
 void FitStump(Context const* ctx, linalg::TensorView<GradientPair const, 2> gpair,
               linalg::VectorView<float> out);
 
-#if !defined(XGBOOST_USE_CUDA)
+#if !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP)
 inline void FitStump(Context const*, linalg::TensorView<GradientPair const, 2>,
                      linalg::VectorView<float>) {
   common::AssertGPUSupport();
 }
-#endif  // !defined(XGBOOST_USE_CUDA)
+#endif  // !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_C
 }  // namespace cuda_impl
 
 void FitStump(Context const* ctx, HostDeviceVector<GradientPair> const& gpair,
diff --git a/src/tree/fit_stump.cu b/src/tree/fit_stump.cu
index 58a1fae82..bc206155f 100644
--- a/src/tree/fit_stump.cu
+++ b/src/tree/fit_stump.cu
@@ -12,7 +12,13 @@
 #include <cstddef>                                // std::size_t
 
 #include "../collective/device_communicator.cuh"  // DeviceCommunicator
+
+#if defined(XGBOOST_USE_CUDA)
 #include "../common/device_helpers.cuh"           // dh::MakeTransformIterator
+#elif defined(XGBOOST_USE_HIP)
+#include "../common/device_helpers.hip.h"         // dh::MakeTransformIterator
+#endif
+
 #include "fit_stump.h"
 #include "xgboost/base.h"     // GradientPairPrecise, GradientPair, XGBOOST_DEVICE
 #include "xgboost/context.h"  // Context
@@ -45,7 +51,13 @@ void FitStump(Context const* ctx, linalg::TensorView<GradientPair const, 2> gpai
   CHECK(d_sum.CContiguous());
 
   dh::XGBCachingDeviceAllocator<char> alloc;
+
+#if defined(XGBOOST_USE_CUDA)
   auto policy = thrust::cuda::par(alloc);
+#elif defined(XGBOOST_USE_HIP)
+  auto policy = thrust::hip::par(alloc);
+#endif
+
   thrust::reduce_by_key(policy, key_it, key_it + gpair.Size(), grad_it,
                         thrust::make_discard_iterator(), dh::tbegin(d_sum.Values()));
 
diff --git a/src/tree/fit_stump.hip b/src/tree/fit_stump.hip
index e69de29bb..6b4ddd0af 100644
--- a/src/tree/fit_stump.hip
+++ b/src/tree/fit_stump.hip
@@ -0,0 +1,4 @@
+
+#if defined(XGBOOST_USE_HIP)
+#include "fit_stump.cu"
+#endif