From 53244bef6fd396a4b17d1eb4100eb949eb031456 Mon Sep 17 00:00:00 2001
From: amdsc21 <96135754+amdsc21@users.noreply.github.com>
Date: Fri, 10 Mar 2023 03:38:09 +0100
Subject: [PATCH] finish simple_dmatrix.cu

---
 src/data/simple_dmatrix.cu  |  5 +++++
 src/data/simple_dmatrix.cuh | 23 +++++++++++++++++++++--
 src/data/simple_dmatrix.hip |  4 ++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/data/simple_dmatrix.cu b/src/data/simple_dmatrix.cu
index 64f308b8c..421e14575 100644
--- a/src/data/simple_dmatrix.cu
+++ b/src/data/simple_dmatrix.cu
@@ -19,7 +19,12 @@ SimpleDMatrix::SimpleDMatrix(AdapterT* adapter, float missing, int32_t /*nthread
   auto device = (adapter->DeviceIdx() < 0 || adapter->NumRows() == 0) ? dh::CurrentDevice()
                                                                       : adapter->DeviceIdx();
   CHECK_GE(device, 0);
+
+#if defined(XGBOOST_USE_CUDA)
   dh::safe_cuda(cudaSetDevice(device));
+#elif defined(XGBOOST_USE_HIP)
+  dh::safe_cuda(hipSetDevice(device));
+#endif
 
   CHECK(adapter->NumRows() != kAdapterUnknownSize);
   CHECK(adapter->NumColumns() != kAdapterUnknownSize);
diff --git a/src/data/simple_dmatrix.cuh b/src/data/simple_dmatrix.cuh
index f3d4d953f..961e2d5d0 100644
--- a/src/data/simple_dmatrix.cuh
+++ b/src/data/simple_dmatrix.cuh
@@ -9,19 +9,38 @@
 #include <thrust/scan.h>
 #include <thrust/execution_policy.h>
 #include "device_adapter.cuh"
+
+#if defined(XGBOOST_USE_CUDA)
 #include "../common/device_helpers.cuh"
+#elif defined(XGBOOST_USE_HIP)
+#include "../common/device_helpers.hip.h"
+#endif
 
 namespace xgboost {
 namespace data {
 
+#if defined(XGBOOST_USE_CUDA)
 template <typename AdapterBatchT>
 struct COOToEntryOp {
   AdapterBatchT batch;
+
   __device__ Entry operator()(size_t idx) {
     const auto& e = batch.GetElement(idx);
     return Entry(e.column_idx, e.value);
   }
 };
+#elif defined(XGBOOST_USE_HIP)
+template <typename AdapterBatchT>
+struct COOToEntryOp : thrust::unary_function<size_t, Entry> {
+  AdapterBatchT batch;
+  COOToEntryOp(AdapterBatchT batch): batch(batch) {};
+
+  __device__ Entry operator()(size_t idx) {
+    const auto& e = batch.GetElement(idx);
+    return Entry(e.column_idx, e.value);
+  }
+};
+#endif
 
 // Here the data is already correctly ordered and simply needs to be compacted
 // to remove missing data
@@ -44,7 +63,7 @@ void CountRowOffsets(const AdapterBatchT& batch, common::Span<bst_row_t> offset,
 
 #if defined(XGBOOST_USE_HIP)
   dh::safe_cuda(hipSetDevice(device_idx));
-#else
+#elif defined(XGBOOST_USE_CUDA)
   dh::safe_cuda(cudaSetDevice(device_idx));
 #endif
 
@@ -66,7 +85,7 @@ void CountRowOffsets(const AdapterBatchT& batch, common::Span<bst_row_t> offset,
       thrust::device_pointer_cast(offset.data()),
       thrust::device_pointer_cast(offset.data() + offset.size()),
       thrust::device_pointer_cast(offset.data()));
-#else
+#elif defined(XGBOOST_USE_CUDA)
   thrust::exclusive_scan(thrust::cuda::par(alloc),
       thrust::device_pointer_cast(offset.data()),
       thrust::device_pointer_cast(offset.data() + offset.size()),
diff --git a/src/data/simple_dmatrix.hip b/src/data/simple_dmatrix.hip
index e69de29bb..9be8187e1 100644
--- a/src/data/simple_dmatrix.hip
+++ b/src/data/simple_dmatrix.hip
@@ -0,0 +1,4 @@
+
+#if defined(XGBOOST_USE_HIP)
+#include "simple_dmatrix.cu"
+#endif