From b9d86d44d6b84dd4155d8fb965fc9400190a1a39 Mon Sep 17 00:00:00 2001
From: amdsc21 <96135754+amdsc21@users.noreply.github.com>
Date: Thu, 9 Mar 2023 20:37:16 +0100
Subject: [PATCH] finish multiclass_metric.cu

---
 src/metric/multiclass_metric.cc  |  2 +-
 src/metric/multiclass_metric.cu  | 41 ++++++++++++++++++++++++++------
 src/metric/multiclass_metric.hip |  4 ++++
 3 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/metric/multiclass_metric.cc b/src/metric/multiclass_metric.cc
index 1257fb0fa..2b6d5a96d 100644
--- a/src/metric/multiclass_metric.cc
+++ b/src/metric/multiclass_metric.cc
@@ -5,4 +5,4 @@
 
 #if !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP)
 #include "multiclass_metric.cu"
-#endif  // !defined(XGBOOST_USE_CUDA)
+#endif  // !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP)
diff --git a/src/metric/multiclass_metric.cu b/src/metric/multiclass_metric.cu
index 4e7c87048..706c0135b 100644
--- a/src/metric/multiclass_metric.cu
+++ b/src/metric/multiclass_metric.cu
@@ -24,7 +24,7 @@
 #endif  // XGBOOST_USE_CUDA
 
 #if defined(XGBOOST_USE_HIP)
-#include <thrust/execution_policy.h>  // thrust::cuda::par
+#include <thrust/execution_policy.h>  // thrust::hip::par
 #include <thrust/functional.h>        // thrust::plus<>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/transform_reduce.h>
@@ -90,7 +90,7 @@ class MultiClassMetricsReduction {
     return res;
   }
 
-#if defined(XGBOOST_USE_CUDA)
+#if defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
 
   PackedReduceResult DeviceReduceMetrics(
       const HostDeviceVector<bst_float>& weights,
@@ -111,6 +111,8 @@ class MultiClassMetricsReduction {
     s_label_error[0] = 0;
 
     dh::XGBCachingDeviceAllocator<char> alloc;
+
+#if defined(XGBOOST_USE_CUDA)
     PackedReduceResult result = thrust::transform_reduce(
         thrust::cuda::par(alloc),
         begin, end,
@@ -128,12 +130,32 @@ class MultiClassMetricsReduction {
         },
         PackedReduceResult(),
         thrust::plus<PackedReduceResult>());
+#elif defined(XGBOOST_USE_HIP)
+    PackedReduceResult result = thrust::transform_reduce(
+        thrust::hip::par(alloc),
+        begin, end,
+        [=] XGBOOST_DEVICE(size_t idx) {
+          bst_float weight = is_null_weight ? 1.0f : s_weights[idx];
+          bst_float residue = 0;
+          auto label = static_cast<int>(s_labels[idx]);
+          if (label >= 0 && label < static_cast<int32_t>(n_class)) {
+            residue = EvalRowPolicy::EvalRow(
+                label, &s_preds[idx * n_class], n_class) * weight;
+          } else {
+            s_label_error[0] = label;
+          }
+          return PackedReduceResult{ residue, weight };
+        },
+        PackedReduceResult(),
+        thrust::plus<PackedReduceResult>());
+#endif
+
     CheckLabelError(s_label_error[0], n_class);
 
     return result;
   }
 
-#endif  // XGBOOST_USE_CUDA
+#endif  // XGBOOST_USE_CUDA || defined(XGBOOST_USE_HIP)
 
   PackedReduceResult Reduce(const Context& tparam, int device, size_t n_class,
                             const HostDeviceVector<bst_float>& weights,
@@ -145,25 +167,30 @@ class MultiClassMetricsReduction {
       result =
           CpuReduceMetrics(weights, labels, preds, n_class, tparam.Threads());
     }
-#if defined(XGBOOST_USE_CUDA)
+#if defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
     else {  // NOLINT
       device_ = tparam.gpu_id;
       preds.SetDevice(device_);
       labels.SetDevice(device_);
       weights.SetDevice(device_);
 
+#if defined(XGBOOST_USE_CUDA)
       dh::safe_cuda(cudaSetDevice(device_));
+#elif defined(XGBOOST_USE_HIP)
+      dh::safe_cuda(hipSetDevice(device_));
+#endif
+
       result = DeviceReduceMetrics(weights, labels, preds, n_class);
     }
-#endif  // defined(XGBOOST_USE_CUDA)
+#endif  // defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
     return result;
   }
 
  private:
-#if defined(XGBOOST_USE_CUDA)
+#if defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
   dh::PinnedMemory label_error_;
   int device_{-1};
-#endif  // defined(XGBOOST_USE_CUDA)
+#endif  // defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
 };
 
 /*!
diff --git a/src/metric/multiclass_metric.hip b/src/metric/multiclass_metric.hip
index e69de29bb..4689644c8 100644
--- a/src/metric/multiclass_metric.hip
+++ b/src/metric/multiclass_metric.hip
@@ -0,0 +1,4 @@
+
+#if defined(XGBOOST_USE_HIP)
+#include "multiclass_metric.cu"
+#endif  // defined(XGBOOST_USE_HIP)