From a45005863b19e849db7ad9986325374e15d80fd4 Mon Sep 17 00:00:00 2001
From: amdsc21 <96135754+amdsc21@users.noreply.github.com>
Date: Wed, 8 Mar 2023 20:15:33 +0100
Subject: [PATCH] fix DispatchScan

---
 src/common/device_helpers.hip.h | 19 ++++++++++++++++++-
 src/data/ellpack_page.cu        | 19 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/common/device_helpers.hip.h b/src/common/device_helpers.hip.h
index 0452d6626..3ac3f6b6a 100644
--- a/src/common/device_helpers.hip.h
+++ b/src/common/device_helpers.hip.h
@@ -1,8 +1,10 @@
-#include "hip/hip_runtime.h"
 /**
  * Copyright 2017-2023 XGBoost contributors
  */
 #pragma once
+
+#include "hip/hip_runtime.h"
+
 #include <thrust/binary_search.h>  // thrust::upper_bound
 #include <thrust/device_malloc_allocator.h>
 #include <thrust/device_ptr.h>
@@ -22,8 +24,11 @@
 #include <algorithm>
 #include <chrono>
 #include <cstddef>  // for size_t
+
 #include <hipcub/hipcub.hpp>
 #include <hipcub/util_allocator.hpp>
+#include <rocprim/rocprim.hpp>
+
 #include <numeric>
 #include <sstream>
 #include <string>
@@ -1153,6 +1158,7 @@ template <typename InputIteratorT, typename OutputIteratorT, typename ScanOpT,
 void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
                    OffsetT num_items) {
   size_t bytes = 0;
+#if 0
 #if THRUST_MAJOR_VERSION >= 2
   safe_cuda((
       hipcub::DispatchScan<InputIteratorT, OutputIteratorT, ScanOpT, hipcub::NullType,
@@ -1165,7 +1171,14 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
                                            hipcub::NullType(), num_items, nullptr,
                                            false)));
 #endif
+#endif
+
+  safe_cuda((rocprim::inclusive_scan<InputIteratorT, OutputIteratorT, ScanOpT>(nullptr,
+                  bytes, d_in, d_out, num_items, scan_op)));
+
   TemporaryArray<char> storage(bytes);
+
+#if 0
 #if THRUST_MAJOR_VERSION >= 2
   safe_cuda((
       hipcub::DispatchScan<InputIteratorT, OutputIteratorT, ScanOpT, hipcub::NullType,
@@ -1179,6 +1192,10 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
                                            d_out, scan_op, hipcub::NullType(),
                                            num_items, nullptr, false)));
 #endif
+#endif
+
+  safe_cuda((rocprim::inclusive_scan<InputIteratorT, OutputIteratorT, ScanOpT>(
+                  storage.data().get(), bytes, d_in, d_out, num_items, scan_op)));
 }
 
 template <typename InIt, typename OutIt, typename Predicate>
diff --git a/src/data/ellpack_page.cu b/src/data/ellpack_page.cu
index 99e17d886..ed84d532f 100644
--- a/src/data/ellpack_page.cu
+++ b/src/data/ellpack_page.cu
@@ -13,6 +13,10 @@
 #include "gradient_index.h"
 #include "xgboost/data.h"
 
+#if defined(__HIP_PLATFORM_AMD__)
+#include <rocprim/rocprim.hpp>
+#endif
+
 namespace xgboost {
 
 EllpackPage::EllpackPage() : impl_{new EllpackPageImpl()} {}
@@ -235,6 +239,8 @@ void CopyDataToEllpack(const AdapterBatchT &batch,
   // Go one level down into cub::DeviceScan API to set OffsetT as 64 bit
   // So we don't crash on n > 2^31
   size_t temp_storage_bytes = 0;
+
+#if defined(__CUDACC__)
   using DispatchScan =
       cub::DispatchScan<decltype(key_value_index_iter), decltype(out),
                         TupleScanOp<Tuple>, cub::NullType, int64_t>;
@@ -257,6 +263,19 @@ void CopyDataToEllpack(const AdapterBatchT &batch,
                          key_value_index_iter, out, TupleScanOp<Tuple>(),
                          cub::NullType(), batch.Size(), nullptr, false);
 #endif
+
+#elif defined (__HIP_PLATFORM_AMD__)
+
+  rocprim::inclusive_scan<decltype(key_value_index_iter), decltype(out), TupleScanOp<Tuple>>
+      (nullptr, temp_storage_bytes, key_value_index_iter, out, batch.Size(), TupleScanOp<Tuple>());
+
+  dh::TemporaryArray<char> temp_storage(temp_storage_bytes);
+
+  rocprim::inclusive_scan<decltype(key_value_index_iter), decltype(out), TupleScanOp<Tuple>>
+      (temp_storage.data().get(), temp_storage_bytes, key_value_index_iter, out, batch.Size(),
+       TupleScanOp<Tuple>());
+
+#endif
 }
 
 void WriteNullValues(EllpackPageImpl* dst, int device_idx,