Test/remove thrust vector usage (#11813)

This PR removes usage of `thrust::device_vector` from almost all of our tests. Since the construction of a device vector is not stream-ordered, we should be using `rmm::device_uvector` instead wherever possible. There is one remaining use of `thrust::device_vector`, but that is in an test explicitly verifying that `device_vector` can convert implicitly to a `device_span` so it's worth keeping that there.

I am working on automated tooling to detect any usage of stream 0 in tests as part of a push to prioritize stream-safety in libcudf, and this PR is a prerequisite to adding such tooling to our CI pipeline since at that point any test using stream 0 would fail. Since there is at least one test where I anticipate stream 0 will always be used (the one described above), I should be able to add specific tests to an allowlist as needed. It's an open question whether the added complexity required by the changes in this PR is a worthwhile tradeoff to be able to programmatically detect stream 0 usage. If reviewers feel that the additional complexity is too high, we can revert some (or all) of these changes and I can just plan for allowing stream 0 usage in all of the necessary tests. This PR demonstrates how we would go about removing it if we choose to do so, though.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Vukasin Milovanovic (https://github.com/vuule)
  - Tobias Ribizel (https://github.com/upsj)
  - Nghia Truong (https://github.com/ttnghia)

URL: #11813

cpp/include/cudf_test/tdigest_utilities.cuh

@@ -118,7 +118,7 @@ void tdigest_minmax_compare(cudf::tdigest::tdigest_column_view const& tdv,
   // verify min/max
   thrust::host_vector<device_span<T const>> h_spans;
   h_spans.push_back({input_values.begin<T>(), static_cast<size_t>(input_values.size())});
-  thrust::device_vector<device_span<T const>> spans(h_spans);
+  auto spans = cudf::detail::make_device_uvector_async(h_spans, cudf::default_stream_value);
 
   auto expected_min = cudf::make_fixed_width_column(
     data_type{type_id::FLOAT64}, spans.size(), mask_state::UNALLOCATED);

cpp/tests/join/conditional_join_tests.cu

@@ -26,10 +26,8 @@
 
 #include <rmm/exec_policy.hpp>
 
-#include <thrust/device_vector.h>
 #include <thrust/equal.h>
-#include <thrust/execution_policy.h>
-#include <thrust/pair.h>
+#include <thrust/iterator/counting_iterator.h>
 #include <thrust/sort.h>
 #include <thrust/transform.h>
 
@@ -127,6 +125,30 @@ gen_random_nullable_repeated_columns(unsigned int N = 10000, unsigned int num_re
                    std::pair(std::move(right), std::move(right_nulls)));
 }
 
+// `rmm::device_uvector<T>` requires that T be trivially copyable. `thrust::pair` does
+// not satisfy this requirement because it defines nontrivial copy/move
+// constructors. Therefore, we need a simple, trivially copyable pair-like
+// object. `index_pair` is a minimal implementation suitable for use in the
+// tests in this file.
+struct index_pair {
+  cudf::size_type first{};
+  cudf::size_type second{};
+  __device__ index_pair(){};
+  __device__ index_pair(cudf::size_type const& first, cudf::size_type const& second)
+    : first(first), second(second){};
+};
+
+__device__ inline bool operator<(const index_pair& lhs, const index_pair& rhs)
+{
+  if (lhs.first > rhs.first) return false;
+  return (lhs.first < rhs.first) || (lhs.second < rhs.second);
+}
+
+__device__ inline bool operator==(const index_pair& lhs, const index_pair& rhs)
+{
+  return lhs.first == rhs.first && lhs.second == rhs.second;
+}
+
 }  // namespace
 
 /**
@@ -253,26 +275,26 @@ struct ConditionalJoinPairReturnTest : public ConditionalJoinTest<T> {
    */
   void _compare_to_hash_join(PairJoinReturn const& result, PairJoinReturn const& reference)
   {
-    thrust::device_vector<thrust::pair<cudf::size_type, cudf::size_type>> result_pairs(
-      result.first->size());
-    thrust::device_vector<thrust::pair<cudf::size_type, cudf::size_type>> reference_pairs(
-      reference.first->size());
+    auto result_pairs =
+      rmm::device_uvector<index_pair>(result.first->size(), cudf::default_stream_value);
+    auto reference_pairs =
+      rmm::device_uvector<index_pair>(reference.first->size(), cudf::default_stream_value);
 
     thrust::transform(rmm::exec_policy(cudf::default_stream_value),
                       result.first->begin(),
                       result.first->end(),
                       result.second->begin(),
                       result_pairs.begin(),
                       [] __device__(cudf::size_type first, cudf::size_type second) {
-                        return thrust::make_pair(first, second);
+                        return index_pair{first, second};
                       });
     thrust::transform(rmm::exec_policy(cudf::default_stream_value),
                       reference.first->begin(),
                       reference.first->end(),
                       reference.second->begin(),
                       reference_pairs.begin(),
                       [] __device__(cudf::size_type first, cudf::size_type second) {
-                        return thrust::make_pair(first, second);
+                        return index_pair{first, second};
                       });
 
     thrust::sort(

cpp/tests/quantiles/tdigest_utilities.cu

@@ -51,23 +51,44 @@ void tdigest_sample_compare(cudf::tdigest::tdigest_column_view const& tdv,
   auto sampled_result_weight = cudf::make_fixed_width_column(
     data_type{type_id::FLOAT64}, h_expected.size(), mask_state::UNALLOCATED);
 
-  rmm::device_vector<expected_value> expected(h_expected.begin(), h_expected.end());
+  auto h_expected_src    = std::vector<size_type>(h_expected.size());
+  auto h_expected_mean   = std::vector<double>(h_expected.size());
+  auto h_expected_weight = std::vector<double>(h_expected.size());
+
+  {
+    auto iter = thrust::make_counting_iterator(0);
+    std::for_each_n(iter, h_expected.size(), [&](size_type const index) {
+      h_expected_src[index]    = thrust::get<0>(h_expected[index]);
+      h_expected_mean[index]   = thrust::get<1>(h_expected[index]);
+      h_expected_weight[index] = thrust::get<2>(h_expected[index]);
+    });
+  }
+
+  auto d_expected_src =
+    cudf::detail::make_device_uvector_async(h_expected_src, cudf::default_stream_value);
+  auto d_expected_mean =
+    cudf::detail::make_device_uvector_async(h_expected_mean, cudf::default_stream_value);
+  auto d_expected_weight =
+    cudf::detail::make_device_uvector_async(h_expected_weight, cudf::default_stream_value);
+
   auto iter = thrust::make_counting_iterator(0);
   thrust::for_each(
     rmm::exec_policy(cudf::default_stream_value),
     iter,
-    iter + expected.size(),
-    [expected            = expected.data().get(),
+    iter + h_expected.size(),
+    [expected_src_in     = d_expected_src.data(),
+     expected_mean_in    = d_expected_mean.data(),
+     expected_weight_in  = d_expected_weight.data(),
      expected_mean       = expected_mean->mutable_view().begin<double>(),
      expected_weight     = expected_weight->mutable_view().begin<double>(),
      result_mean         = result_mean.begin<double>(),
      result_weight       = result_weight.begin<double>(),
      sampled_result_mean = sampled_result_mean->mutable_view().begin<double>(),
      sampled_result_weight =
        sampled_result_weight->mutable_view().begin<double>()] __device__(size_type index) {
-      expected_mean[index]         = thrust::get<1>(expected[index]);
-      expected_weight[index]       = thrust::get<2>(expected[index]);
-      auto const src_index         = thrust::get<0>(expected[index]);
+      expected_mean[index]         = expected_mean_in[index];
+      expected_weight[index]       = expected_weight_in[index];
+      auto const src_index         = expected_src_in[index];
       sampled_result_mean[index]   = result_mean[src_index];
       sampled_result_weight[index] = result_weight[src_index];
     });

cpp/tests/utilities_tests/span_tests.cu

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/span.hpp>
 #include <io/utilities/hostdevice_vector.hpp>
@@ -234,17 +235,14 @@ __global__ void simple_device_kernel(device_span<bool> result) { result[0] = tru
 
 TEST(SpanTest, CanUseDeviceSpan)
 {
-  rmm::device_vector<bool> d_message = std::vector<bool>({false});
+  auto d_message =
+    cudf::detail::make_zeroed_device_uvector_async<bool>(1, cudf::default_stream_value);
 
-  auto d_span = device_span<bool>(d_message.data().get(), d_message.size());
+  auto d_span = device_span<bool>(d_message.data(), d_message.size());
 
   simple_device_kernel<<<1, 1, 0, cudf::default_stream_value.value()>>>(d_span);
 
-  cudaDeviceSynchronize();
-
-  thrust::host_vector<bool> h_message = d_message;
-
-  ASSERT_TRUE(h_message[0]);
+  ASSERT_TRUE(d_message.element(0, cudf::default_stream_value));
 }
 
 class MdSpanTest : public cudf::test::BaseFixture {