From fb296feb317604100a61aab8dceb78cd5ad44c9c Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Tue, 29 Mar 2022 10:39:35 -0700
Subject: [PATCH] Clean up Thrust includes. (#1011)

## Description

This PR cleans up some `#include`s for Thrust. This is meant to help ease the transition to Thrust 1.16 when that is updated in rapids-cmake.

## Context

I opened a PR https://github.com/rapidsai/cudf/pull/10489 that updates cuDF to Thrust 1.16. Notably, Thrust reduced the number of internal header inclusions:
> [#1572](https://github.com/NVIDIA/thrust/pull/1572) Removed several unnecessary header includes. Downstream projects may need to update their includes if they were relying on this behavior.

I spoke with @robertmaynard and he recommended making similar changes to clean up includes ("include what we use," in essence) to make sure we have compatibility with future versions of Thrust across all RAPIDS libraries.

It looks like rmm may be able to build with Thrust 1.16 even without these changes, but I think this changeset may help prevent future problems arising from inconsistency and reliance on `detail` headers.

Authors:
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Mark Harris (https://github.com/harrism)
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Conor Hoekstra (https://github.com/codereport)

URL: https://github.com/rapidsai/rmm/pull/1011
---
 benchmarks/device_uvector/device_uvector_bench.cu  | 4 +++-
 include/rmm/mr/device/thrust_allocator_adaptor.hpp | 3 ++-
 tests/mr/device/thrust_allocator_tests.cu          | 2 ++
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/benchmarks/device_uvector/device_uvector_bench.cu b/benchmarks/device_uvector/device_uvector_bench.cu
index 701e5fceb..454db81a5 100644
--- a/benchmarks/device_uvector/device_uvector_bench.cu
+++ b/benchmarks/device_uvector/device_uvector_bench.cu
@@ -15,7 +15,6 @@
  */
 
 #include "../synchronization/synchronization.hpp"
-#include "thrust/detail/raw_pointer_cast.h"
 
 #include <rmm/cuda_stream.hpp>
 #include <rmm/detail/error.hpp>
@@ -30,6 +29,9 @@
 
 #include <cuda_runtime_api.h>
 
+#include <thrust/device_vector.h>
+#include <thrust/memory.h>
+
 #include <cstdio>
 #include <type_traits>
 
diff --git a/include/rmm/mr/device/thrust_allocator_adaptor.hpp b/include/rmm/mr/device/thrust_allocator_adaptor.hpp
index 56e910801..c09c9f334 100644
--- a/include/rmm/mr/device/thrust_allocator_adaptor.hpp
+++ b/include/rmm/mr/device/thrust_allocator_adaptor.hpp
@@ -19,8 +19,9 @@
 #include <rmm/mr/device/device_memory_resource.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 
-#include <thrust/detail/type_traits/pointer_traits.h>
 #include <thrust/device_malloc_allocator.h>
+#include <thrust/device_ptr.h>
+#include <thrust/memory.h>
 
 namespace rmm::mr {
 /**
diff --git a/tests/mr/device/thrust_allocator_tests.cu b/tests/mr/device/thrust_allocator_tests.cu
index e0d467595..3a56a541b 100644
--- a/tests/mr/device/thrust_allocator_tests.cu
+++ b/tests/mr/device/thrust_allocator_tests.cu
@@ -23,6 +23,8 @@
 
 #include <gtest/gtest.h>
 
+#include <thrust/reduce.h>
+
 // explicit instantiation for test coverage purposes
 template class rmm::mr::thrust_allocator<int>;