From 1a75f433044b9a700249860fe3eca616c44d6bec Mon Sep 17 00:00:00 2001
From: Robert Maynard <robertjmaynard@gmail.com>
Date: Tue, 27 Jul 2021 08:05:33 -0400
Subject: [PATCH] Allow compilation with nvcc 11.4 (#7131)

* Use type aliases for discard iterators

* update to include host_vector as thrust 1.12 doesn't bring it in as a side-effect

* cub::DispatchRadixSort requires signed offset types
---
 src/common/device_helpers.cuh        | 8 ++++----
 src/data/ellpack_page.cu             | 5 +----
 src/metric/auc.cu                    | 5 +----
 src/predictor/gpu_predictor.cu       | 1 +
 src/tree/gpu_hist/row_partitioner.cu | 5 +----
 5 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index 5bb9129a4..f2e6ac629 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -1329,22 +1329,22 @@ void ArgSort(xgboost::common::Span<U> keys, xgboost::common::Span<IdxT> sorted_i
 
   if (accending) {
     void *d_temp_storage = nullptr;
-    safe_cuda((cub::DispatchRadixSort<false, KeyT, ValueT, size_t>::Dispatch(
+    safe_cuda((cub::DispatchRadixSort<false, KeyT, ValueT, std::ptrdiff_t>::Dispatch(
         d_temp_storage, bytes, d_keys, d_values, sorted_idx.size(), 0,
         sizeof(KeyT) * 8, false, nullptr, false)));
     TemporaryArray<char> storage(bytes);
     d_temp_storage = storage.data().get();
-    safe_cuda((cub::DispatchRadixSort<false, KeyT, ValueT, size_t>::Dispatch(
+    safe_cuda((cub::DispatchRadixSort<false, KeyT, ValueT, std::ptrdiff_t>::Dispatch(
         d_temp_storage, bytes, d_keys, d_values, sorted_idx.size(), 0,
         sizeof(KeyT) * 8, false, nullptr, false)));
   } else {
     void *d_temp_storage = nullptr;
-    safe_cuda((cub::DispatchRadixSort<true, KeyT, ValueT, size_t>::Dispatch(
+    safe_cuda((cub::DispatchRadixSort<true, KeyT, ValueT, std::ptrdiff_t>::Dispatch(
         d_temp_storage, bytes, d_keys, d_values, sorted_idx.size(), 0,
         sizeof(KeyT) * 8, false, nullptr, false)));
     TemporaryArray<char> storage(bytes);
     d_temp_storage = storage.data().get();
-    safe_cuda((cub::DispatchRadixSort<true, KeyT, ValueT, size_t>::Dispatch(
+    safe_cuda((cub::DispatchRadixSort<true, KeyT, ValueT, std::ptrdiff_t>::Dispatch(
         d_temp_storage, bytes, d_keys, d_values, sorted_idx.size(), 0,
         sizeof(KeyT) * 8, false, nullptr, false)));
   }
diff --git a/src/data/ellpack_page.cu b/src/data/ellpack_page.cu
index 2315700dc..97eccf682 100644
--- a/src/data/ellpack_page.cu
+++ b/src/data/ellpack_page.cu
@@ -185,10 +185,7 @@ struct TupleScanOp {
 
 // Change the value type of thrust discard iterator so we can use it with cub
 template <typename T>
-class TypedDiscard : public thrust::discard_iterator<T> {
- public:
-  using value_type = T;  // NOLINT
-};
+using TypedDiscard = thrust::discard_iterator<T>;
 
 // Here the data is already correctly ordered and simply needs to be compacted
 // to remove missing data
diff --git a/src/metric/auc.cu b/src/metric/auc.cu
index 708b424f9..2e33a2a5f 100644
--- a/src/metric/auc.cu
+++ b/src/metric/auc.cu
@@ -20,10 +20,7 @@ namespace xgboost {
 namespace metric {
 namespace {
 template <typename T>
-class Discard : public thrust::discard_iterator<T>  {
- public:
-  using value_type = T;  // NOLINT
-};
+using Discard = thrust::discard_iterator<T>;
 
 struct GetWeightOp {
   common::Span<float const> weights;
diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu
index fb9a10588..d1b9c8132 100644
--- a/src/predictor/gpu_predictor.cu
+++ b/src/predictor/gpu_predictor.cu
@@ -5,6 +5,7 @@
 #include <thrust/device_ptr.h>
 #include <thrust/device_vector.h>
 #include <thrust/fill.h>
+#include <thrust/host_vector.h>
 #include <GPUTreeShap/gpu_treeshap.h>
 #include <memory>
 
diff --git a/src/tree/gpu_hist/row_partitioner.cu b/src/tree/gpu_hist/row_partitioner.cu
index ebb3666e0..db9d878e5 100644
--- a/src/tree/gpu_hist/row_partitioner.cu
+++ b/src/tree/gpu_hist/row_partitioner.cu
@@ -61,10 +61,7 @@ struct WriteResultsFunctor {
 };
 
 // Change the value type of thrust discard iterator so we can use it with cub
-class DiscardOverload : public thrust::discard_iterator<IndexFlagTuple> {
- public:
-  using value_type = IndexFlagTuple;  // NOLINT
-};
+using DiscardOverload = thrust::discard_iterator<IndexFlagTuple>;
 
 // Implement partitioning via single scan operation using transform output to
 // write the result