Make internal host_allocator stream ordered

cpp/include/cudf/detail/utilities/host_memory.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -13,13 +13,48 @@
 #include <cstddef>
 
 namespace cudf::detail {
+
+/**
+ * @brief Property tag for the CUDA native host memory pool.
+ *
+ * Advertised by `cuda_host_pinned_pool_memory_resource` (the pool backed by
+ * `cudaMallocFromPoolAsync` / `cudaFreeAsync`).  The typed resource ref below carries it through
+ * type erasure so callers can select the right pool without going through the public
+ * `get_pinned_memory_resource()` API.
+ */
+struct stream_ordered_host_accessible_t {};
+
+/**
+ * @brief Typed resource ref for the stream-ordered pinned host pool.
+ *
+ * Carrying `stream_ordered_host_accessible_t` in the template parameters means callers that
+ * construct `rmm_host_allocator` from this ref get a concrete type that still carries the
+ * property (before it is erased into `rmm::host_async_resource_ref`).  This is the return type
+ * of `get_stream_ordered_pinned_memory_resource()`.
+ */
+using stream_ordered_host_device_async_resource_ref =
+  cuda::mr::resource_ref<cuda::mr::host_accessible,
+                         cuda::mr::device_accessible,
+                         stream_ordered_host_accessible_t>;
+
 /**
  * @brief Get the memory resource to be used for pageable memory allocations.
  *
  * @return Reference to the pageable memory resource
  */
 CUDF_EXPORT rmm::host_async_resource_ref get_pageable_memory_resource();
 
+/**
+ * @brief Get the stream-ordered pinned memory resource.
+ *
+ * The underlying resource uses `cudaMallocFromPoolAsync` / `cudaFreeAsync` so both allocation and
+ * deallocation are truly stream-ordered.
+ *
+ * @return A `stream_ordered_host_device_async_resource_ref` backed by the default pinned pool
+ */
+CUDF_EXPORT stream_ordered_host_device_async_resource_ref
+get_stream_ordered_pinned_memory_resource();
+
 /**
  * @brief Get the allocator to be used for the host memory allocation.
  *
@@ -31,7 +66,7 @@ template <typename T>
 rmm_host_allocator<T> get_host_allocator(std::size_t size, rmm::cuda_stream_view stream)
 {
   if (size * sizeof(T) <= get_allocate_host_as_pinned_threshold()) {
-    return {get_pinned_memory_resource(), stream};
+    return {get_stream_ordered_pinned_memory_resource(), stream};
   }
   return {get_pageable_memory_resource(), stream};
 }

cpp/include/cudf/detail/utilities/host_vector.hpp

@@ -141,9 +141,6 @@ class rmm_host_allocator {
   {
     if (cnt > this->max_size()) { throw std::bad_alloc(); }  // end if
     auto const result = mr.allocate(stream, cnt * sizeof(value_type), alignof(value_type));
-    // Synchronize to ensure the memory is allocated before thrust::host_vector initialization
-    // TODO: replace thrust::host_vector with a type that does not require synchronization
-    stream.synchronize();
     return static_cast<pointer>(result);
   }
 

cpp/include/cudf/detail/utilities/vector_factories.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -494,15 +494,16 @@ host_vector<typename Container::value_type> make_host_vector(Container const& c,
 template <typename T>
 host_vector<T> make_empty_pinned_vector(size_t capacity, rmm::cuda_stream_view stream)
 {
-  auto result = host_vector<T>({cudf::get_pinned_memory_resource(), stream});
+  auto result =
+    host_vector<T>(rmm_host_allocator<T>{get_stream_ordered_pinned_memory_resource(), stream});
   result.reserve(capacity);
   return result;
 }
 
 /**
  * @brief Asynchronously construct a pinned `cudf::detail::host_vector` of the given size
  *
- * @note This function may not synchronize `stream` after the copy.
+ * @note This function does not synchronize `stream`.
  *
  * @tparam T The type of the vector data
  * @param size The number of elements in the created vector
@@ -512,7 +513,8 @@ host_vector<T> make_empty_pinned_vector(size_t capacity, rmm::cuda_stream_view s
 template <typename T>
 host_vector<T> make_pinned_vector_async(size_t size, rmm::cuda_stream_view stream)
 {
-  return host_vector<T>(size, {cudf::get_pinned_memory_resource(), stream});
+  return host_vector<T>(size,
+                        rmm_host_allocator<T>{get_stream_ordered_pinned_memory_resource(), stream});
 }
 
 /**

cpp/include/cudf_test/tdigest_utilities.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -128,10 +128,8 @@ void tdigest_minmax_compare(cudf::tdigest::tdigest_column_view const& tdv,
   auto expected_max = static_cast<double>(max_scalar->value());
 
   double tdv_min, tdv_max;
-  EXPECT_EQ(cudaMemcpy(&tdv_min, tdv.min_begin(), sizeof(double), cudaMemcpyDeviceToHost),
-            cudaSuccess);
-  EXPECT_EQ(cudaMemcpy(&tdv_max, tdv.max_begin(), sizeof(double), cudaMemcpyDeviceToHost),
-            cudaSuccess);
+  EXPECT_EQ(cudaMemcpy(&tdv_min, tdv.min_begin(), sizeof(double), cudaMemcpyDefault), cudaSuccess);
+  EXPECT_EQ(cudaMemcpy(&tdv_max, tdv.max_begin(), sizeof(double), cudaMemcpyDefault), cudaSuccess);
 
   EXPECT_EQ(tdv_min, expected_min);
   EXPECT_EQ(tdv_max, expected_max);

cpp/src/quantiles/tdigest/tdigest_aggregation.cu

@@ -1536,43 +1536,42 @@ std::unique_ptr<column> merge_tdigests(tdigest_column_view const& tdv,
     // if we will be at least partially using the CPU here, move the important values into pinned
     // and reference those instead.
     if (use_cpu_for_cluster_computation(num_groups)) {
-      auto pinned_mr = cudf::get_pinned_memory_resource();
-
-      rmm::device_uvector<size_type> _p_group_offsets(num_groups + 1, stream, pinned_mr);
+      rmm::device_uvector<size_type> d_p_group_offsets(num_groups + 1, stream);
       thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                    group_offsets,
-                   group_offsets + _p_group_offsets.size(),
-                   _p_group_offsets.begin());
-      cudf::device_span<size_type const> p_group_offsets(_p_group_offsets);
-
-      rmm::device_uvector<double> p_cumulative_weights(cumulative_weights, stream, pinned_mr);
-
-      rmm::device_uvector<size_type> p_tdigest_offsets(tdigest_offsets.size(), stream, pinned_mr);
-      thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
-                   tdigest_offsets.begin<size_type>(),
-                   tdigest_offsets.begin<size_type>() + p_tdigest_offsets.size(),
-                   p_tdigest_offsets.begin());
-
-      rmm::device_uvector<size_type> _p_group_labels(num_group_labels, stream, pinned_mr);
+                   group_offsets + d_p_group_offsets.size(),
+                   d_p_group_offsets.begin());
+      auto _p_group_offsets = cudf::detail::make_pinned_vector_async(d_p_group_offsets, stream);
+      cudf::device_span<size_type const> p_group_offsets(_p_group_offsets.data(),
+                                                         _p_group_offsets.size());
+
+      auto _p_tdigest_offsets = cudf::detail::make_pinned_vector_async<size_type>(
+        device_span<size_type const>(tdigest_offsets.data<size_type>(), tdigest_offsets.size()),
+        stream);
+      auto p_tdigest_offsets =
+        cuda::std::span<size_type const>{_p_tdigest_offsets.data(), _p_tdigest_offsets.size()};
+
+      rmm::device_uvector<size_type> d_p_group_labels(num_group_labels, stream);
       thrust::copy(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                    group_labels,
                    group_labels + num_group_labels,
-                   _p_group_labels.begin());
-      cudf::device_span<size_type const> p_group_labels(_p_group_labels);
+                   d_p_group_labels.begin());
+      auto _p_group_labels = cudf::detail::make_pinned_vector_async(d_p_group_labels, stream);
+      cudf::device_span<size_type const> p_group_labels(_p_group_labels.data(),
+                                                        _p_group_labels.size());
 
       stream.synchronize();
+      auto pinned_mr = cudf::get_pinned_memory_resource();
+      rmm::device_uvector<double> p_cumulative_weights(cumulative_weights, stream, pinned_mr);
       return generate_group_cluster_info(
         delta,
         num_groups,
         nearest_value_centroid_weights{
-          p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.begin()},
+          p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.data()},
         centroid_group_info{
-          p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.begin()},
+          p_cumulative_weights.begin(), p_group_offsets, p_tdigest_offsets.data()},
         cumulative_centroid_weight{
-          p_cumulative_weights.begin(),
-          p_group_labels,
-          p_group_offsets,
-          cuda::std::span<size_type const>{p_tdigest_offsets.begin(), p_tdigest_offsets.size()}},
+          p_cumulative_weights.begin(), p_group_labels, p_group_offsets, p_tdigest_offsets},
         has_nulls,
         stream,
         mr);