diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 27d3289e75..ca4e7fcd90 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -136,7 +136,7 @@ extern "C" cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream) { return cuvs::core::translate_exceptions([=] { auto res_ptr = reinterpret_cast(res); - raft::resource::set_cuda_stream(*res_ptr, static_cast(stream)); + raft::resource::set_cuda_stream(*res_ptr, static_cast(stream)); }); } @@ -144,7 +144,7 @@ extern "C" cuvsError_t cuvsStreamGet(cuvsResources_t res, cudaStream_t* stream) { return cuvs::core::translate_exceptions([=] { auto res_ptr = reinterpret_cast(res); - *stream = raft::resource::get_cuda_stream(*res_ptr); + *stream = raft::resource::get_cuda_stream(*res_ptr).get(); }); } diff --git a/c/src/neighbors/nn_descent.cpp b/c/src/neighbors/nn_descent.cpp index 708056144a..3a0f5f3845 100644 --- a/c/src/neighbors/nn_descent.cpp +++ b/c/src/neighbors/nn_descent.cpp @@ -107,7 +107,7 @@ void _get_distances(cuvsResources_t res, cuvsNNDescentIndex_t index, DLManagedTe src->data_handle(), dst.extent(0) * dst.extent(1) * sizeof(float), cudaMemcpyDefault, - raft::resource::get_cuda_stream(*res_ptr)); + raft::resource::get_cuda_stream(*res_ptr).get()); } else { RAFT_FAIL("Unsupported nn-descent index dtype: %d and bits: %d", dtype.code, dtype.bits); diff --git a/c/tests/neighbors/ann_ivf_sq_c.cu b/c/tests/neighbors/ann_ivf_sq_c.cu index 42c1b29999..b2e6496448 100644 --- a/c/tests/neighbors/ann_ivf_sq_c.cu +++ b/c/tests/neighbors/ann_ivf_sq_c.cu @@ -108,7 +108,7 @@ TEST(IvfSqC, BuildSearch) cuvsResources_t res; cuvsResourcesCreate(&res); - cuvsStreamSet(res, stream); + cuvsStreamSet(res, stream.get()); run_ivf_sq(res, n_rows, diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h index 1a276e8cc8..64e43412dd 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h @@ -122,8 +122,8 @@ class configured_raft_resources { */ explicit configured_raft_resources(const std::shared_ptr& shared_res) : shared_res_{shared_res}, - res_{std::make_unique( - rmm::cuda_stream_view(get_stream_from_global_pool()))} + res_{ + std::make_unique(cuda::stream_ref(get_stream_from_global_pool()))} { raft::resource::set_large_workspace_resource( *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()}); diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 1b943afe30..4ad9aea688 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1168,7 +1168,7 @@ auto make_device_dense_row_major_dataset_from_src(raft::resources const& res, RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), 0, out_array.size() * sizeof(ValueT), - raft::resource::get_cuda_stream(res))); + raft::resource::get_cuda_stream(res).get())); raft::copy_matrix(out_array.data_handle(), target_stride, src.data_handle(), diff --git a/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh b/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh index d5c43aaa34..3aa7b78bb9 100644 --- a/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh +++ b/cpp/internal/cuvs_internal/neighbors/refine_helper.cuh @@ -127,7 +127,7 @@ class RefineHelper { public: RefineInputs p; const raft::resources& handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; raft::device_matrix dataset; raft::device_matrix queries; diff --git a/cpp/src/cluster/detail/agglomerative.cuh b/cpp/src/cluster/detail/agglomerative.cuh index bc920f7701..a664bffd04 100644 --- a/cpp/src/cluster/detail/agglomerative.cuh +++ b/cpp/src/cluster/detail/agglomerative.cuh @@ -280,7 +280,7 @@ void extract_flattened_clusters(raft::resources const& handle, rmm::device_uvector levels(n_vertices, stream); value_idx n_blocks = raft::ceildiv(n_vertices, (value_idx)tpb); - write_levels_kernel<<>>(children, levels.data(), n_vertices); + write_levels_kernel<<>>(children, levels.data(), n_vertices); /** * Step 1: Find label roots: * @@ -323,7 +323,7 @@ void extract_flattened_clusters(raft::resources const& handle, */ value_idx cut_level = (n_edges / 2) - (n_clusters - 1); - inherit_labels<<>>( + inherit_labels<<>>( children, levels.data(), n_leaves, tmp_labels.data(), cut_level, n_vertices); // copy tmp labels to actual labels diff --git a/cpp/src/cluster/detail/connectivities.cuh b/cpp/src/cluster/detail/connectivities.cuh index 1737eead12..6d39e621e7 100644 --- a/cpp/src/cluster/detail/connectivities.cuh +++ b/cpp/src/cluster/detail/connectivities.cuh @@ -61,7 +61,7 @@ struct distance_graph_impl { auto thrust_policy = raft::resource::get_thrust_policy(handle); // Need to symmetrize knn into undirected graph - raft::sparse::COO knn_graph_coo(stream); + raft::sparse::COO knn_graph_coo(stream.get()); auto X_view = raft::make_device_matrix_view(X, m, n); cuvs::neighbors::detail::knn_graph( @@ -92,7 +92,7 @@ struct distance_graph_impl { raft::make_const_mdspan(vals_in_view)); raft::sparse::convert::sorted_coo_to_csr( - knn_graph_coo.rows(), knn_graph_coo.nnz, indptr.data(), m + 1, stream); + knn_graph_coo.rows(), knn_graph_coo.nnz, indptr.data(), m + 1, stream.get()); // TODO: Wouldn't need to copy here if we could compute knn // graph directly on the device uvectors @@ -140,7 +140,7 @@ void pairwise_distances(const raft::resources& handle, value_idx nnz = m * m; value_idx blocks = raft::ceildiv(nnz, (value_idx)256); - fill_indices2<<>>(indices, m, nnz); + fill_indices2<<>>(indices, m, nnz); raft::linalg::map_offset(handle, raft::make_device_vector_view(indptr, m), diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..be18152b03 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -95,7 +95,7 @@ void kmeansPlusPlus(raft::resources const& handle, rmm::device_uvector& workspace) { raft::common::nvtx::range fun_scope("kmeansPlusPlus"); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = params.n_clusters; @@ -309,7 +309,7 @@ void initScalableKMeansPlusPlus(raft::resources const& handle, { raft::common::nvtx::range fun_scope( "initScalableKMeansPlusPlus"); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = params.n_clusters; @@ -573,7 +573,7 @@ void kmeans_fit( auto n_features = X.extent(1); auto n_clusters = pams.n_clusters; auto metric = pams.metric; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); if (sample_weight.has_value()) RAFT_EXPECTS(sample_weight.value().extent(0) == n_samples, @@ -1035,7 +1035,7 @@ void kmeans_predict(raft::resources const& handle, raft::common::nvtx::range fun_scope("kmeans_predict"); auto n_samples = X.extent(0); auto n_features = X.extent(1); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); // Check that parameters are valid if (sample_weight.has_value()) RAFT_EXPECTS(sample_weight.value().extent(0) == n_samples, @@ -1186,7 +1186,7 @@ void kmeans_transform(raft::resources const& handle, "kmeans only supports L2Expanded or L2SqrtExpanded distance metrics."); raft::common::nvtx::range fun_scope("kmeans_transform"); raft::default_logger().set_level(pams.verbosity); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = pams.n_clusters; diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 272d45f13d..8149f5e0ed 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -142,7 +142,7 @@ inline std::enable_if_t> predict_core( &beta, distances.data(), n_clusters, - stream); + stream.get()); auto distances_const_view = raft::make_device_matrix_view( distances.data(), n_rows, n_clusters); @@ -286,13 +286,29 @@ void calc_centers_and_sizes(const raft::resources& handle, // Apply mapping only when the data and math types are different. if constexpr (std::is_same_v) { - raft::linalg::reduce_rows_by_key( - dataset, dim, labels, nullptr, n_rows, dim, n_clusters, centers, stream, reset_counters); + raft::linalg::reduce_rows_by_key(dataset, + dim, + labels, + nullptr, + n_rows, + dim, + n_clusters, + centers, + stream.get(), + reset_counters); } else { // todo(lsugy): use iterator from KV output of fusedL2NN thrust::transform_iterator mapping_itr(dataset, mapping_op); - raft::linalg::reduce_rows_by_key( - mapping_itr, dim, labels, nullptr, n_rows, dim, n_clusters, centers, stream, reset_counters); + raft::linalg::reduce_rows_by_key(mapping_itr, + dim, + labels, + nullptr, + n_rows, + dim, + n_clusters, + centers, + stream.get(), + reset_counters); } // Compute weight of each cluster @@ -689,38 +705,39 @@ auto adjust_centers(const raft::resources& handle, search_count.set_value_to_zero_async(stream); const dim3 grid_dim(raft::ceildiv(n_clusters, static_cast(kBlockDimY)), 1, 1); adjust_centers_random_donor_kernel - <<>>(centers, - n_clusters, - dim, - dataset, - n_rows, - labels, - cluster_sizes, - lower_threshold, - static_cast(n_rows / n_clusters), - centroid_offset, - ofst, - search_count.data(), - update_count.data(), - mapping_op); + <<>>(centers, + n_clusters, + dim, + dataset, + n_rows, + labels, + cluster_sizes, + lower_threshold, + static_cast(n_rows / n_clusters), + centroid_offset, + ofst, + search_count.data(), + update_count.data(), + mapping_op); return update_count.value(stream) > 0; // NB: rmm scalar performs the sync } raft::update_device(receiver_clusters.data(), host_receiver_clusters.data(), n_pairs, stream); raft::update_device(donor_clusters.data(), host_donor_clusters.data(), n_pairs, stream); const dim3 grid_dim(raft::ceildiv(n_pairs, static_cast(kBlockDimY)), 1, 1); - adjust_centers_kernel<<>>(centers, - n_pairs, - dim, - dataset, - n_rows, - labels, - receiver_clusters.data(), - donor_clusters.data(), - centroid_offset, - ofst, - update_count.data(), - mapping_op); + adjust_centers_kernel + <<>>(centers, + n_pairs, + dim, + dataset, + n_rows, + labels, + receiver_clusters.data(), + donor_clusters.data(), + centroid_offset, + ofst, + update_count.data(), + mapping_op); auto n_updates = update_count.value(stream); // NB: rmm scalar performs the sync RAFT_EXPECTS(n_updates == n_pairs, "Balanced k-means failed to update all adjusted centers"); return n_updates > 0; @@ -1068,7 +1085,7 @@ auto build_fine_clusters(const raft::resources& handle, } thrust::transform_iterator mapping_itr(dataset_mptr, mapping_op); - raft::matrix::gather(mapping_itr, dim, n_rows, mc_trainset_ids, k, mc_trainset, stream); + raft::matrix::gather(mapping_itr, dim, n_rows, mc_trainset_ids, k, mc_trainset, stream.get()); if (params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded) { diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ab3ef0a05a..5d6696b83e 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -125,7 +125,7 @@ void countLabels(raft::resources const& handle, IndexT n_clusters, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); // CUB::DeviceHistogram requires a signed index type typedef typename std::make_signed_t CubIndexT; @@ -177,7 +177,7 @@ void weightSum( if constexpr (raft::is_device_mdspan_v) { raft::linalg::mapThenSumReduce( - d_wt_sum.data_handle(), n_samples, raft::identity_op{}, stream, weight.data_handle()); + d_wt_sum.data_handle(), n_samples, raft::identity_op{}, stream.get(), weight.data_handle()); if (check_positive) { raft::copy(&wt_sum_h, d_wt_sum.data_handle(), 1, stream); raft::resource::sync_stream(handle); @@ -219,7 +219,7 @@ void computeClusterCost(raft::resources const& handle, MainOpT main_op, ReductionOpT reduction_op) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cuda::transform_iterator itr(minClusterDistance.data_handle(), main_op); @@ -254,7 +254,7 @@ void sampleCentroids(raft::resources const& handle, rmm::device_uvector& inRankCp, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_local_samples = X.extent(0); auto n_features = X.extent(1); @@ -354,7 +354,7 @@ void shuffleAndGather(raft::resources const& handle, uint32_t n_samples_to_gather, uint64_t seed) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = in.extent(0); auto n_features = in.extent(1); @@ -455,7 +455,7 @@ void countSamplesInCluster(raft::resources const& handle, rmm::device_uvector& workspace, raft::device_vector_view sampleCountInCluster) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); @@ -534,7 +534,7 @@ void compute_centroid_adjustments( rmm::device_uvector& workspace, bool reset_sums = true) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); workspace.resize(n_samples, stream); @@ -582,7 +582,7 @@ void finalize_centroids(raft::resources const& handle, raft::device_matrix_view old_centroids, raft::device_matrix_view new_centroids) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); raft::linalg::matrix_vector_op(handle, raft::make_const_mdspan(centroid_sums), @@ -617,7 +617,7 @@ void compute_centroid_shift(raft::resources const& handle, raft::device_matrix_view new_centroids, raft::device_scalar_view sqrd_norm_out) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); raft::linalg::mapThenSumReduce(sqrd_norm_out.data_handle(), old_centroids.size(), raft::sqdiff_op{}, @@ -702,7 +702,7 @@ void process_batch( raft::device_scalar_view clustering_cost, rmm::device_uvector& batch_workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); minClusterAndDistanceCompute(handle, batch_data, diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index ee3cc3cdfd..344f0bf5e5 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -27,7 +27,7 @@ void minClusterAndDistanceCompute( int batch_centroids, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); @@ -242,7 +242,7 @@ void minClusterDistanceCompute(raft::resources const& handle, int batch_centroids, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto n_samples = X.extent(0); auto n_features = X.extent(1); auto n_clusters = centroids.extent(0); diff --git a/cpp/src/cluster/detail/mst.cuh b/cpp/src/cluster/detail/mst.cuh index bbd74022ba..f2b53f01f8 100644 --- a/cpp/src/cluster/detail/mst.cuh +++ b/cpp/src/cluster/detail/mst.cuh @@ -76,7 +76,7 @@ void connect_knn_graph( { auto stream = raft::resource::get_cuda_stream(handle); - raft::sparse::COO connected_edges(stream); + raft::sparse::COO connected_edges(stream.get()); // default row and column batch sizes are chosen for computing cross component nearest neighbors. // Reference: PR #1445 @@ -95,7 +95,7 @@ void connect_knn_graph( rmm::device_uvector indptr2(m + 1, stream); raft::sparse::convert::sorted_coo_to_csr( - connected_edges.rows(), connected_edges.nnz, indptr2.data(), m + 1, stream); + connected_edges.rows(), connected_edges.nnz, indptr2.data(), m + 1, stream.get()); // On the second call, we hand the MST the original colors // and the new set of edges and let it restart the optimization process @@ -107,11 +107,11 @@ void connect_knn_graph( m, connected_edges.nnz, color, - stream, + stream.get(), false, false); - merge_msts(msf, new_mst, stream); + merge_msts(msf, new_mst, stream.get()); } /** @@ -147,10 +147,10 @@ void connect_knn_graph( "FixConnectivitiesRedOp"); auto stream = raft::resource::get_cuda_stream(handle); - int n_components = get_n_components(color, m, stream); + int n_components = get_n_components(color, m, stream.get()); rmm::device_uvector d_color_remapped(m, stream); - raft::label::make_monotonic(d_color_remapped.data(), color, m, stream, true); + raft::label::make_monotonic(d_color_remapped.data(), color, m, stream.get(), true); std::vector h_color(m); raft::copy(handle, @@ -251,7 +251,7 @@ void connect_knn_graph( rmm::device_uvector indptr2(m + 1, stream); raft::sparse::convert::sorted_coo_to_csr( - device_u_indices.data_handle(), new_nnz, indptr2.data(), m + 1, stream); + device_u_indices.data_handle(), new_nnz, indptr2.data(), m + 1, stream.get()); // On the second call, we hand the MST the original colors // and the new set of edges and let it restart the optimization process @@ -263,11 +263,11 @@ void connect_knn_graph( m, new_nnz, color, - stream, + stream.get(), false, false); - merge_msts(msf, new_mst, stream); + merge_msts(msf, new_mst, stream.get()); } /** @@ -317,10 +317,10 @@ void build_sorted_mst( // We want to have MST initialize colors on first call. auto mst_coo = raft::sparse::solver::mst( - handle, indptr, indices, pw_dists, (value_idx)m, nnz, color, stream, false, true); + handle, indptr, indices, pw_dists, (value_idx)m, nnz, color, stream.get(), false, true); int iters = 1; - int n_components = cuvs::sparse::neighbors::get_n_components(color, m, stream); + int n_components = cuvs::sparse::neighbors::get_n_components(color, m, stream.get()); bool data_on_device = raft::memory_type_from_pointer(X) != raft::memory_type::host; @@ -348,7 +348,7 @@ void build_sorted_mst( iters++; - n_components = cuvs::sparse::neighbors::get_n_components(color, m, stream); + n_components = cuvs::sparse::neighbors::get_n_components(color, m, stream.get()); } /** @@ -372,11 +372,11 @@ void build_sorted_mst( max_iter); raft::sparse::op::coo_sort_by_weight( - mst_coo.src.data(), mst_coo.dst.data(), mst_coo.weights.data(), mst_coo.n_edges, stream); + mst_coo.src.data(), mst_coo.dst.data(), mst_coo.weights.data(), mst_coo.n_edges, stream.get()); - raft::copy_async(mst_src, mst_coo.src.data(), mst_coo.n_edges, stream); - raft::copy_async(mst_dst, mst_coo.dst.data(), mst_coo.n_edges, stream); - raft::copy_async(mst_weight, mst_coo.weights.data(), mst_coo.n_edges, stream); + raft::copy_async(mst_src, mst_coo.src.data(), mst_coo.n_edges, stream.get()); + raft::copy_async(mst_dst, mst_coo.dst.data(), mst_coo.n_edges, stream.get()); + raft::copy_async(mst_weight, mst_coo.weights.data(), mst_coo.n_edges, stream.get()); } }; // namespace cuvs::cluster::agglomerative::detail diff --git a/cpp/src/cluster/detail/single_linkage.cuh b/cpp/src/cluster/detail/single_linkage.cuh index dc0a5dca75..9adc295771 100644 --- a/cpp/src/cluster/detail/single_linkage.cuh +++ b/cpp/src/cluster/detail/single_linkage.cuh @@ -76,7 +76,7 @@ void build_mr_linkage( handle, all_neighbors_p, X, inds.view(), dists.view(), core_dists, alpha); // allocate memory after all neighbors build - mr_coo.emplace(stream, min_samples * m * 2); + mr_coo.emplace(stream.get(), min_samples * m * 2); // self-loops get max distance auto coo_rows = raft::make_device_vector(handle, min_samples * m); raft::linalg::map_offset(handle, coo_rows.view(), raft::div_const_op(min_samples)); @@ -92,7 +92,7 @@ void build_mr_linkage( } // scope to drop inds and dists matrices early auto mr_indptr = raft::make_device_vector(handle, m + 1); raft::sparse::convert::sorted_coo_to_csr( - mr_coo.value().rows(), mr_coo.value().nnz, mr_indptr.data_handle(), m + 1, stream); + mr_coo.value().rows(), mr_coo.value().nnz, mr_indptr.data_handle(), m + 1, stream.get()); auto rows_view = raft::make_device_vector_view(mr_coo.value().rows(), mr_coo.value().nnz); diff --git a/cpp/src/cluster/detail/spectral.cuh b/cpp/src/cluster/detail/spectral.cuh index 52513afe26..9106ac426d 100644 --- a/cpp/src/cluster/detail/spectral.cuh +++ b/cpp/src/cluster/detail/spectral.cuh @@ -49,7 +49,7 @@ void fit_predict(raft::resources const& handle, embedding_row_major.data_handle(), n_samples, config.n_components, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); cuvs::cluster::kmeans::fit_predict(handle, kmeans_config, diff --git a/cpp/src/cluster/gmm_impl.cuh b/cpp/src/cluster/gmm_impl.cuh index 1f52a2eba0..06769dcbb1 100644 --- a/cpp/src/cluster/gmm_impl.cuh +++ b/cpp/src/cluster/gmm_impl.cuh @@ -218,7 +218,7 @@ void e_step(raft::resources const& handle, T* resp, T* log_prob_norm) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); covariance_type ct = params.cov_type; if (ct == covariance_type::FULL || (ct == covariance_type::TIED && d > 128)) { @@ -392,7 +392,7 @@ struct MStepWorkspace { dB_ptrs(0, raft::resource::get_cuda_stream(handle)), lwork(0) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); covariance_type ct = params.cov_type; ones.resize(n, stream); thrust::fill(thrust::cuda::par.on(stream), ones.data(), ones.data() + n, T(1)); @@ -441,7 +441,7 @@ void m_accumulate(raft::resources const& handle, MStepWorkspace& ws, T beta) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); covariance_type ct = params.cov_type; T one = T(1); @@ -491,7 +491,7 @@ void m_finalize(raft::resources const& handle, T* means, T* covariances) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); covariance_type ct = params.cov_type; T one = T(1), zero = T(0); T eps = std::numeric_limits::epsilon(); @@ -563,7 +563,7 @@ void m_cov_full_pass(raft::resources const& handle, MStepWorkspace& ws, T beta) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); T one = T(1); int threads = 256; int blocks = (int)(((size_t)n * d + threads - 1) / threads); @@ -598,7 +598,7 @@ void m_finalize_cov_full(raft::resources const& handle, MStepWorkspace& ws, T* covariances) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); T eps = std::numeric_limits::epsilon(); detail::m_step_finalize_cov_full_kernel <<>>(ws.N_k.data(), covariances, T(params.reg_covar), eps, d, K); @@ -611,7 +611,7 @@ template void precision_cholesky_one( raft::resources const& handle, const T* cov, T* prec_chol, int d, MStepWorkspace& ws) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cusolverDnHandle_t solver = raft::resource::get_cusolver_dn_handle(handle); cublasHandle_t cublas = raft::resource::get_cublas_handle(handle); @@ -661,7 +661,7 @@ void precision_cholesky_full_batched(raft::resources const& handle, int K, MStepWorkspace& ws) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cusolverDnHandle_t solver = raft::resource::get_cusolver_dn_handle(handle); cublasHandle_t cublas = raft::resource::get_cublas_handle(handle); @@ -715,7 +715,7 @@ void update_precisions(raft::resources const& handle, T* log_det, MStepWorkspace& ws) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); covariance_type ct = params.cov_type; if (ct == covariance_type::FULL) { precision_cholesky_full_batched(handle, covariances, prec_chol, d, K, ws); @@ -763,7 +763,7 @@ void compute_precisions(raft::resources const& handle, int K, T* precisions) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); covariance_type ct = params.cov_type; T one = T(1), zero = T(0); if (ct == covariance_type::FULL) { @@ -797,7 +797,7 @@ void kmeans_assign(raft::resources const& handle, uint64_t init_seed, int* labels_out) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cuvs::cluster::kmeans::params kp; kp.n_clusters = K; // KMeansPlusPlus: seeding labels only (max_iter=0, no Lloyd); KMeans: full Lloyd. @@ -834,7 +834,7 @@ void kmeans_assign(raft::resources const& handle, template T mean_device(raft::resources const& handle, const T* v, int n) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); rmm::device_scalar d_sum(stream); raft::linalg::reduce( handle, @@ -864,7 +864,7 @@ void fit_impl(raft::resources const& handle, bool& converged, bool warm_start) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); int K = params.n_components; covariance_type ct = params.cov_type; size_t cn = cov_elems(ct, d, K); @@ -1118,7 +1118,7 @@ void infer(raft::resources const& handle, T* resp, T* log_prob_norm) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); int K = params.n_components; covariance_type ct = params.cov_type; @@ -1172,7 +1172,7 @@ void launch_estep_tiled(raft::resources const& handle, int* labels, T* log_prob = nullptr) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); constexpr int FEAT = 32; constexpr int CELL = (sizeof(T) == 4) ? 64 : 32; constexpr int TPB = 256; @@ -1204,7 +1204,7 @@ void fused_score(raft::resources const& handle, T* log_prob_norm, int* labels) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); int K = params.n_components; covariance_type ct = params.cov_type; @@ -1329,7 +1329,7 @@ void predict_impl(raft::resources const& handle, const T* precisions_chol, int* labels) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); rmm::device_uvector lpn(n, stream); fused_score(handle, params, X, n, d, weights, means, precisions_chol, lpn.data(), labels); raft::resource::sync_stream(handle); @@ -1346,7 +1346,7 @@ void predict_proba_impl(raft::resources const& handle, const T* precisions_chol, T* resp) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); rmm::device_uvector lpn(n, stream); infer(handle, params, X, n, d, weights, means, precisions_chol, resp, lpn.data()); raft::resource::sync_stream(handle); diff --git a/cpp/src/core/mnmg_comms.cuh b/cpp/src/core/mnmg_comms.cuh index 02b90ce381..6238368502 100644 --- a/cpp/src/core/mnmg_comms.cuh +++ b/cpp/src/core/mnmg_comms.cuh @@ -59,7 +59,7 @@ class mnmg_comms { : dev_res_(dev_res), use_nccl_(use_nccl), nccl_comm_(nccl_comm), - stream_(raft::resource::get_cuda_stream(dev_res_)) + stream_(raft::resource::get_cuda_stream(dev_res_).get()) { } diff --git a/cpp/src/distance/detail/compress_to_bits.cuh b/cpp/src/distance/detail/compress_to_bits.cuh index a4917156b4..89e29419ba 100644 --- a/cpp/src/distance/detail/compress_to_bits.cuh +++ b/cpp/src/distance/detail/compress_to_bits.cuh @@ -106,7 +106,7 @@ void compress_to_bits(raft::resources const& handle, dim3 grid(num_SMs * blocks_per_sm); dim3 block(128); - compress_to_bits_kernel<<>>(in, out); + compress_to_bits_kernel<<>>(in, out); RAFT_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/distance/detail/distance.cuh b/cpp/src/distance/detail/distance.cuh index a8b73d58e3..a1e5849310 100644 --- a/cpp/src/distance/detail/distance.cuh +++ b/cpp/src/distance/detail/distance.cuh @@ -95,7 +95,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); } @@ -118,7 +118,7 @@ void distance_impl(raft::resources const& handle, ASSERT(!(worksize < 2 * (m + n) * sizeof(AccT)), "workspace size error"); ASSERT(workspace != nullptr, "workspace is null"); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); AccT* x_norm = workspace; AccT* y_norm = workspace; @@ -232,7 +232,7 @@ void distance_impl(raft::resources const& handle, ASSERT(!(worksize < (m + n) * sizeof(AccT)), "workspace size error"); ASSERT(workspace != nullptr, "workspace is null"); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); OutT* x_norm = reinterpret_cast(workspace); OutT* y_norm = reinterpret_cast(workspace); @@ -297,7 +297,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); @@ -318,7 +318,7 @@ void distance_impl(raft::resources const& handle, bool is_row_major, DataT) // metric_arg unused { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); raft::linalg::gemm(handle, out, const_cast(x), @@ -377,7 +377,7 @@ void distance_impl(raft::resources const& handle, raft::make_const_mdspan(raft::make_device_vector_view(start, union_size))); } - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); // Calculate Hellinger distance ops::hellinger_distance_op distance_op{}; @@ -436,7 +436,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); @@ -457,7 +457,7 @@ void distance_impl(raft::resources const& handle, bool is_row_major, DataT) // metric_arg unused { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto unaryOp_lambda = [] __device__(DataT input) { auto input_ = raft::to_float(input); @@ -528,7 +528,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); } @@ -609,7 +609,7 @@ void distance_impl(raft::resources const& handle, DataT) // metric_arg unused { bool perform_sqrt = false; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); distance_impl_l2_expanded( perform_sqrt, x, y, out, m, n, k, workspace, worksize, fin_op, stream, is_row_major); } @@ -630,7 +630,7 @@ void distance_impl(raft::resources const& handle, DataT) // metric_arg unused { bool perform_sqrt = true; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); distance_impl_l2_expanded( perform_sqrt, x, y, out, m, n, k, workspace, worksize, fin_op, stream, is_row_major); } @@ -657,7 +657,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( l2_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); @@ -685,7 +685,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( l2_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); @@ -711,7 +711,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); @@ -737,7 +737,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); @@ -763,7 +763,7 @@ void distance_impl(raft::resources const& handle, const OutT* x_norm = nullptr; const OutT* y_norm = nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); pairwise_matrix_dispatch( distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major); diff --git a/cpp/src/distance/detail/kernels/gram_matrix.cu b/cpp/src/distance/detail/kernels/gram_matrix.cu index 6d5b8b1567..1406d6ab9c 100644 --- a/cpp/src/distance/detail/kernels/gram_matrix.cu +++ b/cpp/src/distance/detail/kernels/gram_matrix.cu @@ -366,7 +366,7 @@ void GramMatrixBase::linear(raft::resources const& handle, &beta, out.data_handle(), ld_out, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } else { // #TODO: Use mdspan-based API when stride-capable // https://github.com/rapidsai/raft/issues/875 @@ -384,7 +384,7 @@ void GramMatrixBase::linear(raft::resources const& handle, &beta, out.data_handle(), ld_out, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } } diff --git a/cpp/src/distance/detail/kernels/kernel_matrices.cu b/cpp/src/distance/detail/kernels/kernel_matrices.cu index 9ed25f959c..3de2b713a0 100644 --- a/cpp/src/distance/detail/kernels/kernel_matrices.cu +++ b/cpp/src/distance/detail/kernels/kernel_matrices.cu @@ -182,7 +182,7 @@ void PolynomialKernel::evaluate(raft::resources const& handle, out.extent(0), out.extent(1), is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate kernel matrix using polynomial kernel. @@ -214,7 +214,7 @@ void PolynomialKernel::evaluate(raft::resources const& handle, out.extent(0), out.extent(1), is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate kernel matrix using polynomial kernel. @@ -246,7 +246,7 @@ void PolynomialKernel::evaluate(raft::resources const& handle, out.extent(0), out.extent(1), is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate the Gram matrix using the legacy interface. @@ -333,7 +333,7 @@ void TanhKernel::evaluate(raft::resources const& handle, out.extent(0), out.extent(1), is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate kernel matrix using tanh kernel. @@ -365,7 +365,7 @@ void TanhKernel::evaluate(raft::resources const& handle, out.extent(0), out.extent(1), is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate kernel matrix using tanh kernel. @@ -397,7 +397,7 @@ void TanhKernel::evaluate(raft::resources const& handle, out.extent(0), out.extent(1), is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate the Gram matrix using the legacy interface. @@ -519,7 +519,7 @@ void RBFKernel::evaluate(raft::resources const& handle, math_t* norm_x1, math_t* norm_x2) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); // lazy compute norms if not given rmm::device_uvector tmp_norm_x1(0, stream); rmm::device_uvector tmp_norm_x2(0, stream); @@ -545,7 +545,7 @@ void RBFKernel::evaluate(raft::resources const& handle, norm_x1, norm_x2, is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate kernel matrix using RBF kernel. @@ -569,7 +569,7 @@ void RBFKernel::evaluate(raft::resources const& handle, math_t* norm_x1, math_t* norm_x2) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); // lazy compute norms if not given rmm::device_uvector tmp_norm_x1(0, stream); @@ -596,7 +596,7 @@ void RBFKernel::evaluate(raft::resources const& handle, norm_x1, norm_x2, is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate kernel matrix using RBF kernel. @@ -620,7 +620,7 @@ void RBFKernel::evaluate(raft::resources const& handle, math_t* norm_x1, math_t* norm_x2) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); // lazy compute norms if not given rmm::device_uvector tmp_norm_x1(0, stream); @@ -647,7 +647,7 @@ void RBFKernel::evaluate(raft::resources const& handle, norm_x1, norm_x2, is_row_major, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } /** Evaluate the Gram matrix using the legacy interface. diff --git a/cpp/src/distance/detail/masked_nn.cuh b/cpp/src/distance/detail/masked_nn.cuh index a3a187e8c4..eb9c0ec2c2 100644 --- a/cpp/src/distance/detail/masked_nn.cuh +++ b/cpp/src/distance/detail/masked_nn.cuh @@ -273,7 +273,8 @@ void masked_l2_nn_impl(raft::resources const& handle, dim3 grid(raft::ceildiv(m, P::Nthreads)); dim3 block(P::Nthreads); - initKernel<<>>(out, m, maxVal, redOp); + initKernel + <<>>(out, m, maxVal, redOp); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -293,24 +294,24 @@ void masked_l2_nn_impl(raft::resources const& handle, dim3 block(P::Nthreads); dim3 grid = launchConfigGenerator

(m, n, smemSize, kernel); - kernel<<>>(out, - x, - y, - xn, - yn, - ws_adj64.data(), - group_idxs, - num_groups, - m, - n, - k, - sqrt, - maxVal, - ws_fused_nn.data(), - redOp, - pairRedOp, - core_lambda, - fin_op); + kernel<<>>(out, + x, + y, + xn, + yn, + ws_adj64.data(), + group_idxs, + num_groups, + m, + n, + k, + sqrt, + maxVal, + ws_fused_nn.data(), + redOp, + pairRedOp, + core_lambda, + fin_op); RAFT_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/distance/detail/sparse/bin_distance.cuh b/cpp/src/distance/detail/sparse/bin_distance.cuh index 760bab724a..a9052e7608 100644 --- a/cpp/src/distance/detail/sparse/bin_distance.cuh +++ b/cpp/src/distance/detail/sparse/bin_distance.cuh @@ -127,7 +127,7 @@ class jaccard_expanded_distances_t : public distances_t { config_->a_nrows, search_coo_rows.data(), config_->a_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); compute_bin_distance(out_dists, search_coo_rows.data(), @@ -138,7 +138,7 @@ class jaccard_expanded_distances_t : public distances_t { config_->b_nnz, config_->a_nrows, config_->b_nrows, - raft::resource::get_cuda_stream(config_->handle), + raft::resource::get_cuda_stream(config_->handle).get(), [] __device__ __host__(value_t dot, value_t q_norm, value_t r_norm) { value_t q_r_union = q_norm + r_norm; value_t denom = q_r_union - dot; @@ -186,7 +186,7 @@ class dice_expanded_distances_t : public distances_t { config_->a_nrows, search_coo_rows.data(), config_->a_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); compute_bin_distance(out_dists, search_coo_rows.data(), @@ -197,7 +197,7 @@ class dice_expanded_distances_t : public distances_t { config_->b_nnz, config_->a_nrows, config_->b_nrows, - raft::resource::get_cuda_stream(config_->handle), + raft::resource::get_cuda_stream(config_->handle).get(), [] __device__ __host__(value_t dot, value_t q_norm, value_t r_norm) { value_t q_r_union = q_norm + r_norm; value_t dice = (2 * dot) / q_r_union; diff --git a/cpp/src/distance/detail/sparse/coo_spmv_strategies/base_strategy.cuh b/cpp/src/distance/detail/sparse/coo_spmv_strategies/base_strategy.cuh index f674a1e901..e91aef74f7 100644 --- a/cpp/src/distance/detail/sparse/coo_spmv_strategies/base_strategy.cuh +++ b/cpp/src/distance/detail/sparse/coo_spmv_strategies/base_strategy.cuh @@ -56,25 +56,26 @@ class coo_spmv_strategy { cudaFuncCachePreferShared)); balanced_coo_generalized_spmv_kernel - <<>>(strategy, - a_indptr, - config.a_indices, - config.a_data, - config.a_nnz, - coo_rows_b, - config.b_indices, - config.b_data, - config.a_nrows, - config.b_nrows, - smem_dim, - config.b_nnz, - out_dists, - n_blocks_per_row, - chunk_size, - config.b_ncols, - product_func, - accum_func, - write_func); + <<>>( + strategy, + a_indptr, + config.a_indices, + config.a_data, + config.a_nnz, + coo_rows_b, + config.b_indices, + config.b_data, + config.a_nrows, + config.b_nrows, + smem_dim, + config.b_nnz, + out_dists, + n_blocks_per_row, + chunk_size, + config.b_ncols, + product_func, + accum_func, + write_func); } template - <<>>(strategy, - b_indptr, - config.b_indices, - config.b_data, - config.b_nnz, - coo_rows_a, - config.a_indices, - config.a_data, - config.b_nrows, - config.a_nrows, - smem_dim, - config.a_nnz, - out_dists, - n_blocks_per_row, - chunk_size, - config.a_ncols, - product_func, - accum_func, - write_func); + <<>>( + strategy, + b_indptr, + config.b_indices, + config.b_data, + config.b_nnz, + coo_rows_a, + config.a_indices, + config.a_data, + config.b_nrows, + config.a_nrows, + smem_dim, + config.a_nnz, + out_dists, + n_blocks_per_row, + chunk_size, + config.a_ncols, + product_func, + accum_func, + write_func); } protected: diff --git a/cpp/src/distance/detail/sparse/coo_spmv_strategies/hash_strategy.cuh b/cpp/src/distance/detail/sparse/coo_spmv_strategies/hash_strategy.cuh index 30872a4e46..14c52d089e 100644 --- a/cpp/src/distance/detail/sparse/coo_spmv_strategies/hash_strategy.cuh +++ b/cpp/src/distance/detail/sparse/coo_spmv_strategies/hash_strategy.cuh @@ -101,7 +101,7 @@ class hash_strategy : public coo_spmv_strategy { this->config.a_nrows, mask_indptr, n_rows_divided, - raft::resource::get_cuda_stream(this->config.handle)); + raft::resource::get_cuda_stream(this->config.handle).get()); auto less_rows = std::get<0>(n_rows_divided); if (less_rows > 0) { @@ -127,21 +127,23 @@ class hash_strategy : public coo_spmv_strategy { more_rows + 1, raft::resource::get_cuda_stream(this->config.handle)); rmm::device_uvector chunk_indices( 0, raft::resource::get_cuda_stream(this->config.handle)); - chunked_mask_row_it::init(this->config.a_indptr, - mask_indptr.data() + less_rows, - more_rows, - capacity_threshold * map_size, - n_chunks_per_row, - chunk_indices, - raft::resource::get_cuda_stream(this->config.handle)); - - chunked_mask_row_it more(this->config.a_indptr, - more_rows, - mask_indptr.data() + less_rows, - capacity_threshold * map_size, - n_chunks_per_row.data(), - chunk_indices.data(), - raft::resource::get_cuda_stream(this->config.handle)); + chunked_mask_row_it::init( + this->config.a_indptr, + mask_indptr.data() + less_rows, + more_rows, + capacity_threshold * map_size, + n_chunks_per_row, + chunk_indices, + raft::resource::get_cuda_stream(this->config.handle).get()); + + chunked_mask_row_it more( + this->config.a_indptr, + more_rows, + mask_indptr.data() + less_rows, + capacity_threshold * map_size, + n_chunks_per_row.data(), + chunk_indices.data(), + raft::resource::get_cuda_stream(this->config.handle).get()); auto n_more_blocks = more.total_row_blocks * n_blocks_per_row; this->_dispatch_base(*this, @@ -175,7 +177,7 @@ class hash_strategy : public coo_spmv_strategy { this->config.b_nrows, mask_indptr, n_rows_divided, - raft::resource::get_cuda_stream(this->config.handle)); + raft::resource::get_cuda_stream(this->config.handle).get()); auto less_rows = std::get<0>(n_rows_divided); if (less_rows > 0) { @@ -201,21 +203,23 @@ class hash_strategy : public coo_spmv_strategy { more_rows + 1, raft::resource::get_cuda_stream(this->config.handle)); rmm::device_uvector chunk_indices( 0, raft::resource::get_cuda_stream(this->config.handle)); - chunked_mask_row_it::init(this->config.b_indptr, - mask_indptr.data() + less_rows, - more_rows, - capacity_threshold * map_size, - n_chunks_per_row, - chunk_indices, - raft::resource::get_cuda_stream(this->config.handle)); - - chunked_mask_row_it more(this->config.b_indptr, - more_rows, - mask_indptr.data() + less_rows, - capacity_threshold * map_size, - n_chunks_per_row.data(), - chunk_indices.data(), - raft::resource::get_cuda_stream(this->config.handle)); + chunked_mask_row_it::init( + this->config.b_indptr, + mask_indptr.data() + less_rows, + more_rows, + capacity_threshold * map_size, + n_chunks_per_row, + chunk_indices, + raft::resource::get_cuda_stream(this->config.handle).get()); + + chunked_mask_row_it more( + this->config.b_indptr, + more_rows, + mask_indptr.data() + less_rows, + capacity_threshold * map_size, + n_chunks_per_row.data(), + chunk_indices.data(), + raft::resource::get_cuda_stream(this->config.handle).get()); auto n_more_blocks = more.total_row_blocks * n_blocks_per_row; this->_dispatch_base_rev(*this, diff --git a/cpp/src/distance/detail/sparse/ip_distance.cuh b/cpp/src/distance/detail/sparse/ip_distance.cuh index f3ba399199..b1f183ad09 100644 --- a/cpp/src/distance/detail/sparse/ip_distance.cuh +++ b/cpp/src/distance/detail/sparse/ip_distance.cuh @@ -43,7 +43,7 @@ class ip_distances_t : public distances_t { config_->b_nrows, coo_rows_b.data(), config_->b_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); } /** diff --git a/cpp/src/distance/detail/sparse/l2_distance.cuh b/cpp/src/distance/detail/sparse/l2_distance.cuh index 3b06dcb895..41cc50f43d 100644 --- a/cpp/src/distance/detail/sparse/l2_distance.cuh +++ b/cpp/src/distance/detail/sparse/l2_distance.cuh @@ -237,7 +237,7 @@ class l2_expanded_distances_t : public distances_t { config_->a_nrows, search_coo_rows.data(), config_->a_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); compute_l2(out_dists, search_coo_rows.data(), @@ -248,7 +248,7 @@ class l2_expanded_distances_t : public distances_t { config_->b_nnz, config_->a_nrows, config_->b_nrows, - raft::resource::get_cuda_stream(config_->handle), + raft::resource::get_cuda_stream(config_->handle).get(), [] __device__ __host__(value_t dot, value_t q_norm, value_t r_norm) { return -2 * dot + q_norm + r_norm; }); @@ -312,7 +312,7 @@ class correlation_expanded_distances_t : public distances_t { config_->a_nrows, search_coo_rows.data(), config_->a_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); compute_corr(out_dists, search_coo_rows.data(), @@ -324,7 +324,7 @@ class correlation_expanded_distances_t : public distances_t { config_->a_nrows, config_->b_nrows, config_->b_ncols, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); } ~correlation_expanded_distances_t() = default; @@ -361,7 +361,7 @@ class cosine_expanded_distances_t : public distances_t { config_->a_nrows, search_coo_rows.data(), config_->a_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); compute_l2(out_dists, search_coo_rows.data(), @@ -372,7 +372,7 @@ class cosine_expanded_distances_t : public distances_t { config_->b_nnz, config_->a_nrows, config_->b_nrows, - raft::resource::get_cuda_stream(config_->handle), + raft::resource::get_cuda_stream(config_->handle).get(), [] __device__ __host__(value_t dot, value_t q_norm, value_t r_norm) { value_t norms = raft::sqrt(q_norm) * raft::sqrt(r_norm); // deal with potential for 0 in denominator by forcing 0/1 instead @@ -418,7 +418,7 @@ class hellinger_expanded_distances_t : public distances_t { config_->b_nrows, coo_rows.data(), config_->b_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); balanced_coo_pairwise_generalized_spmv( out_dists, diff --git a/cpp/src/distance/detail/sparse/lp_distance.cuh b/cpp/src/distance/detail/sparse/lp_distance.cuh index 38025329b9..9940e45938 100644 --- a/cpp/src/distance/detail/sparse/lp_distance.cuh +++ b/cpp/src/distance/detail/sparse/lp_distance.cuh @@ -47,7 +47,7 @@ void unexpanded_lp_distances(value_t* out_dists, config_->b_nrows, coo_rows.data(), config_->b_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); balanced_coo_pairwise_generalized_spmv( out_dists, *config_, coo_rows.data(), product_func, accum_func, write_func); @@ -56,7 +56,7 @@ void unexpanded_lp_distances(value_t* out_dists, config_->a_nrows, coo_rows.data(), config_->a_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); balanced_coo_pairwise_generalized_spmv_rev( out_dists, *config_, coo_rows.data(), product_func, accum_func, write_func); @@ -292,7 +292,7 @@ class kl_divergence_unexpanded_distances_t : public distances_t { config_->b_nrows, coo_rows.data(), config_->b_nnz, - raft::resource::get_cuda_stream(config_->handle)); + raft::resource::get_cuda_stream(config_->handle).get()); balanced_coo_pairwise_generalized_spmv( out_dists, diff --git a/cpp/src/distance/detail/sparse/utils.cuh b/cpp/src/distance/detail/sparse/utils.cuh index 50213d0f55..d6656d1964 100644 --- a/cpp/src/distance/detail/sparse/utils.cuh +++ b/cpp/src/distance/detail/sparse/utils.cuh @@ -114,7 +114,7 @@ void faster_dot_on_csr(raft::resources const& handle, dim3 blocks(block_x, block_y, 1); faster_dot_on_csr_kernel - <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); + <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); } else if (dim < 256) { constexpr int tpb = 128; @@ -126,7 +126,7 @@ void faster_dot_on_csr(raft::resources const& handle, dim3 blocks(block_x, block_y, 1); faster_dot_on_csr_kernel - <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); + <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); } else if (dim < 512) { constexpr int tpb = 256; cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -137,7 +137,7 @@ void faster_dot_on_csr(raft::resources const& handle, dim3 blocks(block_x, block_y, 1); faster_dot_on_csr_kernel - <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); + <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); } else { constexpr int tpb = 512; cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -148,7 +148,7 @@ void faster_dot_on_csr(raft::resources const& handle, dim3 blocks(block_x, block_y, 1); faster_dot_on_csr_kernel - <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); + <<>>(dot, indptr, cols, A, B, nnz, n_rows, dim); } RAFT_CUDA_TRY(cudaPeekAtLastError()); diff --git a/cpp/src/distance/distance-inl.cuh b/cpp/src/distance/distance-inl.cuh index 5a6f1f288d..c76722aef3 100644 --- a/cpp/src/distance/distance-inl.cuh +++ b/cpp/src/distance/distance-inl.cuh @@ -238,7 +238,7 @@ void pairwise_distance(raft::resources const& handle, bool isRowMajor = true, DistT metric_arg = 2.0f) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); auto dispatch = [&](auto distance_type) { auto worksize = getWorkspaceSize(x, y, m, n, k); diff --git a/cpp/src/distance/kde.cu b/cpp/src/distance/kde.cu index dc66c808b3..3bc8e12590 100644 --- a/cpp/src/distance/kde.cu +++ b/cpp/src/distance/kde.cu @@ -558,7 +558,7 @@ void kde(raft::resources const& handle, const T* weights_ptr = weights.has_value() ? weights->data_handle() : nullptr; T* output_ptr = output.data_handle(); - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); T log_norm = std::log(sum_weights) + norm_factor(kernel, bandwidth, d); // Cap feature tile to the actual dimension to avoid wasted shared memory diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh index be44f08aa5..015ec407d0 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors_merge.cuh @@ -197,7 +197,7 @@ void merge_subgraphs(raft::resources const& res, { if (num_elems <= 128) { merge_subgraphs_kernel - <<>>( + <<>>( inverted_indices_d, k, num_data_in_cluster, @@ -208,7 +208,7 @@ void merge_subgraphs(raft::resources const& res, select_min); } else if (num_elems <= 512) { merge_subgraphs_kernel - <<>>( + <<>>( inverted_indices_d, k, num_data_in_cluster, @@ -219,7 +219,7 @@ void merge_subgraphs(raft::resources const& res, select_min); } else if (num_elems <= 1024) { merge_subgraphs_kernel - <<>>( + <<>>( inverted_indices_d, k, num_data_in_cluster, @@ -230,7 +230,7 @@ void merge_subgraphs(raft::resources const& res, select_min); } else if (num_elems <= 2048) { merge_subgraphs_kernel - <<>>( + <<>>( inverted_indices_d, k, num_data_in_cluster, diff --git a/cpp/src/neighbors/ball_cover/ball_cover.cuh b/cpp/src/neighbors/ball_cover/ball_cover.cuh index c39756f7d5..d3bcaf7eb2 100644 --- a/cpp/src/neighbors/ball_cover/ball_cover.cuh +++ b/cpp/src/neighbors/ball_cover/ball_cover.cuh @@ -131,7 +131,7 @@ void construct_landmark_1nn(raft::resources const& handle, index.m, index.get_R_indptr().data_handle(), index.n_landmarks + 1, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); // reorder X to allow aligned access raft::matrix::copy_rows( diff --git a/cpp/src/neighbors/ball_cover/registers.cuh b/cpp/src/neighbors/ball_cover/registers.cuh index 4d381d8ca2..528261621c 100644 --- a/cpp/src/neighbors/ball_cover/registers.cuh +++ b/cpp/src/neighbors/ball_cover/registers.cuh @@ -1005,7 +1005,7 @@ void rbc_low_dim_pass_one(raft::resources const& handle, { if (k <= 32) block_rbc_kernel_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1024,7 +1024,7 @@ void rbc_low_dim_pass_one(raft::resources const& handle, else if (k <= 64) block_rbc_kernel_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1042,7 +1042,7 @@ void rbc_low_dim_pass_one(raft::resources const& handle, weight); else if (k <= 128) block_rbc_kernel_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1061,7 +1061,7 @@ void rbc_low_dim_pass_one(raft::resources const& handle, else if (k <= 256) block_rbc_kernel_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1080,7 +1080,7 @@ void rbc_low_dim_pass_one(raft::resources const& handle, else if (k <= 512) block_rbc_kernel_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1099,7 +1099,7 @@ void rbc_low_dim_pass_one(raft::resources const& handle, else if (k <= 1024) block_rbc_kernel_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1142,22 +1142,22 @@ void rbc_low_dim_pass_two(raft::resources const& handle, <<>>(query, - index.n, - R_knn_inds, - R_knn_dists, - index.get_R_radius().data_handle(), - index.get_R().data_handle(), - index.n_landmarks, - bitset_size, - k, - index.metric, - bitset.data(), - weight); + raft::resource::get_cuda_stream(handle).get()>>>(query, + index.n, + R_knn_inds, + R_knn_dists, + index.get_R_radius().data_handle(), + index.get_R().data_handle(), + index.n_landmarks, + bitset_size, + k, + index.metric, + bitset.data(), + weight); if (k <= 32) compute_final_dists_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1174,7 +1174,7 @@ void rbc_low_dim_pass_two(raft::resources const& handle, index.metric); else if (k <= 64) compute_final_dists_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1191,7 +1191,7 @@ void rbc_low_dim_pass_two(raft::resources const& handle, index.metric); else if (k <= 128) compute_final_dists_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1208,7 +1208,7 @@ void rbc_low_dim_pass_two(raft::resources const& handle, index.metric); else if (k <= 256) compute_final_dists_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1225,7 +1225,7 @@ void rbc_low_dim_pass_two(raft::resources const& handle, index.metric); else if (k <= 512) compute_final_dists_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1242,7 +1242,7 @@ void rbc_low_dim_pass_two(raft::resources const& handle, index.metric); else if (k <= 1024) compute_final_dists_registers - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, index.n, @@ -1271,7 +1271,7 @@ void rbc_eps_pass(raft::resources const& handle, value_idx* vd) { block_rbc_kernel_eps_dense - <<>>( + <<>>( index.get_X_reordered().data_handle(), query, n_query_rows, @@ -1296,7 +1296,7 @@ void rbc_eps_pass(raft::resources const& handle, &sum, sizeof(value_idx), cudaMemcpyHostToDevice, - raft::resource::get_cuda_stream(handle))); + raft::resource::get_cuda_stream(handle).get())); } raft::resource::sync_stream(handle); @@ -1325,44 +1325,44 @@ void rbc_eps_pass(raft::resources const& handle, <<(n_query_rows, 2), 64, 0, - raft::resource::get_cuda_stream(handle)>>>(index.get_X_reordered().data_handle(), - query, - n_query_rows, - index.n, - R, - index.m, - eps, - index.n_landmarks, - index.get_R_indptr().data_handle(), - index.get_R_1nn_cols().data_handle(), - index.get_R_1nn_dists().data_handle(), - index.get_R_radius().data_handle(), - dfunc, - vd_ptr, - nullptr, - false, - index.n); + raft::resource::get_cuda_stream(handle).get()>>>(index.get_X_reordered().data_handle(), + query, + n_query_rows, + index.n, + R, + index.m, + eps, + index.n_landmarks, + index.get_R_indptr().data_handle(), + index.get_R_1nn_cols().data_handle(), + index.get_R_1nn_dists().data_handle(), + index.get_R_radius().data_handle(), + dfunc, + vd_ptr, + nullptr, + false, + index.n); } else { block_rbc_kernel_eps_csr_pass <<(n_query_rows, 2), 64, 0, - raft::resource::get_cuda_stream(handle)>>>(index.get_X_reordered().data_handle(), - query, - n_query_rows, - index.n, - R, - index.m, - eps, - index.n_landmarks, - index.get_R_indptr().data_handle(), - index.get_R_1nn_cols().data_handle(), - index.get_R_1nn_dists().data_handle(), - index.get_R_radius().data_handle(), - dfunc, - vd_ptr, - nullptr, - false); + raft::resource::get_cuda_stream(handle).get()>>>(index.get_X_reordered().data_handle(), + query, + n_query_rows, + index.n, + R, + index.m, + eps, + index.n_landmarks, + index.get_R_indptr().data_handle(), + index.get_R_1nn_cols().data_handle(), + index.get_R_1nn_dists().data_handle(), + index.get_R_radius().data_handle(), + dfunc, + vd_ptr, + nullptr, + false); } thrust::exclusive_scan(raft::resource::get_thrust_policy(handle), @@ -1378,7 +1378,29 @@ void rbc_eps_pass(raft::resources const& handle, <<(n_query_rows, 2), 64, 0, - raft::resource::get_cuda_stream(handle)>>>(index.get_X_reordered().data_handle(), + raft::resource::get_cuda_stream(handle).get()>>>(index.get_X_reordered().data_handle(), + query, + n_query_rows, + index.n, + R, + index.m, + eps, + index.n_landmarks, + index.get_R_indptr().data_handle(), + index.get_R_1nn_cols().data_handle(), + index.get_R_1nn_dists().data_handle(), + index.get_R_radius().data_handle(), + dfunc, + adj_ia, + adj_ja, + true, + index.n); + } else { + block_rbc_kernel_eps_csr_pass + <<(n_query_rows, 2), + 64, + 0, + resource::get_cuda_stream(handle).get()>>>(index.get_X_reordered().data_handle(), query, n_query_rows, index.n, @@ -1393,27 +1415,7 @@ void rbc_eps_pass(raft::resources const& handle, dfunc, adj_ia, adj_ja, - true, - index.n); - } else { - block_rbc_kernel_eps_csr_pass - <<(n_query_rows, 2), 64, 0, resource::get_cuda_stream(handle)>>>( - index.get_X_reordered().data_handle(), - query, - n_query_rows, - index.n, - R, - index.m, - eps, - index.n_landmarks, - index.get_R_indptr().data_handle(), - index.get_R_1nn_cols().data_handle(), - index.get_R_1nn_dists().data_handle(), - index.get_R_radius().data_handle(), - dfunc, - adj_ia, - adj_ja, - true); + true); } } } else { @@ -1424,23 +1426,25 @@ void rbc_eps_pass(raft::resources const& handle, raft::resource::get_cuda_stream(handle)); block_rbc_kernel_eps_max_k - <<(n_query_rows, 2), 64, 0, raft::resource::get_cuda_stream(handle)>>>( - index.get_X_reordered().data_handle(), - query, - n_query_rows, - index.n, - R, - index.m, - eps, - index.n_landmarks, - index.get_R_indptr().data_handle(), - index.get_R_1nn_cols().data_handle(), - index.get_R_1nn_dists().data_handle(), - index.get_R_radius().data_handle(), - dfunc, - vd_ptr, - max_k_in, - tmp.data()); + <<(n_query_rows, 2), + 64, + 0, + raft::resource::get_cuda_stream(handle).get()>>>(index.get_X_reordered().data_handle(), + query, + n_query_rows, + index.n, + R, + index.m, + eps, + index.n_landmarks, + index.get_R_indptr().data_handle(), + index.get_R_1nn_cols().data_handle(), + index.get_R_1nn_dists().data_handle(), + index.get_R_radius().data_handle(), + dfunc, + vd_ptr, + max_k_in, + tmp.data()); int64_t actual_max = thrust::reduce(raft::resource::get_thrust_policy(handle), vd_ptr, @@ -1465,7 +1469,7 @@ void rbc_eps_pass(raft::resources const& handle, (value_idx)0); block_rbc_kernel_eps_max_k_copy - <<>>( + <<>>( max_k_in, adj_ia, tmp.data(), adj_ja); // return 'new' max-k @@ -1478,7 +1482,7 @@ void rbc_eps_pass(raft::resources const& handle, adj_ia + n_query_rows, sizeof(value_idx), cudaMemcpyDeviceToDevice, - raft::resource::get_cuda_stream(handle))); + raft::resource::get_cuda_stream(handle).get())); } raft::resource::sync_stream(handle); diff --git a/cpp/src/neighbors/brute_force.cu b/cpp/src/neighbors/brute_force.cu index d499df5796..e4ee656dff 100644 --- a/cpp/src/neighbors/brute_force.cu +++ b/cpp/src/neighbors/brute_force.cu @@ -119,7 +119,7 @@ index::index(raft::resources const& res, dataset_.data_handle(), dataset_view.extent(0), dataset_view.extent(1), - raft::resource::get_cuda_stream(res)); + raft::resource::get_cuda_stream(res).get()); dataset_view_ = raft::make_const_mdspan(dataset_.view()); } @@ -149,7 +149,7 @@ index::index(raft::resources const& res, dataset_.data_handle(), dataset_view.extent(0), dataset_view.extent(1), - raft::resource::get_cuda_stream(res)); + raft::resource::get_cuda_stream(res).get()); dataset_view_ = raft::make_const_mdspan(dataset_.view()); } diff --git a/cpp/src/neighbors/composite/index.cu b/cpp/src/neighbors/composite/index.cu index 4b9150ba01..124c2cd28a 100644 --- a/cpp/src/neighbors/composite/index.cu +++ b/cpp/src/neighbors/composite/index.cu @@ -84,7 +84,7 @@ void composite_index::search( temp_neighbors[i].data(), offset, temp_neighbors[i].size(), - stream); + stream.get()); } raft::copy_matrix( diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh index be2d77aa79..a803ed3ca5 100644 --- a/cpp/src/neighbors/detail/ann_utils.cuh +++ b/cpp/src/neighbors/detail/ann_utils.cuh @@ -224,15 +224,15 @@ HDI constexpr auto mapping::operator()(const float& x) const -> int8_t * @param[in] n_bytes */ template -inline void memzero(T* ptr, IdxT n_elems, rmm::cuda_stream_view stream) +inline void memzero(T* ptr, IdxT n_elems, cuda::stream_ref stream) { switch (check_pointer_residency(ptr)) { case pointer_residency::host_and_device: case pointer_residency::device_only: { - RAFT_CUDA_TRY(cudaMemsetAsync(ptr, 0, n_elems * sizeof(T), stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(ptr, 0, n_elems * sizeof(T), stream.get())); } break; case pointer_residency::host_only: { - stream.synchronize(); + stream.sync(); ::memset(ptr, 0, n_elems * sizeof(T)); } break; default: RAFT_FAIL("memset: unreachable code"); @@ -298,14 +298,14 @@ void block_copy(const IdxT* in_offsets, const T* in_data, T* out_data, IdxT n_mult, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { IdxT in_size; update_host(&in_size, in_offsets + n_blocks, 1, stream); - stream.synchronize(); + stream.sync(); dim3 threads(128, 1, 1); dim3 blocks(raft::ceildiv(in_size * n_mult, threads.x), 1, 1); - block_copy_kernel<<>>( + block_copy_kernel<<>>( in_offsets, out_offsets, n_blocks, in_data, out_data, n_mult); } @@ -325,11 +325,11 @@ void block_copy(const IdxT* in_offsets, * @param stream */ template -void outer_add(const T* a, IdxT len_a, const T* b, IdxT len_b, T* c, rmm::cuda_stream_view stream) +void outer_add(const T* a, IdxT len_a, const T* b, IdxT len_b, T* c, cuda::stream_ref stream) { dim3 threads(128, 1, 1); dim3 blocks(raft::ceildiv(len_a * len_b, threads.x), 1, 1); - outer_add_kernel<<>>(a, len_a, b, len_b, c); + outer_add_kernel<<>>(a, len_a, b, len_b, c); } template @@ -370,25 +370,25 @@ void copy_selected(IdxT n_rows, IdxT ld_src, T* dst, IdxT ld_dst, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { switch (check_pointer_residency(src, dst, row_ids)) { case pointer_residency::host_and_device: case pointer_residency::device_only: { IdxT block_dim = 128; IdxT grid_dim = raft::ceildiv(n_rows * n_cols, block_dim); - copy_selected_kernel - <<>>(n_rows, n_cols, src, row_ids, ld_src, dst, ld_dst); + copy_selected_kernel<<>>( + n_rows, n_cols, src, row_ids, ld_src, dst, ld_dst); } break; case pointer_residency::host_only: { - stream.synchronize(); + stream.sync(); for (IdxT i_dst = 0; i_dst < n_rows; i_dst++) { auto i_src = static_cast(row_ids[i_dst]); for (IdxT j = 0; j < n_cols; j++) { dst[ld_dst * i_dst + j] = mapping{}(src[ld_src * i_src + j]); } } - stream.synchronize(); + stream.sync(); } break; default: RAFT_FAIL("All pointers must reside on the same side, host or device."); } @@ -403,8 +403,7 @@ void copy_selected(IdxT n_rows, * the main stream itself is returned with `false`, and the caller should treat prefetch as a * no-op (no overlap is possible on a single stream). */ -inline auto get_prefetch_stream(raft::resources const& res) - -> std::pair +inline auto get_prefetch_stream(raft::resources const& res) -> std::pair { if (res.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) && raft::resource::get_stream_pool_size(res) >= 1) { @@ -535,7 +534,7 @@ struct batch_load_iterator { } // Stream is shared with the iterator; it must be sync'd before the underlying buffers (or, // in the passthrough case, the source mdspan) can be safely reused. - copy_stream_.synchronize(); + copy_stream_.sync(); } [[nodiscard]] auto row_width() const -> size_type { return row_width_; } @@ -600,7 +599,7 @@ struct batch_load_iterator { batch(raft::resources const& res, MdspanT input_view, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch, bool initialize, @@ -703,7 +702,7 @@ struct batch_load_iterator { prefetch_pos_.reset(); // Ensure prefetch_next_batch()'s queued H2D into this slot (and any prior D2H of the // slot from the previous overwrite) finished before the user kernel reads it. - copy_stream_.synchronize(); + copy_stream_.sync(); } else { // Non-pipelined fast path (prefetch_=false, or prefetch_pos_ didn't match). if (host_writeback_ && dirty_cur_ && pos_.has_value()) { @@ -711,7 +710,7 @@ struct batch_load_iterator { dirty_cur_ = false; } if (initialize_) { queue_h2d(dev_ptr_, row_offset, len); } - copy_stream_.synchronize(); + copy_stream_.sync(); } pos_.emplace(pos); batch_len_ = len; @@ -789,7 +788,7 @@ struct batch_load_iterator { source_ + src_row_offset * row_width_, n_bytes, cudaMemcpyHostToDevice, - copy_stream_)); + copy_stream_.get())); } void queue_d2h(element_type* src, size_type pos) @@ -803,10 +802,10 @@ struct batch_load_iterator { src, n_bytes, cudaMemcpyDeviceToHost, - copy_stream_)); + copy_stream_.get())); } - rmm::cuda_stream_view copy_stream_; + cuda::stream_ref copy_stream_; raft::resources const* res_; MdspanT input_view_; element_type* source_; @@ -860,7 +859,7 @@ struct batch_load_iterator { batch_load_iterator(raft::resources const& res, MdspanT input_view, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch = false, bool initialize = true, @@ -876,7 +875,7 @@ struct batch_load_iterator { batch_load_iterator(raft::resources const& res, MdspanT input_view, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, bool prefetch = false, bool initialize = true, bool host_writeback = false) @@ -1025,7 +1024,7 @@ class batch_load_iterator_dyn { IdxT n_rows, IdxT row_width, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch = false, bool initialize = true, @@ -1050,7 +1049,7 @@ class batch_load_iterator_dyn { IdxT n_rows, IdxT row_width, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, bool prefetch = false, bool initialize = true, bool host_writeback = false) @@ -1156,7 +1155,7 @@ class batch_load_iterator_dyn { IdxT n_rows, IdxT row_width, size_type batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch, bool initialize, @@ -1219,7 +1218,7 @@ auto make_batch_load_iterator(raft::resources const& res, detail::type_identity_t n_rows, detail::type_identity_t row_width, size_t batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, rmm::device_async_resource_ref mr, bool prefetch = false, bool initialize = true, @@ -1244,7 +1243,7 @@ auto make_batch_load_iterator(raft::resources const& res, detail::type_identity_t n_rows, detail::type_identity_t row_width, size_t batch_size, - rmm::cuda_stream_view copy_stream, + cuda::stream_ref copy_stream, bool prefetch = false, bool initialize = true, bool host_writeback = false) -> batch_load_iterator_dyn diff --git a/cpp/src/neighbors/detail/cagra/add_nodes.cuh b/cpp/src/neighbors/detail/cagra/add_nodes.cuh index c19314a835..7f2b372592 100644 --- a/cpp/src/neighbors/detail/cagra/add_nodes.cuh +++ b/cpp/src/neighbors/detail/cagra/add_nodes.cuh @@ -47,7 +47,7 @@ void add_node_core( idx.graph().data_handle(), old_size * degree, 1, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); auto host_num_incoming_edges = raft::make_host_vector(new_size); raft::copy(handle, host_num_incoming_edges.view(), dev_num_incoming_edges.view()); diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 44770fda8f..021a96d457 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -195,7 +195,7 @@ void ace_get_partition_labels( size_t n_partitions = partition_histogram.extent(0); RAFT_EXPECTS(labels_dim == 2, "Labels must have 2 columns"); RAFT_EXPECTS(partition_histogram.extent(1) == 2, "Partition histogram must have 2 columns"); - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); // Sampling vectors from dataset. Uses float conversion on host instead of // raft::matrix::sample_rows to minimize GPU memory usage. diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index aa00f73cfe..a935df9b6f 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -158,7 +158,7 @@ cuvs::neighbors::cagra::index merge_rebuild( return index; }; - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); if (bitset_filtered) { auto staging = raft::make_device_mdarray( @@ -421,7 +421,7 @@ auto merge_fastener(raft::resources const& handle, 0, static_cast(stride - preflight.dim) * sizeof(T), static_cast(preflight.rows), - raft::resource::get_cuda_stream(handle))); + raft::resource::get_cuda_stream(handle).get())); } copy_input_datasets( handle, indices, preflight.offsets, preflight.dim, stride, destination); diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index 44e215189f..897991d2d5 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -630,7 +630,7 @@ void assign_bucket(raft::resources const& res, // Gather the batch's tile rows and leader vectors int point_blocks = strided_grid_size(static_cast(batch_size * point_elements)); - manyway_gather_tile_points_kernel<<>>( + manyway_gather_tile_points_kernel<<>>( dataset.data_handle(), dim, row_stride, @@ -642,7 +642,7 @@ void assign_bucket(raft::resources const& res, RAFT_CUDA_TRY(cudaGetLastError()); int leader_blocks = strided_grid_size(static_cast(batch_size * leader_elements)); - manyway_gather_tile_leaders_kernel<<>>( + manyway_gather_tile_leaders_kernel<<>>( dataset.data_handle(), dim, row_stride, @@ -954,7 +954,7 @@ auto build_leaf_neighbors(raft::resources const& res, size_t batch_size = std::min(batch_capacity, leaves.starts_host.size() - leaf_offset); int gather_blocks = strided_grid_size(static_cast(batch_size * vector_elements_per_leaf)); - manyway_gather_leaf_vectors_kernel<<>>( + manyway_gather_leaf_vectors_kernel<<>>( dataset.data_handle(), input_dimension, row_stride, @@ -1102,7 +1102,7 @@ auto build(raft::resources const& res, split_context context(res, rows, dim); int norm_blocks = static_cast((rows + ROW_WARPS_PER_BLOCK - 1) / ROW_WARPS_PER_BLOCK); - manyway_l2_norms_kernel<<>>( + manyway_l2_norms_kernel<<>>( dataset.data_handle(), rows, dim, dataset.extent(1), context.norms.data_handle()); RAFT_CUDA_TRY(cudaGetLastError()); diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold_shared.cu b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold_shared.cu index e79999b209..5386d1e274 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold_shared.cu +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold_shared.cu @@ -338,7 +338,8 @@ void launch_initialize_root_memberships(raft::resources const& res, initialize_root_memberships_kernel<<>>(memberships, rows); + raft::resource::get_cuda_stream(res).get()>>>(memberships, + rows); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -352,7 +353,7 @@ void launch_carry_completed_parents(raft::resources const& res, carry_completed_parents_kernel<<>>( + raft::resource::get_cuda_stream(res).get()>>>( input, spans, output, output_keys); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -371,7 +372,7 @@ void launch_materialize_tile_distances(raft::resources const& res, materialize_tile_distances_kernel<<>>( + raft::resource::get_cuda_stream(res).get()>>>( dots, batch_size, tile_rows, padded_leaders, norms, leader_ids, input_memberships, tiles); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -391,15 +392,15 @@ void launch_emit_tile_assignments(raft::resources const& res, emit_tile_assignments_kernel<<>>(selected_leaders, - batch_size, - tile_rows, - fanout, - occurrence_stride, - input_memberships, - tiles, - output_keys, - output_memberships); + raft::resource::get_cuda_stream(res).get()>>>(selected_leaders, + batch_size, + tile_rows, + fanout, + occurrence_stride, + input_memberships, + tiles, + output_keys, + output_memberships); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -461,7 +462,7 @@ void launch_initialize_self_scaffold(raft::resources const& res, initialize_self_scaffold_kernel<<>>( + raft::resource::get_cuda_stream(res).get()>>>( graph, rows, graph_degree, scaffold_offset, scaffold_degree); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -484,19 +485,19 @@ void launch_leaf_gram_knn(raft::resources const& res, leaf_gram_knn_kernel<<(leaf_count), leaf_size, 0, - raft::resource::get_cuda_stream(res)>>>(gram, - memberships, - origins, - leaf_starts, - leaf_counts, - leaf_strides, - leaf_offset, - leaf_count, - leaf_size, - leaf_degree, - graph_degree, - scaffold_offset, - graph); + raft::resource::get_cuda_stream(res).get()>>>(gram, + memberships, + origins, + leaf_starts, + leaf_counts, + leaf_strides, + leaf_offset, + leaf_count, + leaf_size, + leaf_degree, + graph_degree, + scaffold_offset, + graph); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -504,7 +505,10 @@ void launch_initialize_origins( raft::resources const& res, uint32_t* origins, int64_t start, int64_t rows, uint32_t origin) { auto blocks = static_cast(raft::div_rounding_up_safe(rows, THREADS_PER_BLOCK)); - initialize_origins_kernel<<>>( + initialize_origins_kernel<<>>( origins, start, rows, origin); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -523,7 +527,7 @@ void launch_copy_partition_graph(raft::resources const& res, copy_partition_graph_kernel<<>>( + raft::resource::get_cuda_stream(res).get()>>>( source, source_rows, source_degree, destination, destination_degree, base_degree, offset); RAFT_CUDA_TRY(cudaGetLastError()); } @@ -540,7 +544,7 @@ void launch_deduplicate_graph_prefix(raft::resources const& res, deduplicate_graph_prefix_kernel<<>>( + raft::resource::get_cuda_stream(res).get()>>>( input, rows, input_degree, output, output_degree); RAFT_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 2d7683816c..40645ffb09 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -434,7 +434,7 @@ void search_multi_partition( auto plan_desc = dataset_descriptor_init_with_cache( res, params, indices[0]->dataset(), metric, dataset_norms_ptr0); - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); // Cap the per-launch query count. num_queries maps to grid.y in the multi-partition kernels, // which is bounded by maxGridSize[1]; chunking also bounds the intermediate workspaces, which diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index a99ec64bc0..024eaf82ea 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -211,9 +211,8 @@ struct dataset_descriptor_host { // Codebook type is determined by DataT for VPQ (always half for now) struct state { - using ready_t = std::tuple; - using init_f = - std::tuple, size_t>; + using ready_t = std::tuple; + using init_f = std::tuple, size_t>; std::mutex mutex; std::atomic ready; // Not sure if std::holds_alternative is thread-safe @@ -230,31 +229,33 @@ struct dataset_descriptor_host { { if (std::holds_alternative(value)) { auto& [ptr, stream] = std::get(value); - RAFT_CUDA_TRY_NO_THROW(cudaFreeAsync(ptr, stream)); + RAFT_CUDA_TRY_NO_THROW(cudaFreeAsync(ptr, stream.get())); } RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(ready_event)); } - void eval(rmm::cuda_stream_view stream) + void eval(cuda::stream_ref stream) { std::lock_guard lock(mutex); if (std::holds_alternative(value)) { auto& [fun, size] = std::get(value); dev_descriptor_t* ptr = nullptr; - RAFT_CUDA_TRY(cudaMallocAsync(&ptr, size, stream)); + RAFT_CUDA_TRY(cudaMallocAsync(&ptr, size, stream.get())); fun(ptr, stream); - RAFT_CUDA_TRY(cudaEventRecord(ready_event, stream)); + RAFT_CUDA_TRY(cudaEventRecord(ready_event, stream.get())); value = std::make_tuple(ptr, stream); ready.store(true, std::memory_order_release); } } - auto get(rmm::cuda_stream_view stream) -> dev_descriptor_t* + auto get(cuda::stream_ref stream) -> dev_descriptor_t* { if (!ready.load(std::memory_order_acquire)) { eval(stream); } // value is immutable at this point. auto& [ptr, ready_stream] = std::get(value); - if (ready_stream != stream) { RAFT_CUDA_TRY(cudaStreamWaitEvent(stream, ready_event, 0)); } + if (ready_stream != stream.get()) { + RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); + } return ptr; } }; @@ -286,12 +287,12 @@ struct dataset_descriptor_host { /** * Return the device pointer, possibly evaluating it in the given thread. */ - [[nodiscard]] auto dev_ptr(rmm::cuda_stream_view stream) const -> const dev_descriptor_t* + [[nodiscard]] auto dev_ptr(cuda::stream_ref stream) const -> const dev_descriptor_t* { return value_->get(stream); } - [[nodiscard]] auto dev_ptr(rmm::cuda_stream_view stream) -> dev_descriptor_t* + [[nodiscard]] auto dev_ptr(cuda::stream_ref stream) -> dev_descriptor_t* { return value_->get(stream); } diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh b/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh index cde42f849b..b0dd95f9ea 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh +++ b/cpp/src/neighbors/detail/cagra/compute_distance_standard-impl.cuh @@ -131,23 +131,23 @@ standard_descriptor_spec* dev_ptr, - rmm::cuda_stream_view stream) { - standard_dataset_descriptor_init_kernel - <<<1, 1, 0, stream>>>(dev_ptr, ptr, size, dim, ld, dataset_norms); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - }, - Metric, - DatasetBlockDim, - false, // is_vpq - 0, // pq_bits - 0}; // pq_len + return host_type{ + desc_type{ptr, size, dim, ld, dataset_norms}, + [=](dataset_descriptor_base_t* dev_ptr, cuda::stream_ref stream) { + standard_dataset_descriptor_init_kernel + <<<1, 1, 0, stream.get()>>>(dev_ptr, ptr, size, dim, ld, dataset_norms); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + }, + Metric, + DatasetBlockDim, + false, // is_vpq + 0, // pq_bits + 0}; // pq_len } } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh b/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh index d286e0b07e..e42fa6e2b7 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh +++ b/cpp/src/neighbors/detail/cagra/compute_distance_vpq-impl.cuh @@ -221,8 +221,7 @@ vpq_descriptor_spec* dev_ptr, - rmm::cuda_stream_view stream) { + [=](dataset_descriptor_base_t* dev_ptr, cuda::stream_ref stream) { vpq_dataset_descriptor_init_kernel<<<1, 1, 0, stream>>>(dev_ptr, - encoded_dataset_ptr, - encoded_dataset_dim, - vq_code_book_ptr, - pq_code_book_ptr, - size, - dim); + SmemDType><<<1, 1, 0, stream.get()>>>(dev_ptr, + encoded_dataset_ptr, + encoded_dataset_dim, + vq_code_book_ptr, + pq_code_book_ptr, + size, + dim); RAFT_CUDA_TRY(cudaPeekAtLastError()); }, Metric, diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 5c5c2d5920..6138412414 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -769,16 +769,18 @@ void merge_graph_gpu( auto mst_graph_num_edges_view = (*d_mst_graph_num_edges).view(); auto output_view = (*d_output_graph).view(); kern_merge_graph - <<>>( - output_view, - d_rev_graph, - d_rev_graph_count, - mst_graph_view, - mst_graph_num_edges_view, - batch_size, - i_batch, - guarantee_connectivity, - d_check_num_protected_edges.data_handle()); + <<>>(output_view, + d_rev_graph, + d_rev_graph_count, + mst_graph_view, + mst_graph_num_edges_view, + batch_size, + i_batch, + guarantee_connectivity, + d_check_num_protected_edges.data_handle()); d_output_graph.prefetch_next_batch(); d_mst_graph.prefetch_next_batch(); @@ -828,7 +830,7 @@ void make_reverse_graph_gpu( dim3 threads(256, 1, 1); dim3 blocks(1024, 1, 1); for (uint64_t k = 0; k < output_graph_degree; k++) { - kern_make_rev_graph_k<<>>( + kern_make_rev_graph_k<<>>( output_graph, d_rev_graph, d_rev_graph_count, k); } } else { @@ -843,7 +845,7 @@ void make_reverse_graph_gpu( dim3 threads(256, 1, 1); dim3 blocks(1024, 1, 1); - kern_make_rev_graph_k<<>>( + kern_make_rev_graph_k<<>>( d_dest_nodes.view(), d_rev_graph, d_rev_graph_count, 0); raft::resource::sync_stream(res); RAFT_LOG_DEBUG("# Making reverse graph on GPUs: %lu / %u \r", k, output_graph_degree); @@ -1175,7 +1177,7 @@ void mst_optimization( 1 * sizeof(IdxT), // width graph_size, cudaMemcpyDeviceToHost, - raft::resource::get_cuda_stream(res))); + raft::resource::get_cuda_stream(res).get())); raft::resource::sync_stream(res); // FIXME: use submdspan and raft::copy once supported @@ -1199,7 +1201,7 @@ void mst_optimization( constexpr uint64_t n_threads = 256; const dim3 threads(n_threads, 1, 1); const dim3 blocks(raft::ceildiv(graph_size, n_threads), 1, 1); - kern_mst_opt_update_graph<<>>( + kern_mst_opt_update_graph<<>>( d_mst_graph_ptr, d_candidate_edges_ptr, d_outgoing_num_edges_ptr, @@ -1245,7 +1247,7 @@ void mst_optimization( constexpr uint64_t n_threads = 256; const dim3 threads(n_threads, 1, 1); const dim3 blocks((graph_size + n_threads - 1) / n_threads, 1, 1); - kern_mst_opt_labeling<<>>( + kern_mst_opt_labeling<<>>( d_label_ptr, d_mst_graph_ptr, graph_size, mst_graph_degree, d_stats_ptr); raft::copy(res, @@ -1279,7 +1281,7 @@ void mst_optimization( constexpr uint64_t n_threads = 256; const dim3 threads(n_threads, 1, 1); const dim3 blocks(raft::ceildiv(graph_size, n_threads), 1, 1); - kern_mst_opt_cluster_size<<>>( + kern_mst_opt_cluster_size<<>>( d_cluster_size_ptr, d_label_ptr, graph_size, d_stats_ptr); raft::copy(res, @@ -1320,7 +1322,10 @@ void mst_optimization( constexpr uint64_t n_threads = 256; const dim3 threads(n_threads, 1, 1); const dim3 blocks((graph_size + n_threads - 1) / n_threads, 1, 1); - kern_mst_opt_postprocessing<<>>( + kern_mst_opt_postprocessing<<>>( d_outgoing_num_edges_ptr, d_incoming_num_edges_ptr, d_outgoing_max_edges_ptr, @@ -1497,13 +1502,15 @@ void prune_graph_gpu( for (uint32_t i_batch = 0; i_batch < num_batch; i_batch++) { auto output_view = (*d_output_graph).view(); kern_fused_prune - <<>>( - input_view, - output_view, - batch_size, - i_batch, - d_invalid_neighbor_list.data_handle(), - dev_stats.data_handle()); + <<>>(input_view, + output_view, + batch_size, + i_batch, + d_invalid_neighbor_list.data_handle(), + dev_stats.data_handle()); d_output_graph.prefetch_next_batch(); ++d_output_graph; diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index 3ed3d9c4b9..c088b8c477 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -151,7 +151,7 @@ void launch_sort_knn_graph_impl(raft::resources const& res, constexpr uint32_t block_size = 256; auto const warps = block_size / raft::WarpSize; auto const blocks = (dataset_size + warps - 1) / warps; - kernel<<>>( + kernel<<>>( dataset, dataset_dim, knn_graph, dataset_size, graph_degree, metric); RAFT_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta.cuh index 20913dcf39..8a0ccbf69f 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta.cuh @@ -219,7 +219,7 @@ struct search SAMPLE_FILTER_T sample_filter) { auto source_indices_ptr = source_indices.has_value() ? source_indices->data_handle() : nullptr; - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); select_and_run(dataset_desc, graph, source_indices_ptr, @@ -307,7 +307,7 @@ struct search DISTANCE_T* intermediate_distances_ptr, SampleFilterT_ sample_filter) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); // Scale the cross-CTA traversed hashmap to (num_queries * num_partitions) rows. const size_t traversed_hash_size = hashmap::get_size(hash_bitlen); diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh index 4edfb264fd..e299e1451f 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh @@ -415,15 +415,20 @@ struct search workspace, sort, hints, - stream); + stream.get()); } if (ldIK > numElements) { if (input_keys_storage.size() != sizeBatch * numElements) { input_keys_storage.resize(sizeBatch * numElements, stream); } - batched_memcpy( - input_keys_storage.data(), numElements, inputKeys, ldIK, numElements, sizeBatch, stream); + batched_memcpy(input_keys_storage.data(), + numElements, + inputKeys, + ldIK, + numElements, + sizeBatch, + stream.get()); inputKeys = input_keys_storage.data(); } @@ -432,8 +437,13 @@ struct search input_values_storage.resize(sizeBatch * numElements, stream); } - batched_memcpy( - input_values_storage.data(), numElements, inputVals, ldIV, numElements, sizeBatch, stream); + batched_memcpy(input_values_storage.data(), + numElements, + inputVals, + ldIV, + numElements, + sizeBatch, + stream.get()); inputVals = input_values_storage.data(); } @@ -457,11 +467,13 @@ struct search sort); if (ldOK > topK) { - batched_memcpy(outputKeys, ldOK, output_keys_storage.data(), topK, topK, sizeBatch, stream); + batched_memcpy( + outputKeys, ldOK, output_keys_storage.data(), topK, topK, sizeBatch, stream.get()); } if (ldOV > topK) { - batched_memcpy(outputVals, ldOV, output_values_storage.data(), topK, topK, sizeBatch, stream); + batched_memcpy( + outputVals, ldOV, output_values_storage.data(), topK, topK, sizeBatch, stream.get()); } } @@ -479,7 +491,7 @@ struct search SAMPLE_FILTER_T sample_filter) { // Init hashmap - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); const uint32_t hash_size = hashmap::get_size(hash_bitlen); set_value_batch( hashmap.data(), hash_size, utils::get_max_value(), hash_size, num_queries, stream); diff --git a/cpp/src/neighbors/detail/cagra/search_plan.cuh b/cpp/src/neighbors/detail/cagra/search_plan.cuh index 8a9d79e177..766ca4ae17 100644 --- a/cpp/src/neighbors/detail/cagra/search_plan.cuh +++ b/cpp/src/neighbors/detail/cagra/search_plan.cuh @@ -38,8 +38,8 @@ namespace cuvs::neighbors::cagra::detail { template struct lightweight_uvector { private: - using raft_res_type = const raft::resources*; - using rmm_res_type = std::tuple; + using raft_res_type = const raft::resources*; + using rmm_res_type = std::tuple; static constexpr size_t kAlign = 256; std::variant res_; @@ -68,21 +68,21 @@ struct lightweight_uvector { } auto copy_size = std::min(size_, new_size); if (copy_size > 0) { - cudaMemcpyAsync(new_ptr, ptr_, copy_size * sizeof(T), cudaMemcpyDefault, s); + cudaMemcpyAsync(new_ptr, ptr_, copy_size * sizeof(T), cudaMemcpyDefault, s.get()); } if (size_ > 0) { r.deallocate(s, ptr_, size_ * sizeof(T), kAlign); } ptr_ = new_ptr; size_ = new_size; } - void resize(size_t new_size, rmm::cuda_stream_view stream) + void resize(size_t new_size, cuda::stream_ref stream) { if (new_size == size_) { return; } if (std::holds_alternative(res_)) { auto& h = std::get(res_); res_ = rmm_res_type{raft::resource::get_workspace_resource_ref(*h), stream}; } else { - std::get(std::get(res_)) = stream; + std::get(std::get(res_)) = stream; } resize(new_size); } diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index 62fe431964..8fd50a1549 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -240,7 +240,7 @@ struct search uint32_t topk, SampleFilterT sample_filter) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); // Allocate global hashmap when small-hash is disabled via the workspace pool // (no cudaMallocAsync/cudaFreeAsync after pool warmup). @@ -289,7 +289,7 @@ struct search uint32_t topk, SAMPLE_FILTER_T sample_filter) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); constexpr uintptr_t kOutputIndexTag = raft::Pow2::Log2; const auto result_indices_uintptr = reinterpret_cast(result_indices_ptr); static_assert(kOutputIndexTag <= 3, "OutputIndexT can't be more than 8 bytes"); diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index e66e60d941..f3d23e4460 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -297,7 +297,7 @@ auto deserialize_device_dense(raft::resources const& res, std::istream& is) 0, (metadata.stride - metadata.dim) * sizeof(DataT), metadata.n_rows, - stream)); + stream.get())); work_queued = true; } if (metadata.elements > 0) { diff --git a/cpp/src/neighbors/detail/dynamic_batching.cuh b/cpp/src/neighbors/detail/dynamic_batching.cuh index 9839a7fe19..1239192e2b 100644 --- a/cpp/src/neighbors/detail/dynamic_batching.cuh +++ b/cpp/src/neighbors/detail/dynamic_batching.cuh @@ -914,7 +914,7 @@ class batch_runner { kernel_progress_counters_.data_handle(), 0, sizeof(*kernel_progress_counters_.data_handle()) * kernel_progress_counters_.size(), - raft::resource::get_cuda_stream(res_))); + raft::resource::get_cuda_stream(res_).get())); // Make sure to initialize the atomic values in the batch_state structs. for (uint32_t i = 0; i < n_queues_; i++) { auto seq_id = batch_queue_.push(); @@ -1026,7 +1026,7 @@ class batch_runner { rem_time_us_ref.store(static_cast(params.dispatch_timeout_ms * 1000), cuda::std::memory_order_relaxed); // run the gather kernel before submitting the data to reduce the latency - gather_inputs<<>>( + gather_inputs<<>>( slice_3d(batch_id, queries_), request_ptrs, &rem_time_us_ref, @@ -1075,13 +1075,13 @@ class batch_runner { // next_batch_token); auto bs = dim3(128, 8, 1); scatter_outputs - <<<1, bs, 0, stream>>>(request_ptrs, - batch_neighbors, - batch_distances, - kernel_progress_counters_.data_handle() + batch_id, - &next_token_ref, - batch_queue::make_seq_batch_id(next_seq_id, batch_id)); - RAFT_CUDA_TRY(cudaEventRecord(completion_events_[batch_id].value(), stream)); + <<<1, bs, 0, stream.get()>>>(request_ptrs, + batch_neighbors, + batch_distances, + kernel_progress_counters_.data_handle() + batch_id, + &next_token_ref, + batch_queue::make_seq_batch_id(next_seq_id, batch_id)); + RAFT_CUDA_TRY(cudaEventRecord(completion_events_[batch_id].value(), stream.get())); dispatch_sequence_id_ref.store(seq_id.value, cuda::std::memory_order_release); dispatch_sequence_id_ref.notify_all(); @@ -1094,7 +1094,7 @@ class batch_runner { dispatched_id_observed = dispatch_sequence_id_ref.load(cuda::std::memory_order_acquire); } // Now we can safely record the event - RAFT_CUDA_TRY(cudaStreamWaitEvent(stream, completion_events_[batch_id].value())); + RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), completion_events_[batch_id].value())); } n_queries -= queries_committed; diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index bac127294e..d39d525f80 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -1097,7 +1097,6 @@ from_cagra(raft::resources const& res, std::optional> dataset) { common::nvtx::range fun_scope("hnsw::from_cagra"); - auto stream = raft::resource::get_cuda_stream(res); auto num_threads = params.num_threads == 0 ? cuvs::core::omp::get_max_threads() : params.num_threads; @@ -1327,7 +1326,7 @@ from_cagra(raft::resources const& res, degree * sizeof(uint32_t), n_rows, cudaMemcpyDefault, - raft::resource::get_cuda_stream(res))); + raft::resource::get_cuda_stream(res).get())); #pragma omp parallel for num_threads(num_threads) for (int64_t i = 0; i < n_rows; i++) { appr_algo->setListCount(appr_algo->get_linklist0(i), degree); diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index 76ed71fe57..e54b2aa89c 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -426,7 +426,7 @@ void brute_force_knn_impl( rmm::device_uvector search_row_major(0, userStream); if (!rowMajorQuery) { search_row_major.resize(n * D, userStream); - raft::linalg::transpose(handle, search, search_row_major.data(), n, D, userStream); + raft::linalg::transpose(handle, search, search_row_major.data(), n, D, userStream.get()); search = search_row_major.data(); } @@ -465,7 +465,7 @@ void brute_force_knn_impl( k, rowMajorIndex, rowMajorQuery, - stream, + stream.get(), metric, input_norms ? (*input_norms)[i] : nullptr, search_norms); @@ -490,7 +490,7 @@ void brute_force_knn_impl( "Haversine distance requires 2 dimensions " "(latitude / longitude)."); - haversine_knn(out_i_ptr, out_d_ptr, input[i], search_items, sizes[i], n, k, stream); + haversine_knn(out_i_ptr, out_d_ptr, input[i], search_items, sizes[i], n, k, stream.get()); break; default: // Create a new handle with the current stream from the stream pool @@ -501,7 +501,7 @@ void brute_force_knn_impl( if (!rowMajorIndex) { index = index_row_major.data() + total_rows_processed * D; total_rows_processed += sizes[i]; - raft::linalg::transpose(handle, input[i], index, sizes[i], D, stream); + raft::linalg::transpose(handle, input[i], index, sizes[i], D, stream.get()); } tiled_brute_force_knn(stream_pool_handle, @@ -872,7 +872,7 @@ void brute_force_search_filtered( compressed_csr_view.get_n_rows(), rows.data(), compressed_csr_view.get_nnz(), - stream); + stream.get()); cuvs::neighbors::detail::epilogue_on_csr( res, csr.get_elements().data(), diff --git a/cpp/src/neighbors/detail/knn_utils.cuh b/cpp/src/neighbors/detail/knn_utils.cuh index 72df390f8e..d6c61066b6 100644 --- a/cpp/src/neighbors/detail/knn_utils.cuh +++ b/cpp/src/neighbors/detail/knn_utils.cuh @@ -47,7 +47,7 @@ void epilogue_on_csr(raft::resources const& handle, int blocks = raft::ceildiv((size_t)nnz, tpb); if (metric == cuvs::distance::DistanceType::L2Expanded) { - epilogue_on_csr_kernel<<>>( + epilogue_on_csr_kernel<<>>( compressed_C, rows, cols, @@ -62,7 +62,7 @@ void epilogue_on_csr(raft::resources const& handle, } }); } else if (metric == cuvs::distance::DistanceType::L2SqrtExpanded) { - epilogue_on_csr_kernel<<>>( + epilogue_on_csr_kernel<<>>( compressed_C, rows, cols, @@ -77,7 +77,7 @@ void epilogue_on_csr(raft::resources const& handle, } }); } else if (metric == cuvs::distance::DistanceType::CosineExpanded) { - epilogue_on_csr_kernel<<>>( + epilogue_on_csr_kernel<<>>( compressed_C, rows, cols, diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index 432cde7ffc..ec92e1a36f 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1513,7 +1513,7 @@ void GNND::build(Data_t* data, "Data type needs to be int8 or uint8 for NN Descent to run with BitwiseHamming distance."); } - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); nrow_ = nrow; graph_.nrow = nrow; graph_.bloom_filter.set_nrow(nrow); diff --git a/cpp/src/neighbors/detail/sparse_knn.cuh b/cpp/src/neighbors/detail/sparse_knn.cuh index 2f20c30b97..fe43c9fc40 100644 --- a/cpp/src/neighbors/detail/sparse_knn.cuh +++ b/cpp/src/neighbors/detail/sparse_knn.cuh @@ -174,7 +174,7 @@ class sparse_knn_t { raft::resource::get_cuda_stream(handle)); value_idx n_query_batch_nnz = query_batcher.get_batch_csr_indptr_nnz( - query_batch_indptr.data(), raft::resource::get_cuda_stream(handle)); + query_batch_indptr.data(), raft::resource::get_cuda_stream(handle).get()); rmm::device_uvector query_batch_indices(n_query_batch_nnz, raft::resource::get_cuda_stream(handle)); @@ -183,7 +183,7 @@ class sparse_knn_t { query_batcher.get_batch_csr_indices_data(query_batch_indices.data(), query_batch_data.data(), - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); // A 3-partition temporary merge space to scale the batching. 2 parts for subsequent // batches and 1 space for the results of the merge, which get copied back to the top @@ -216,13 +216,14 @@ class sparse_knn_t { rmm::device_uvector idx_batch_data(0, raft::resource::get_cuda_stream(handle)); value_idx idx_batch_nnz = idx_batcher.get_batch_csr_indptr_nnz( - idx_batch_indptr.data(), raft::resource::get_cuda_stream(handle)); + idx_batch_indptr.data(), raft::resource::get_cuda_stream(handle).get()); idx_batch_indices.resize(idx_batch_nnz, raft::resource::get_cuda_stream(handle)); idx_batch_data.resize(idx_batch_nnz, raft::resource::get_cuda_stream(handle)); - idx_batcher.get_batch_csr_indices_data( - idx_batch_indices.data(), idx_batch_data.data(), raft::resource::get_cuda_stream(handle)); + idx_batcher.get_batch_csr_indices_data(idx_batch_indices.data(), + idx_batch_data.data(), + raft::resource::get_cuda_stream(handle).get()); /** * Compute distances @@ -253,8 +254,10 @@ class sparse_knn_t { // populate batch indices array value_idx batch_rows = query_batcher.batch_rows(), batch_cols = idx_batcher.batch_rows(); - iota_fill( - batch_indices.data(), batch_rows, batch_cols, raft::resource::get_cuda_stream(handle)); + iota_fill(batch_indices.data(), + batch_rows, + batch_cols, + raft::resource::get_cuda_stream(handle).get()); /** * Perform k-selection on batch & merge with other k-selections diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index fc262b3311..67d027b94d 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -159,12 +159,13 @@ void batched_insert_vamana( raft::make_extents(max_batchsize, visited_size)); // Assign memory to query_list structures and initialize - init_query_candidate_list<<<256, blockD, 0, stream>>>(query_list, - visited_ids.data_handle(), - visited_dists.data_handle(), - (int)max_batchsize, - visited_size, - 1); + init_query_candidate_list + <<<256, blockD, 0, stream.get()>>>(query_list, + visited_ids.data_handle(), + visited_dists.data_handle(), + (int)max_batchsize, + visited_size, + 1); auto topk_pq_mem = raft::make_device_mdarray>(res, raft::resource::get_large_workspace_resource_ref(res), @@ -241,20 +242,20 @@ void batched_insert_vamana( res, raft::make_device_vector_view(query_ids.data_handle(), int64_t(step_size)), raft::make_host_vector_view(insert_order.data() + start, int64_t(step_size))); - set_query_ids<<>>( + set_query_ids<<>>( query_list_ptr.data_handle(), query_ids.data_handle(), step_size); // Call greedy search to get candidates for every vector being inserted GreedySearchKernel - <<>>(d_graph.view(), - dataset, - query_list_ptr.data_handle(), - step_size, - *medoid_id, - visited_size, - metric, - queue_size, - topk_pq_mem.data_handle()); + <<>>(d_graph.view(), + dataset, + query_list_ptr.data_handle(), + step_size, + *medoid_id, + visited_size, + metric, + queue_size, + topk_pq_mem.data_handle()); RAFT_CUDA_TRY(cudaPeekAtLastError()); #if KERNEL_TIMING @@ -266,7 +267,7 @@ void batched_insert_vamana( #endif // Segmented sort on query list - SortPairsKernel<<>>( + SortPairsKernel<<>>( query_list_ptr.data_handle(), step_size, visited_size); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -280,18 +281,18 @@ void batched_insert_vamana( // Run on candidates of vectors being inserted RobustPruneKernel - <<>>(d_graph.view(), - dataset, - query_list_ptr.data_handle(), - step_size, - visited_size, - metric, - alpha, - s_coords_mem.data_handle()); + <<>>(d_graph.view(), + dataset, + query_list_ptr.data_handle(), + step_size, + visited_size, + metric, + alpha, + s_coords_mem.data_handle()); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Segmented sort on query list - SortPairsKernel<<>>( + SortPairsKernel<<>>( query_list_ptr.data_handle(), step_size, degree); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -304,7 +305,7 @@ void batched_insert_vamana( #endif // Write results from first prune to graph edge list - write_graph_edges_kernel<<>>( + write_graph_edges_kernel<<>>( d_graph.view(), query_list_ptr.data_handle(), degree, step_size); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -320,7 +321,8 @@ void batched_insert_vamana( // auto d_total_edges = raft::make_device_mdarray( // res, raft::resource::get_workspace_resource_ref(res), raft::make_extents(1)); rmm::device_scalar d_total_edges(stream); - prefix_sums_sizes<<<1, 1, 0, stream>>>(query_list, step_size, d_total_edges.data()); + prefix_sums_sizes + <<<1, 1, 0, stream.get()>>>(query_list, step_size, d_total_edges.data()); RAFT_CUDA_TRY(cudaPeekAtLastError()); int total_edges = d_total_edges.value(stream); @@ -343,11 +345,11 @@ void batched_insert_vamana( // Create reverse edge list create_reverse_edge_list - <<>>(query_list_ptr.data_handle(), - step_size, - degree, - edge_src.data_handle(), - edge_dist_pair.data_handle()); + <<>>(query_list_ptr.data_handle(), + step_size, + degree, + edge_src.data_handle(), + edge_dist_pair.data_handle()); RAFT_CUDA_TRY(cudaPeekAtLastError()); { @@ -361,7 +363,7 @@ void batched_insert_vamana( edge_src.data_handle(), total_edges, CmpDist(), - stream); + stream.get()); RAFT_LOG_DEBUG("Temp storage needed for sorting dist (bytes): %lu", temp_storage_bytes); @@ -377,7 +379,7 @@ void batched_insert_vamana( edge_src.data_handle(), total_edges, CmpDist(), - stream); + stream.get()); } /* @@ -400,7 +402,7 @@ void batched_insert_vamana( edge_src.data_handle(), total_edges, CmpEdge(), - stream); + stream.get()); RAFT_LOG_DEBUG("Temp storage needed for sorting (bytes): %lu", temp_storage_bytes); @@ -416,11 +418,11 @@ void batched_insert_vamana( edge_src.data_handle(), total_edges, CmpEdge(), - stream); + stream.get()); // Get number of unique node destinations IdxT unique_dests = - cuvs::sparse::neighbors::get_n_components(edge_dest.data_handle(), total_edges, stream); + cuvs::sparse::neighbors::get_n_components(edge_dest.data_handle(), total_edges, stream.get()); // Find which node IDs have reverse edges and their indices in the reverse edge list thrust::device_vector edge_dest_vec(edge_dest.data_handle(), @@ -466,52 +468,53 @@ void batched_insert_vamana( QueryCandidates* reverse_list = static_cast*>(reverse_list_ptr.data_handle()); - init_query_candidate_list<<<256, blockD, 0, stream>>>(reverse_list, - rev_ids.data_handle(), - rev_dists.data_handle(), - (int)reverse_batch, - visited_size); + init_query_candidate_list + <<<256, blockD, 0, stream.get()>>>(reverse_list, + rev_ids.data_handle(), + rev_dists.data_handle(), + (int)reverse_batch, + visited_size); // May need more blocks for reverse list num_blocks = min(maxBlocks, reverse_batch); // Populate reverse list ids and candidate lists from edge_src and edge_dest populate_reverse_list_struct - <<>>(reverse_list, - edge_src.data_handle(), - edge_dest.data_handle(), - unique_indices.data_handle(), - unique_dests, - total_edges, - dataset.extent(0), - rev_start, - reverse_batch); + <<>>(reverse_list, + edge_src.data_handle(), + edge_dest.data_handle(), + unique_indices.data_handle(), + unique_dests, + total_edges, + dataset.extent(0), + rev_start, + reverse_batch); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Recompute distances (avoided keeping it during sorting) recompute_reverse_dists - <<>>(reverse_list, dataset, reverse_batch, metric); + <<>>(reverse_list, dataset, reverse_batch, metric); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Call 2nd RobustPrune on reverse query_list - RobustPruneKernel - <<>>(d_graph.view(), - raft::make_const_mdspan(dataset), - reverse_list_ptr.data_handle(), - reverse_batch, - visited_size, - metric, - alpha, - s_coords_mem.data_handle()); + RobustPruneKernel<<>>( + d_graph.view(), + raft::make_const_mdspan(dataset), + reverse_list_ptr.data_handle(), + reverse_batch, + visited_size, + metric, + alpha, + s_coords_mem.data_handle()); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Segmented sort on reverse_list - SortPairsKernel<<>>( + SortPairsKernel<<>>( reverse_list_ptr.data_handle(), reverse_batch, degree); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Write new edge lists to graph - write_graph_edges_kernel<<>>( + write_graph_edges_kernel<<>>( d_graph.view(), reverse_list_ptr.data_handle(), degree, reverse_batch); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/src/neighbors/detail/vpq_dataset.cuh b/cpp/src/neighbors/detail/vpq_dataset.cuh index 1ccf59dd81..4d27c5424a 100644 --- a/cpp/src/neighbors/detail/vpq_dataset.cuh +++ b/cpp/src/neighbors/detail/vpq_dataset.cuh @@ -488,14 +488,15 @@ void process_and_fill_codes( } dim3 blocks( raft::div_rounding_up_safe(dataset_view.extent(0), kBlockSize / threads_per_vec), 1, 1); - kernel<<>>(codes_view, - dataset_view, - pq_centers, - vq_centers, - raft::make_const_mdspan(labels_view), - rows_in_shared_memory, - pq_bits, - inline_vq_labels); + kernel<<>>( + codes_view, + dataset_view, + pq_centers, + vq_centers, + raft::make_const_mdspan(labels_view), + rows_in_shared_memory, + pq_bits, + inline_vq_labels); RAFT_CUDA_TRY(cudaPeekAtLastError()); }; auto batch_labels = raft::make_device_vector(res, 0); @@ -891,13 +892,14 @@ void process_and_fill_codes_subspaces( } dim3 blocks( raft::div_rounding_up_safe(dataset_view.extent(0), kBlockSize / threads_per_vec), 1, 1); - kernel<<>>(codes_view, - dataset_view, - pq_centers, - vq_centers, - raft::make_const_mdspan(labels_view), - pq_bits, - shared_memory_size > 0); + kernel<<>>( + codes_view, + dataset_view, + pq_centers, + vq_centers, + raft::make_const_mdspan(labels_view), + pq_bits, + shared_memory_size > 0); RAFT_CUDA_TRY(cudaPeekAtLastError()); }; if (!need_batching && !need_copy_to_device) { diff --git a/cpp/src/neighbors/epsilon_neighborhood.cu b/cpp/src/neighbors/epsilon_neighborhood.cu index d05dc21d20..af90eb92c9 100644 --- a/cpp/src/neighbors/epsilon_neighborhood.cu +++ b/cpp/src/neighbors/epsilon_neighborhood.cu @@ -46,7 +46,7 @@ void compute(raft::resources const& handle, y.extent(0), x.extent(1), eps, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); } // Explicit template instantiations diff --git a/cpp/src/neighbors/ivf_common.cu b/cpp/src/neighbors/ivf_common.cu index b87a14f7c3..6e6e0d4177 100644 --- a/cpp/src/neighbors/ivf_common.cu +++ b/cpp/src/neighbors/ivf_common.cu @@ -53,7 +53,7 @@ void calc_chunk_indices::configured::operator()(const uint32_t* cluster_sizes, const uint32_t* clusters_to_probe, uint32_t* chunk_indices, uint32_t* n_samples, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { void* kernel = nullptr; switch (block_dim.x) { @@ -70,7 +70,7 @@ void calc_chunk_indices::configured::operator()(const uint32_t* cluster_sizes, void* args[] = // NOLINT {&n_probes, &cluster_sizes, &clusters_to_probe, &chunk_indices, &n_samples}; - RAFT_CUDA_TRY(cudaLaunchKernel(kernel, grid_dim, block_dim, args, 0, stream)); + RAFT_CUDA_TRY(cudaLaunchKernel(kernel, grid_dim, block_dim, args, 0, stream.get())); } // Helper function to sort cluster sizes using CUB, extracted from template to avoid @@ -78,17 +78,23 @@ void calc_chunk_indices::configured::operator()(const uint32_t* cluster_sizes, void sort_cluster_sizes_descending(uint32_t* input, uint32_t* output, uint32_t n_lists, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref tmp_res) { int begin_bit = 0; int end_bit = sizeof(uint32_t) * 8; size_t cub_workspace_size = 0; cub::DeviceRadixSort::SortKeysDescending( - nullptr, cub_workspace_size, input, output, n_lists, begin_bit, end_bit, stream); + nullptr, cub_workspace_size, input, output, n_lists, begin_bit, end_bit, stream.get()); rmm::device_buffer cub_workspace(cub_workspace_size, stream, tmp_res); - cub::DeviceRadixSort::SortKeysDescending( - cub_workspace.data(), cub_workspace_size, input, output, n_lists, begin_bit, end_bit, stream); + cub::DeviceRadixSort::SortKeysDescending(cub_workspace.data(), + cub_workspace_size, + input, + output, + n_lists, + begin_bit, + end_bit, + stream.get()); } } // namespace cuvs::neighbors::ivf::detail diff --git a/cpp/src/neighbors/ivf_common.cuh b/cpp/src/neighbors/ivf_common.cuh index e466a13fd7..b3cf54f79a 100644 --- a/cpp/src/neighbors/ivf_common.cuh +++ b/cpp/src/neighbors/ivf_common.cuh @@ -19,7 +19,7 @@ namespace cuvs::neighbors::ivf::detail { void sort_cluster_sizes_descending(uint32_t* input, uint32_t* output, uint32_t n_lists, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, rmm::device_async_resource_ref tmp_res); /** @@ -57,7 +57,7 @@ struct calc_chunk_indices { const uint32_t* clusters_to_probe, uint32_t* chunk_indices, uint32_t* n_samples, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); }; static inline auto configure(uint32_t n_probes, uint32_t n_queries) -> configured @@ -153,19 +153,19 @@ void postprocess_neighbors(IdxT* neighbors_out, // [n_queries, to uint32_t n_queries, uint32_t n_probes, uint32_t topk, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { constexpr int kPNThreads = 256; const int pn_blocks = raft::div_rounding_up_unsafe(n_queries * topk, kPNThreads); postprocess_neighbors_kernel - <<>>(neighbors_out, - neighbors_in, - db_indices, - clusters_to_probe, - chunk_indices, - n_queries, - n_probes, - topk); + <<>>(neighbors_out, + neighbors_in, + db_indices, + clusters_to_probe, + chunk_indices, + n_queries, + n_probes, + topk); } /** diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh index 1566056eb6..f5639f0b35 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_build.cuh @@ -261,7 +261,7 @@ void extend(raft::resources const& handle, new_labels.data_handle(), n_rows, 1, - stream); + stream.get()); raft::linalg::add(handle, raft::make_device_vector_view(list_sizes_ptr, n_lists), raft::make_device_vector_view( @@ -315,16 +315,16 @@ void extend(raft::resources const& handle, const dim3 block_dim(256); const dim3 grid_dim(raft::ceildiv(batch.size(), block_dim.x)); build_index_kernel - <<>>(new_labels.data_handle() + batch.offset(), - batch_data_view.data_handle(), - idx_batch->data(), - index->data_ptrs().data_handle(), - index->inds_ptrs().data_handle(), - list_sizes_ptr, - batch.size(), - dim, - index->veclen(), - batch.offset()); + <<>>(new_labels.data_handle() + batch.offset(), + batch_data_view.data_handle(), + idx_batch->data(), + index->data_ptrs().data_handle(), + index->inds_ptrs().data_handle(), + list_sizes_ptr, + batch.size(), + dim, + index->veclen(), + batch.offset()); vec_batches.prefetch_next_batch(); // User needs to make sure kernel finishes its work before we overwrite batch in the next // iteration if different streams are used for kernel and copy. @@ -501,15 +501,15 @@ inline void fill_refinement_index(raft::resources const& handle, const dim3 block_dim(256); const dim3 grid_dim(raft::ceildiv(n_queries * n_candidates, block_dim.x)); build_index_kernel - <<>>(new_labels.data(), - dataset, - candidate_idx, - refinement_index->data_ptrs().data_handle(), - refinement_index->inds_ptrs().data_handle(), - list_sizes_ptr, - n_queries * n_candidates, - refinement_index->dim(), - refinement_index->veclen()); + <<>>(new_labels.data(), + dataset, + candidate_idx, + refinement_index->data_ptrs().data_handle(), + refinement_index->inds_ptrs().data_handle(), + list_sizes_ptr, + n_queries * n_candidates, + refinement_index->dim(), + refinement_index->veclen()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_helpers.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_helpers.cuh index d123835c2a..97e6f127de 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_helpers.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_helpers.cuh @@ -99,7 +99,7 @@ void pack_list_data( dim3 blocks(raft::div_rounding_up_safe(n_rows, kBlockSize), 1, 1); dim3 threads(kBlockSize, 1, 1); auto stream = raft::resource::get_cuda_stream(res); - pack_interleaved_list_kernel<<>>( + pack_interleaved_list_kernel<<>>( codes.data_handle(), list_data.data_handle(), n_rows, dim, veclen, offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -120,7 +120,7 @@ void unpack_list_data( dim3 blocks(raft::div_rounding_up_safe(n_rows, kBlockSize), 1, 1); dim3 threads(kBlockSize, 1, 1); auto stream = raft::resource::get_cuda_stream(res); - unpack_interleaved_list_kernel<<>>( + unpack_interleaved_list_kernel<<>>( list_data.data_handle(), codes.data_handle(), n_rows, dim, veclen, offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh index 052b7bfe9a..f5f61bc786 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_explicit_inst.cuh @@ -36,7 +36,7 @@ uint32_t* neighbors, \ float* distances, \ uint32_t& grid_dim_x, \ - rmm::cuda_stream_view stream, \ + cuda::stream_ref stream, \ const std::optional& metric_udf); #define COMMA , diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh index 3a782822b4..e5d8da76d6 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_ext.cuh @@ -35,7 +35,7 @@ void ivfflat_interleaved_scan(const index& index, uint32_t* neighbors, float* distances, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, const std::optional& metric_udf) RAFT_EXPLICIT; #define CUVS_INST_IVF_FLAT_INTERLEAVED_SCAN(T, IdxT, SampleFilterT) \ @@ -58,7 +58,7 @@ void ivfflat_interleaved_scan(const index& index, uint32_t* neighbors, \ float* distances, \ uint32_t& grid_dim_x, \ - rmm::cuda_stream_view stream, \ + cuda::stream_ref stream, \ const std::optional& metric_udf); CUVS_INST_IVF_FLAT_INTERLEAVED_SCAN(float, int64_t, cuvs::neighbors::filtering::none_sample_filter); diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh index f6ff6fde64..c9bfc93e7f 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan_jit.cuh @@ -151,7 +151,7 @@ void launch_kernel(const index& index, uint32_t* neighbors, float* distances, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, const std::optional& metric_udf) { RAFT_EXPECTS(Veclen == index.veclen(), @@ -218,7 +218,7 @@ void launch_kernel(const index& index, n_probes, smem_size); kernel_launcher->dispatch>( - stream, + stream.get(), grid_dim, block_dim, smem_size, @@ -435,7 +435,7 @@ void ivfflat_interleaved_scan(const index& index, uint32_t* neighbors, float* distances, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, const std::optional& metric_udf) { const uint32_t n_probes_clamped = std::min(n_probes, index.n_lists()); diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh index 960d48c818..a80fef017d 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh @@ -159,7 +159,7 @@ void search_impl(raft::resources const& handle, &beta, distance_buffer_dev.data(), index.n_lists(), - stream); + stream.get()); if (effective_metric == cuvs::distance::DistanceType::CosineExpanded) { auto n_lists = index.n_lists(); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in index 05892b9333..320f6bd98e 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_compute_similarity_run_inst.cu.in @@ -17,7 +17,7 @@ namespace cuvs::neighbors::ivf_pq::detail { template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( cuvs::neighbors::ivf_pq::detail::selected s, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, uint32_t dim, uint32_t n_probes, uint32_t pq_dim, diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu index 5fac491532..d215c69202 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_contiguous_list_data.cu @@ -16,7 +16,7 @@ void unpack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { unpack_contiguous_list_data_impl( codes, list_data, n_rows, pq_dim, offset_or_indices, pq_bits, stream); @@ -31,7 +31,7 @@ void pack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { pack_contiguous_list_data_impl( list_data, codes, n_rows, pq_dim, offset_or_indices, pq_bits, stream); diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu index 34a7244a67..c439298f04 100644 --- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu +++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_list_data.cu @@ -13,7 +13,7 @@ void unpack_list_data(raft::device_matrix_view list_data, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { unpack_list_data_impl(codes, list_data, offset_or_indices, pq_bits, stream); }; @@ -24,7 +24,7 @@ void pack_list_data(raft::device_mdspan codes, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { pack_list_data_impl(list_data, codes, offset_or_indices, pq_bits, stream); }; diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index e8f3dae911..187b382b9f 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -101,7 +101,7 @@ void select_residuals(raft::resources const& handle, // need to know it, any strictly positive number would work. thrust::transform_iterator, const T*, thrust::use_default, float> mapping_itr(dataset, utils::mapping{}); - raft::matrix::gather(mapping_itr, (IdxT)dim, n_rows, row_ids, n_rows, tmp.data(), stream); + raft::matrix::gather(mapping_itr, (IdxT)dim, n_rows, row_ids, n_rows, tmp.data(), stream.get()); raft::matrix::linewise_op( handle, @@ -126,7 +126,7 @@ void select_residuals(raft::resources const& handle, &beta, residuals, rot_dim, - stream); + stream.get()); } /** @@ -202,7 +202,7 @@ void flat_compute_residuals( &beta, residuals, rot_dim, - stream); + stream.get()); } template @@ -235,7 +235,7 @@ auto calculate_offsets_and_indices(IdxT n_rows, const uint32_t* cluster_sizes, IdxT* cluster_offsets, IdxT* data_indices, - rmm::cuda_stream_view stream) -> uint32_t + cuda::stream_ref stream) -> uint32_t { auto exec_policy = rmm::exec_policy(stream); // Calculate the offsets @@ -246,7 +246,7 @@ auto calculate_offsets_and_indices(IdxT n_rows, raft::update_host(&cumsum, cluster_offsets + n_lists, 1, stream); uint32_t max_cluster_size = *thrust::max_element(exec_policy, cluster_sizes, cluster_sizes + n_lists); - stream.synchronize(); + stream.sync(); RAFT_EXPECTS(cumsum == n_rows, "cluster sizes do not add up."); RAFT_LOG_DEBUG("Max cluster size %d", max_cluster_size); rmm::device_uvector data_offsets_buf(n_lists, stream); @@ -255,7 +255,7 @@ auto calculate_offsets_and_indices(IdxT n_rows, constexpr uint32_t n_threads = 128; // NOLINT const IdxT n_blocks = raft::div_rounding_up_unsafe(n_rows, n_threads); fill_indices_kernel - <<>>(n_rows, data_indices, data_offsets, labels); + <<>>(n_rows, data_indices, data_offsets, labels); return max_cluster_size; } @@ -379,7 +379,7 @@ void train_per_subset(raft::resources const& handle, &beta, sub_trainset.data(), impl->pq_len(), - stream); + stream.get()); // train PQ codebook for this subspace auto sub_trainset_view = raft::make_device_matrix_view( @@ -441,7 +441,7 @@ void train_per_cluster(raft::resources const& handle, labels, n_rows, 1, - stream); + stream.get()); auto cluster_offsets = offsets_buf.data(); auto indices = indices_buf.data(); @@ -630,13 +630,14 @@ void reconstruct_list_data(raft::resources const& res, default: RAFT_FAIL("Invalid pq_bits (%u), the value must be within [4, 8]", pq_bits); } }(index.pq_bits()); - kernel<<>>(tmp.view(), - typed_list->data.view(), - index.pq_centers(), - index.centers_rot(), - index.codebook_kind(), - label, - offset_or_indices); + kernel<<>>( + tmp.view(), + typed_list->data.view(), + index.pq_centers(), + index.centers_rot(), + index.codebook_kind(), + label, + offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); float* out_float_ptr = nullptr; @@ -666,7 +667,7 @@ void reconstruct_list_data(raft::resources const& res, &beta, out_float_ptr, index.dim(), - raft::resource::get_cuda_stream(res)); + raft::resource::get_cuda_stream(res).get()); // Transform the data to the original type, if necessary if constexpr (!std::is_same_v) { raft::linalg::map( @@ -757,7 +758,7 @@ void encode_list_data(raft::resources const& res, default: RAFT_FAIL("Invalid pq_bits (%u), the value must be within [4, 8]", pq_bits); } }(index->pq_bits()); - kernel<<>>( + kernel<<>>( index->lists()[label]->data_ptr(), new_vectors_residual.view(), index->pq_centers(), @@ -777,7 +778,7 @@ void encode_list_data(raft::resources const& res, } }(index->pq_bits()); auto typed_list = std::static_pointer_cast>(index->lists()[label]); - kernel<<>>( + kernel<<>>( typed_list->data.view(), new_vectors_residual.view(), index->pq_centers(), @@ -1162,7 +1163,7 @@ void extend(raft::resources const& handle, new_data_labels.data(), n_rows, 1, - stream); + stream.get()); raft::linalg::add( handle, raft::make_device_vector_view(list_sizes, n_clusters), diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu index 2deab7a073..1d702b78ae 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu @@ -248,7 +248,7 @@ void make_rotation_matrix(raft::resources const& handle, rmm::device_uvector buf(inplace ? 0 : n * n, stream); float* mat = inplace ? rotation_matrix : buf.data(); raft::random::normal(handle, rng, mat, n * n, 0.0f, 1.0f); - raft::linalg::detail::qrGetQ_inplace(handle, mat, n, n, stream); + raft::linalg::detail::qrGetQ_inplace(handle, mat, n, n, stream.get()); if (!inplace) { raft::copy_matrix(rotation_matrix, n_cols, mat, n, n_cols, n_rows, stream); } } else { uint32_t stride = n + 1; @@ -375,7 +375,7 @@ void rotate_padded_centers( &beta, rotated_centers.data_handle(), rot_dim, // ldc (leading dim of output) - stream); + stream.get()); } void resize_list(raft::resources const& res, diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh index c55e3cdbc0..34e9b760eb 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity.cuh @@ -14,7 +14,7 @@ #include #include // cuvs::neighbors::ivf_pq::codebook_gen #include // RAFT_WEAK_FUNCTION -#include // rmm::cuda_stream_view +#include // cuda::stream_ref #include // __half @@ -28,7 +28,7 @@ auto RAFT_WEAK_FUNCTION is_local_topk_feasible(uint32_t k, uint32_t n_probes, ui template void compute_similarity_run(selected s, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, uint32_t dim, uint32_t n_probes, uint32_t pq_dim, @@ -138,7 +138,7 @@ auto compute_similarity_select(const cudaDeviceProp& dev_props, \ extern template void cuvs::neighbors::ivf_pq::detail::compute_similarity_run( \ cuvs::neighbors::ivf_pq::detail::selected s, \ - rmm::cuda_stream_view stream, \ + cuda::stream_ref stream, \ uint32_t dim, \ uint32_t n_probes, \ uint32_t pq_dim, \ diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh index 8aa8ce3c30..8550e080d2 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh @@ -21,7 +21,7 @@ #include // raft::Pow2 #include // raft::TxN_t -#include // rmm::cuda_stream_view +#include // cuda::stream_ref namespace cuvs::neighbors::ivf_pq::detail { @@ -265,7 +265,7 @@ struct occupancy_t { template void compute_similarity_run(selected s, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, uint32_t dim, uint32_t n_probes, uint32_t pq_dim, @@ -290,7 +290,7 @@ void compute_similarity_run(selected s, OutT* _out_scores, uint32_t* _out_indices) { - s.launcher->template dispatch>(stream, + s.launcher->template dispatch>(stream.get(), s.grid_dim, s.block_dim, s.smem_size, diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh index 87c4c4ea79..507528e905 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data.cuh @@ -20,7 +20,7 @@ void unpack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); template void unpack_contiguous_list_data(raft::resources const& res, @@ -53,7 +53,7 @@ void pack_contiguous_list_data( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); template void pack_contiguous_list_data(raft::resources const& res, diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh index bc477004b0..593797f1c2 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_contiguous_list_data_impl.cuh @@ -72,7 +72,7 @@ inline void unpack_contiguous_list_data_impl( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { if (n_rows == 0) { return; } @@ -89,7 +89,7 @@ inline void unpack_contiguous_list_data_impl( default: RAFT_FAIL("Invalid pq_bits (%u), the value must be within [4, 8]", pq_bits); } }(); - kernel<<>>(codes, list_data, n_rows, pq_dim, offset_or_indices); + kernel<<>>(codes, list_data, n_rows, pq_dim, offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -154,7 +154,7 @@ inline void pack_contiguous_list_data_impl( uint32_t pq_dim, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { if (n_rows == 0) { return; } @@ -171,7 +171,7 @@ inline void pack_contiguous_list_data_impl( default: RAFT_FAIL("Invalid pq_bits (%u), the value must be within [4, 8]", pq_bits); } }(); - kernel<<>>(list_data, codes, n_rows, pq_dim, offset_or_indices); + kernel<<>>(list_data, codes, n_rows, pq_dim, offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } // namespace cuvs::neighbors::ivf_pq::detail diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp index ac6bbec662..774a2a4ef5 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data.hpp @@ -18,7 +18,7 @@ void unpack_list_data(raft::device_matrix_view list_data, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); void pack_list_data(raft::device_mdspan::list_extents, @@ -26,7 +26,7 @@ void pack_list_data(raft::device_mdspan codes, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); /** Unpack the list data; see the public interface for the api and usage. */ template diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh index e03fb79398..2f6d0504b3 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_list_data_impl.cuh @@ -68,7 +68,7 @@ inline void unpack_list_data_impl( raft::row_major> list_data, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { auto n_rows = codes.extent(0); if (n_rows == 0) { return; } @@ -86,7 +86,7 @@ inline void unpack_list_data_impl( default: RAFT_FAIL("Invalid pq_bits (%u), the value must be within [4, 8]", pq_bits); } }(); - kernel<<>>(codes, list_data, offset_or_indices); + kernel<<>>(codes, list_data, offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -143,7 +143,7 @@ inline void pack_list_data_impl( raft::device_matrix_view codes, std::variant offset_or_indices, uint32_t pq_bits, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { auto n_rows = codes.extent(0); if (n_rows == 0) { return; } @@ -161,7 +161,7 @@ inline void pack_list_data_impl( default: RAFT_FAIL("Invalid pq_bits (%u), the value must be within [4, 8]", pq_bits); } }(); - kernel<<>>(list_data, codes, offset_or_indices); + kernel<<>>(list_data, codes, offset_or_indices); RAFT_CUDA_TRY(cudaPeekAtLastError()); } }; // namespace cuvs::neighbors::ivf_pq::detail diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_process_and_fill_codes_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_process_and_fill_codes_impl.cuh index 01f37dcd1a..1e6e3e4bc9 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_process_and_fill_codes_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_process_and_fill_codes_impl.cuh @@ -95,16 +95,17 @@ void launch_process_and_fill_codes_kernel( } }(index.pq_bits()); - kernel<<>>(new_vectors_residual, - src_offset_or_indices, - new_labels, - index.list_sizes(), - index.inds_ptrs(), - index.data_ptrs(), - index.pq_centers(), - index.codebook_kind(), - index.codes_layout(), - bytes_per_vector); + kernel<<>>( + new_vectors_residual, + src_offset_or_indices, + new_labels, + index.list_sizes(), + index.inds_ptrs(), + index.data_ptrs(), + index.pq_centers(), + index.codebook_kind(), + index.codes_layout(), + bytes_per_vector); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 1dd3d7663d..6518284614 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -154,7 +154,7 @@ void select_clusters(raft::resources const& handle, &beta, qc_distances.data(), n_lists, - stream); + stream.get()); // Select neighbor clusters for each query. rmm::device_uvector cluster_dists(size_t(n_queries) * size_t(n_probes), stream, mr); @@ -240,7 +240,7 @@ void select_clusters(raft::resources const& handle, &beta, qc_distances.data(), n_lists, - stream); + stream.get()); // Select neighbor clusters for each query. rmm::device_uvector cluster_dists(size_t(n_queries) * size_t(n_probes), stream, mr); @@ -325,7 +325,7 @@ void select_clusters(raft::resources const& handle, &beta, qc_distances.data(), n_lists, - stream); + stream.get()); // Select neighbor clusters for each query. rmm::device_uvector cluster_dists(size_t(n_queries) * size_t(n_probes), stream, mr); @@ -508,7 +508,7 @@ void ivfpq_search_worker(raft::resources const& handle, n_queries_probes, begin_bit, end_bit, - stream); + stream.get()); rmm::device_buffer cub_workspace(cub_workspace_size, stream, mr); cub::DeviceRadixSort::SortPairs(cub_workspace.data(), cub_workspace_size, @@ -519,7 +519,7 @@ void ivfpq_search_worker(raft::resources const& handle, n_queries_probes, begin_bit, end_bit, - stream); + stream.get()); } // select and run the main search kernel @@ -1014,7 +1014,7 @@ inline void search(raft::resources const& handle, &beta, rot_queries.data(), index.rot_dim(), - stream); + stream.get()); }, gemm_queries); if (index.metric() == distance::DistanceType::CosineExpanded) { diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh index e0032a4ac5..c01019ace8 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_transform.cuh @@ -90,11 +90,12 @@ void transform_batch(raft::resources const& res, } }(index.pq_bits()); - kernel<<>>(dataset_residuals.view(), - output_labels, - output_dataset, - index.pq_centers(), - index.codebook_kind()); + kernel<<>>( + dataset_residuals.view(), + output_labels, + output_dataset, + index.pq_centers(), + index.codebook_kind()); } template diff --git a/cpp/src/neighbors/ivf_rabitq.cu b/cpp/src/neighbors/ivf_rabitq.cu index d5572d4039..e5ba0102f1 100644 --- a/cpp/src/neighbors/ivf_rabitq.cu +++ b/cpp/src/neighbors/ivf_rabitq.cu @@ -192,7 +192,7 @@ void search(raft::resources const& handle, raft::device_matrix_view neighbors, raft::device_matrix_view distances) { - auto stream = raft::resource::get_cuda_stream(handle).value(); + auto stream = raft::resource::get_cuda_stream(handle).get(); size_t NQ = queries.extent(0); size_t k = neighbors.extent(1); diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh index d2bbbcf345..27690d64af 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/initializer_gpu.cuh @@ -81,7 +81,7 @@ class InitializerGPU { size_t D; // Dimension size_t K; // Num of Centroids raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ }; diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index 8a6e23223d..03d8cc50ea 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -493,7 +493,7 @@ void IVFGPU::construct_on_gpu(const float* device_data, lower_level, upper_level, num_vectors, - stream_)); + stream_.get())); { rmm::device_buffer d_temp_storage(temp_storage_bytes, stream_); @@ -505,7 +505,7 @@ void IVFGPU::construct_on_gpu(const float* device_data, lower_level, upper_level, num_vectors, - stream_)); + stream_.get())); } // ------------------------- @@ -514,8 +514,12 @@ void IVFGPU::construct_on_gpu(const float* device_data, rmm::device_uvector d_offsets(num_centroids + 1, stream_); temp_storage_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( - nullptr, temp_storage_bytes, d_histogram.data(), d_offsets.data(), num_centroids, stream_)); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + temp_storage_bytes, + d_histogram.data(), + d_offsets.data(), + num_centroids, + stream_.get())); { rmm::device_buffer d_temp_storage(temp_storage_bytes, stream_); @@ -524,7 +528,7 @@ void IVFGPU::construct_on_gpu(const float* device_data, d_histogram.data(), d_offsets.data(), num_centroids, - stream_)); + stream_.get())); } // Set the last offset element @@ -536,7 +540,7 @@ void IVFGPU::construct_on_gpu(const float* device_data, GPUClusterMeta* d_cluster_meta_temp = cluster_meta_.data_handle(); num_blocks = (num_centroids + block_size - 1) / block_size; - build_cluster_meta_kernel<<>>( + build_cluster_meta_kernel<<>>( d_cluster_meta_temp, d_histogram.data(), d_offsets.data(), num_centroids); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -550,7 +554,7 @@ void IVFGPU::construct_on_gpu(const float* device_data, raft::copy(d_atomic_counters.data(), d_offsets.data(), num_centroids, stream_); num_blocks = (num_vectors + block_size - 1) / block_size; - scatter_pids_kernel<<>>( + scatter_pids_kernel<<>>( d_flat_pids, device_cluster_ids, d_offsets.data(), d_atomic_counters.data(), num_vectors); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -645,7 +649,7 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, lower_level, upper_level, num_vectors, - stream_)); + stream_.get())); auto d_temp_storage_vec = raft::make_device_vector(handle_, temp_storage_bytes); d_temp_storage = d_temp_storage_vec.data_handle(); @@ -658,7 +662,7 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, lower_level, upper_level, num_vectors, - stream_)); + stream_.get())); // ------------------------- // 4. Compute prefix sum (offsets) on GPU using CUB @@ -672,7 +676,7 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, d_histogram.data_handle(), d_offsets.data_handle(), num_centroids, - stream_)); + stream_.get())); d_temp_storage_vec = raft::make_device_vector(handle_, temp_storage_bytes); d_temp_storage = d_temp_storage_vec.data_handle(); @@ -682,7 +686,7 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, d_histogram.data_handle(), d_offsets.data_handle(), num_centroids, - stream_)); + stream_.get())); // Set the last offset element raft::copy(d_offsets.data_handle() + num_centroids, &num_vectors, 1, stream_); @@ -694,7 +698,7 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, int block_size = 256; int num_blocks = (num_centroids + block_size - 1) / block_size; - build_cluster_meta_kernel<<>>( + build_cluster_meta_kernel<<>>( d_cluster_meta_temp, d_histogram.data_handle(), d_offsets.data_handle(), num_centroids); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -708,11 +712,11 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, raft::copy(d_atomic_counters.data_handle(), d_offsets.data_handle(), num_centroids, stream_); num_blocks = (num_vectors + block_size - 1) / block_size; - scatter_pids_kernel<<>>(d_flat_pids, - device_cluster_ids, - d_offsets.data_handle(), - d_atomic_counters.data_handle(), - num_vectors); + scatter_pids_kernel<<>>(d_flat_pids, + device_cluster_ids, + d_offsets.data_handle(), + d_atomic_counters.data_handle(), + num_vectors); RAFT_CUDA_TRY(cudaPeekAtLastError()); // ------------------------- @@ -902,7 +906,7 @@ void sort_cluster_query_pairs(raft::resources const& handle, int batch_size, int nprobe) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); int total_pairs = batch_size * nprobe; // Allocate temporary arrays for sorting @@ -1033,7 +1037,7 @@ void IVFGPU::PrepareClusterSearchInputs( raft::device_vector& d_G_kbxSumq) { raft::resources const& searcher_handle = searcher.get_handle(); - rmm::cuda_stream_view searcher_stream = searcher.get_stream(); + cuda::stream_ref searcher_stream = searcher.get_stream(); const size_t batch_size = queries.extent(0); // Compute ||q - c||^2 = -2 * q . c + ||q||^2 + ||c||^2 into centroid_distances: @@ -1091,7 +1095,7 @@ void IVFGPU::PrepareClusterSearchInputs( d_G_k1xSumq = raft::make_device_vector(searcher_handle, batch_size); d_G_kbxSumq = raft::make_device_vector(searcher_handle, batch_size); computeQueryFactors( - queries, d_G_k1xSumq.view(), d_G_kbxSumq.view(), ex_bits, searcher_stream); + queries, d_G_k1xSumq.view(), d_G_kbxSumq.view(), ex_bits, searcher_stream.get()); // Sync here to ensure all outputs are visible before SearcherGPU reads them. raft::resource::sync_stream(searcher_handle); diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh index 75ddaee865..4f8c94380c 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cuh @@ -342,7 +342,7 @@ class IVFGPU { void AllocateHostMemory(); raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ // Device pointers for each data array. diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index af0876acca..8b80842034 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -501,7 +501,7 @@ void data_transformation_batch_opt(const float* d_data, int blockSize = D < 256 ? 128 : 256; size_t totalPadElements = (num_points + 1) * D; int gridPadSize = (totalPadElements + blockSize - 1) / blockSize; - gatherAndPadKernel<<>>( + gatherAndPadKernel<<>>( d_data, d_IDs, d_centroid, d_X_and_C_pad, num_points, DIM, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -526,13 +526,14 @@ void data_transformation_batch_opt(const float* d_data, size_t sharedMemSize = FusedBlockSize * sizeof(float); subtract_normalize_binarize_Kernel - <<>>(d_XP, // Input: Rotated data - d_CP, // Input: Rotated centroid - d_XP_output, // Output 1: Final residuals - d_XP_norm, // Output 2: Normalized residuals - d_bin_XP, // Output 3: Binarized data - num_points, - D); + <<>>( + d_XP, // Input: Rotated data + d_CP, // Input: Rotated centroid + d_XP_output, // Output 1: Final residuals + d_XP_norm, // Output 2: Normalized residuals + d_bin_XP, // Output 3: Binarized data + num_points, + D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -552,7 +553,7 @@ void rabitq_codes_and_factors_fused(const float* d_rotated_c, dim3 grid(num_points); dim3 block(threads_per_block); - pack_and_compute_factors_kernel<<>>( + pack_and_compute_factors_kernel<<>>( d_rotated_c, d_bin_XP, d_XP, @@ -592,17 +593,17 @@ void exrabitq_codes_and_factors_fused(const int* d_bin_XP, BlockSize * sizeof(float); // s_partials for reduction exrabitq_fused_kernel_batch - <<>>(d_bin_XP, - d_XP_norm, - d_XP, - d_centroid, - num_points, - D, - EX_BITS, - const_scaling_factor, - 1.9f, // kConstEpsilon - d_long_code, - d_ex_factor); + <<>>(d_bin_XP, + d_XP_norm, + d_XP, + d_centroid, + num_points, + D, + EX_BITS, + const_scaling_factor, + 1.9f, // kConstEpsilon + d_long_code, + d_ex_factor); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::resource::sync_stream(handle); } @@ -694,7 +695,7 @@ void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, int blockSize = D < 256 ? 128 : 256; size_t totalPadElements = (num_points + 1) * D; int gridPadSize = (totalPadElements + blockSize - 1) / blockSize; - gatherAndPadKernel<<>>( + gatherAndPadKernel<<>>( d_contiguous_data, nullptr, d_centroid, d_X_and_C_pad, num_points, DIM, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); @@ -719,13 +720,14 @@ void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, size_t sharedMemSize = FusedBlockSize * sizeof(float); subtract_normalize_binarize_Kernel - <<>>(d_XP, // Input: Rotated data - d_CP, // Input: Rotated centroid - d_XP_output, // Output 1: Final residuals - d_XP_norm, // Output 2: Normalized residuals - d_bin_XP, // Output 3: Binarized data - num_points, - D); + <<>>( + d_XP, // Input: Rotated data + d_CP, // Input: Rotated centroid + d_XP_output, // Output 1: Final residuals + d_XP_norm, // Output 2: Normalized residuals + d_bin_XP, // Output 3: Binarized data + num_points, + D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -1256,16 +1258,16 @@ void exrabitq_codes_and_factors_fused_ori(const int* d_bin_XP, } exrabitq_fused_kernel_batch_ori - <<>>(d_bin_XP, - d_XP_norm, - d_XP, - d_centroid, - num_points, - D, - EX_BITS, - 1.9f, // kConstEpsilon - d_long_code, - d_ex_factor); + <<>>(d_bin_XP, + d_XP_norm, + d_XP, + d_centroid, + num_points, + D, + EX_BITS, + 1.9f, // kConstEpsilon + d_long_code, + d_ex_factor); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::resource::sync_stream(handle); } @@ -1342,7 +1344,7 @@ float DataQuantizerGPU::get_const_scaling_factors_fully_gpu(size_t dim, size_t e unsigned long long seed = time(nullptr); auto kernel_fn = fully_fused_kernel; // decay to function pointer for std::atomic auto const& kernel_launcher = [&]() { - kernel_fn<<>>( + kernel_fn<<>>( d_factors.data(), kConstNum, dim, ex_bits, seed); }; cudaKernel_t cuda_kernel; @@ -1354,7 +1356,7 @@ float DataQuantizerGPU::get_const_scaling_factors_fully_gpu(size_t dim, size_t e // Use CUB for reduction - handles any size optimally size_t temp_storage_bytes = 0; cub::DeviceReduce::Sum( - nullptr, temp_storage_bytes, d_factors.data(), d_sum.data(), kConstNum, stream_); + nullptr, temp_storage_bytes, d_factors.data(), d_sum.data(), kConstNum, stream_.get()); { rmm::device_buffer d_temp_storage(temp_storage_bytes, stream_); @@ -1363,7 +1365,7 @@ float DataQuantizerGPU::get_const_scaling_factors_fully_gpu(size_t dim, size_t e d_factors.data(), d_sum.data(), kConstNum, - stream_); + stream_.get()); } RAFT_CUDA_TRY(cudaGetLastError()); diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh index 90ef71ed12..55a78cd2f3 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cuh @@ -155,7 +155,7 @@ class DataQuantizerGPU { // RAFT resources raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ // Device temporary buffers for quantization diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu index e095c35205..6ba4ff6388 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/rotator_gpu.cu @@ -28,7 +28,7 @@ RotatorGPU::RotatorGPU(raft::resources const& handle, uint32_t dim) raft::random::RngState rng(7ULL); raft::random::normal(handle, rng, rotation_matrix_.data_handle(), D * D, 0.0f, 1.0f); raft::linalg::detail::qrGetQ_inplace( - handle, rotation_matrix_.data_handle(), D, D, raft::resource::get_cuda_stream(handle)); + handle, rotation_matrix_.data_handle(), D, D, raft::resource::get_cuda_stream(handle).get()); } size_t RotatorGPU::size() const { return D; } diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu index 22bb74c682..14457335ed 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cu @@ -115,13 +115,14 @@ void launchPrecomputeLUTs(const float* d_query, float* d_lut_for_queries, size_t num_queries, size_t D, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { // Launch precompute kernel dim3 gridDim(num_queries, 1, 1); dim3 blockDim(256, 1, 1); // Can tune this - precomputeAllLUTs<<>>(d_query, d_lut_for_queries, num_queries, D); + precomputeAllLUTs<<>>( + d_query, d_lut_for_queries, num_queries, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -141,7 +142,7 @@ void SearcherGPU::SearchClusterQueryPairs( // First allocate space for LUT size_t lut_size = num_queries * (D / BITS_PER_CHUNK) * LUT_SIZE * sizeof(float); rmm::device_uvector d_lut_for_queries(lut_size / sizeof(float), stream_); - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_lut_for_queries.data(), d_lut_for_queries.data() + d_lut_for_queries.size(), -std::numeric_limits::infinity()); @@ -182,20 +183,20 @@ void SearcherGPU::SearchClusterQueryPairs( auto d_topk_pids = raft::make_device_matrix(handle_, num_queries, n_cols); // initialize distances - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_topk_dists.data_handle(), d_topk_dists.data_handle() + d_topk_dists.size(), std::numeric_limits::infinity()); rmm::device_uvector d_query_write_counters(num_queries, stream_); - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_query_write_counters.data(), d_query_write_counters.data() + num_queries, 0); rmm::device_uvector d_topk_threshold_batch(use_block_sort ? num_queries : 0, stream_); if (use_block_sort) { - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_topk_threshold_batch.data(), d_topk_threshold_batch.data() + num_queries, std::numeric_limits::infinity()); @@ -250,7 +251,7 @@ void SearcherGPU::SearchClusterQueryPairs( cur_ivf.get_ex_bits(), /*with_ex=*/true); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( - stream_, gridDim, blockDim, shared_mem_size, kernelParams); + stream_.get(), gridDim, blockDim, shared_mem_size, kernelParams); }; cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< compute_inner_products_with_lut_func_t>( @@ -266,7 +267,7 @@ void SearcherGPU::SearchClusterQueryPairs( /*ex_bits=*/0, /*with_ex=*/false); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( - stream_, gridDim, blockDim, shared_mem_size, kernelParams); + stream_.get(), gridDim, blockDim, shared_mem_size, kernelParams); }; cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< compute_inner_products_with_lut_func_t>( diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh index 73557f57ea..1535af97b1 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu.cuh @@ -52,7 +52,7 @@ class SearcherGPU { // Getter methods std::string const& get_mode() { return mode_; } raft::resources const& get_handle() const { return handle_; } - rmm::cuda_stream_view get_stream() const { return stream_; } + cuda::stream_ref get_stream() const { return stream_; } float* get_centroid_distances() { return centroid_distances_.data_handle(); } float* get_q_norms() { return q_norms_.data_handle(); } @@ -110,7 +110,7 @@ class SearcherGPU { private: raft::resources const& handle_; // reusable resource handle - rmm::cuda_stream_view stream_ = + cuda::stream_ref stream_ = raft::resource::get_cuda_stream(handle_); // CUDA stream obtained from handle_ size_t D; // number of dimension const float* query_ = nullptr; // rotated query (non-owning) diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu index 92906fb0e6..a69a13d820 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -356,28 +356,28 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( const int grid_size = num_queries; size_t shared_mem = D * sizeof(float) + block_size * sizeof(float); exrabitq_quantize_query - <<>>(queries.data_handle(), - num_queries, - D, - num_bits, - best_rescaling_factor, - 1.9f, - d_quantized_queries.data_handle(), - d_widths.data_handle()); + <<>>(queries.data_handle(), + num_queries, + D, + num_bits, + best_rescaling_factor, + 1.9f, + d_quantized_queries.data_handle(), + d_widths.data_handle()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } else { // scalar quantize // Step 1: Find min/max for each query const int block_size = 256; const int grid_size = num_queries; - findQueryRanges<<>>(queries.data_handle(), - d_query_ranges.data_handle(), - num_queries, - cur_ivf.get_num_padded_dim()); + findQueryRanges<<>>(queries.data_handle(), + d_query_ranges.data_handle(), + num_queries, + cur_ivf.get_num_padded_dim()); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Step 2: Quantize queries to int8_t with BQ=8 if (use_4bit) { - quantizeQueriesToInt4<<>>( + quantizeQueriesToInt4<<>>( queries.data_handle(), d_query_ranges.data_handle(), d_quantized_queries.data_handle(), @@ -386,7 +386,7 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( cur_ivf.get_num_padded_dim()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } else { - quantizeQueriesToInt8<<>>( + quantizeQueriesToInt8<<>>( queries.data_handle(), d_query_ranges.data_handle(), d_quantized_queries.data_handle(), @@ -404,14 +404,14 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( const int grid_size = (num_queries * num_bits * num_words + block_size - 1) / block_size; if (use_4bit) { - packInt4QueryBitPlanes<<>>( + packInt4QueryBitPlanes<<>>( d_quantized_queries.data_handle(), d_packed_queries.data_handle(), num_queries, cur_ivf.get_num_padded_dim()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } else { - packInt8QueryBitPlanes<<>>( + packInt8QueryBitPlanes<<>>( d_quantized_queries.data_handle(), d_packed_queries.data_handle(), num_queries, @@ -450,18 +450,18 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( auto d_topk_pids = raft::make_device_matrix(handle_, num_queries, n_cols); // initialize distances - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_topk_dists.data_handle(), d_topk_dists.data_handle() + d_topk_dists.size(), std::numeric_limits::infinity()); - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_query_write_counters.data_handle(), d_query_write_counters.data_handle() + num_queries, 0); if (use_block_sort) { - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_topk_threshold_batch.data_handle(), d_topk_threshold_batch.data_handle() + num_queries, std::numeric_limits::infinity()); @@ -526,7 +526,7 @@ void SearcherGPU::SearchClusterQueryPairsQuantizeQuery( cur_ivf.get_ex_bits(), with_ex); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( - stream_, gridDim, blockDim, shared_mem_size, kernelParams); + stream_.get(), gridDim, blockDim, shared_mem_size, kernelParams); }; cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< compute_inner_products_with_lut_func_t>( diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu index 8357af5fbb..d3a05434b7 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_shared_mem_opt.cu @@ -102,7 +102,7 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( // Initialize with -infinity (convert to FP16) lut_dtype neg_inf_fp16 = __float2half(-std::numeric_limits::infinity()); - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_lut_for_queries.data(), d_lut_for_queries.data() + lut_elements, neg_inf_fp16); @@ -112,7 +112,7 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( d_lut_for_queries.data(), num_queries, cur_ivf.get_num_padded_dim(), - stream_); + stream_.get()); // check if the inner products kernel should use block sort to keep a top-k priority queue vs. // outputting distances from all vectors in probed clusters @@ -148,20 +148,20 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( auto d_topk_pids = raft::make_device_matrix(handle_, num_queries, n_cols); // initialize distances - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_topk_dists.data_handle(), d_topk_dists.data_handle() + d_topk_dists.size(), std::numeric_limits::infinity()); rmm::device_uvector d_query_write_counters(num_queries, stream_); - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_query_write_counters.data(), d_query_write_counters.data() + num_queries, 0); rmm::device_uvector d_topk_threshold_batch(use_block_sort ? num_queries : 0, stream_); if (use_block_sort) { - thrust::fill(thrust::cuda::par.on(stream_), + thrust::fill(thrust::cuda::par.on(stream_.get()), d_topk_threshold_batch.data(), d_topk_threshold_batch.data() + num_queries, std::numeric_limits::infinity()); @@ -221,7 +221,7 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( cur_ivf.get_ex_bits(), /*with_ex=*/true); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( - stream_, gridDim, blockDim, shared_mem_size, kernelParams); + stream_.get(), gridDim, blockDim, shared_mem_size, kernelParams); }; cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< compute_inner_products_with_lut_func_t>( @@ -240,7 +240,7 @@ void SearcherGPU::SearchClusterQueryPairsSharedMemOpt( /*ex_bits=*/0, /*with_ex=*/false); auto const& kernel_launcher = [&]() -> void { jit_launcher->dispatch( - stream_, gridDim, blockDim, shared_mem_size, kernelParams); + stream_.get(), gridDim, blockDim, shared_mem_size, kernelParams); }; cuvs::neighbors::detail::safely_launch_kernel_with_smem_size< compute_inner_products_with_lut_func_t>( diff --git a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu index 2bb9fb2174..5ff99326b5 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu +++ b/cpp/src/neighbors/ivf_rabitq/utils/searcher_gpu_utils.cu @@ -40,8 +40,10 @@ void get_max_probed_cluster_size_and_vectors_count( auto d_max_probed_cluster_size_ptr = d_max_probed_cluster_size.data_handle(); auto d_probed_vectors_count_ptr = d_probed_vectors_count.data_handle(); if (get_max_probed_vectors_count) { - RAFT_CUDA_TRY(cudaMemsetAsync( - d_probed_vectors_count_ptr, 0, num_queries * sizeof(size_t), stream)); // Initialize to 0 + RAFT_CUDA_TRY(cudaMemsetAsync(d_probed_vectors_count_ptr, + 0, + num_queries * sizeof(size_t), + stream.get())); // Initialize to 0 } auto count = thrust::make_counting_iterator(0); diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh index 9d8bb20764..88d2dfda3a 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_build.cuh @@ -342,9 +342,9 @@ void extend_inplace(raft::resources const& handle, new_labels.data_handle(), n_rows, 1, - stream); + stream.get()); raft::linalg::add( - list_sizes_ptr, list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream); + list_sizes_ptr, list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream.get()); std::vector new_list_sizes(n_lists); std::vector old_list_sizes(n_lists); @@ -385,18 +385,18 @@ void extend_inplace(raft::resources const& handle, const dim3 block_dim(kEncodeBlockSize); const dim3 grid_dim(raft::ceildiv(bs, int64_t(kEncodeWarpsPerBlk))); encode_and_fill_kernel - <<>>(new_labels.data_handle() + batch.offset(), - batch.data(), - index->centers().data_handle(), - idx_batch->data(), - index->data_ptrs().data_handle(), - index->inds_ptrs().data_handle(), - list_sizes_ptr, - index->sq_vmin().data_handle(), - index->sq_delta().data_handle(), - bs, - dim, - batch.offset()); + <<>>(new_labels.data_handle() + batch.offset(), + batch.data(), + index->centers().data_handle(), + idx_batch->data(), + index->data_ptrs().data_handle(), + index->inds_ptrs().data_handle(), + list_sizes_ptr, + index->sq_vmin().data_handle(), + index->sq_delta().data_handle(), + bs, + dim, + batch.offset()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -419,11 +419,14 @@ void extend_inplace(raft::resources const& handle, index->centers().data_handle(), dim, n_lists, - stream, + stream.get(), raft::sqrt_op{}); } else { - raft::linalg::rowNorm( - index->center_norms()->data_handle(), index->centers().data_handle(), dim, n_lists, stream); + raft::linalg::rowNorm(index->center_norms()->data_handle(), + index->centers().data_handle(), + dim, + n_lists, + stream.get()); } }; @@ -492,11 +495,11 @@ inline auto build( constexpr int kResidualBlockSize = 256; compute_residuals_inplace_kernel - <<>>(trainset.data_handle(), - idx.centers().data_handle(), - train_labels.data_handle(), - n_rows_train, - dim); + <<>>(trainset.data_handle(), + idx.centers().data_handle(), + train_labels.data_handle(), + n_rows_train, + dim); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -510,7 +513,7 @@ inline auto build( constexpr int kMinMaxBlockSize = 256; launch_fused_column_minmax( - residuals.data_handle(), vmin_ptr, vmax_ptr, n_rows_train, dim, stream); + residuals.data_handle(), vmin_ptr, vmax_ptr, n_rows_train, dim, stream.get()); RAFT_CUDA_TRY(cudaPeekAtLastError()); // Expand the observed range by a small margin to reduce clipping on unseen data, diff --git a/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh b/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh index 3f0dfc62d9..246a82c892 100644 --- a/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh +++ b/cpp/src/neighbors/ivf_sq/ivf_sq_search.cuh @@ -149,7 +149,7 @@ void launch_kernel(const index& idx, float* out_distances, uint32_t* out_indices, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream, + cuda::stream_ref stream, IvfSampleFilterT sample_filter) { static_assert(std::is_same_v, "IVF-SQ JIT-LTO scan only supports CodeT=uint8_t"); @@ -247,7 +247,7 @@ void launch_kernel(const index& idx, const float* sq_vmin = idx.sq_vmin().data_handle(); const float* sq_delta = idx.sq_delta().data_handle(); - kernel_launcher->dispatch>(stream, + kernel_launcher->dispatch>(stream.get(), grid, block, smem, @@ -291,7 +291,7 @@ void ivf_sq_scan(raft::resources const& handle, uint32_t* out_indices, IvfSampleFilterT sample_filter, uint32_t& grid_dim_x, - rmm::cuda_stream_view stream) + cuda::stream_ref stream) { // Determine the fused top-k capacity (0 = disabled / fallback to materialization) int capacity = is_local_topk_feasible(k) ? raft::bound_by_power_of_two(int(k)) : 0; @@ -390,7 +390,7 @@ void search_impl(raft::resources const& handle, converted_queries_ptr = const_cast(queries); } else { raft::linalg::unaryOp( - converted_queries_ptr, queries, n_queries * dim, utils::mapping{}, stream); + converted_queries_ptr, queries, n_queries * dim, utils::mapping{}, stream.get()); } auto distance_buffer_dev_view = raft::make_device_matrix_view( @@ -412,7 +412,7 @@ void search_impl(raft::resources const& handle, converted_queries_ptr, static_cast(dim), static_cast(n_queries), - stream); + stream.get()); utils::outer_add(query_norm_dev->data_handle(), (int64_t)n_queries, index.center_norms()->data_handle(), @@ -426,7 +426,7 @@ void search_impl(raft::resources const& handle, converted_queries_ptr, static_cast(dim), static_cast(n_queries), - stream, + stream.get(), raft::sqrt_op{}); alpha = -1.0f; beta = 0.0f; @@ -454,7 +454,7 @@ void search_impl(raft::resources const& handle, &beta, distance_buffer_dev.data(), index.n_lists(), - stream); + stream.get()); if (index.metric() == cuvs::distance::DistanceType::CosineExpanded) { auto n_lists_local = index.n_lists(); diff --git a/cpp/src/neighbors/ivf_sq_index.cpp b/cpp/src/neighbors/ivf_sq_index.cpp index c4a75b84e3..7a6501669f 100644 --- a/cpp/src/neighbors/ivf_sq_index.cpp +++ b/cpp/src/neighbors/ivf_sq_index.cpp @@ -46,12 +46,12 @@ index::index(raft::resources const& res, check_consistency(); auto stream = raft::resource::get_cuda_stream(res); std::memset(accum_sorted_sizes_.data_handle(), 0, accum_sorted_sizes_.size() * sizeof(int64_t)); + RAFT_CUDA_TRY(cudaMemsetAsync( + list_sizes_.data_handle(), 0, list_sizes_.size() * sizeof(uint32_t), stream.get())); RAFT_CUDA_TRY( - cudaMemsetAsync(list_sizes_.data_handle(), 0, list_sizes_.size() * sizeof(uint32_t), stream)); - RAFT_CUDA_TRY( - cudaMemsetAsync(data_ptrs_.data_handle(), 0, data_ptrs_.size() * sizeof(CodeT*), stream)); - RAFT_CUDA_TRY( - cudaMemsetAsync(inds_ptrs_.data_handle(), 0, inds_ptrs_.size() * sizeof(int64_t*), stream)); + cudaMemsetAsync(data_ptrs_.data_handle(), 0, data_ptrs_.size() * sizeof(CodeT*), stream.get())); + RAFT_CUDA_TRY(cudaMemsetAsync( + inds_ptrs_.data_handle(), 0, inds_ptrs_.size() * sizeof(int64_t*), stream.get())); } template diff --git a/cpp/src/neighbors/knn_merge_parts.cu b/cpp/src/neighbors/knn_merge_parts.cu index da7110a475..ebf7083f03 100644 --- a/cpp/src/neighbors/knn_merge_parts.cu +++ b/cpp/src/neighbors/knn_merge_parts.cu @@ -28,7 +28,7 @@ void _knn_merge_parts(raft::resources const& res, rows, parts, k, - raft::resource::get_cuda_stream(res), + raft::resource::get_cuda_stream(res).get(), translations.data_handle()); } } // namespace diff --git a/cpp/src/neighbors/mg/snmg.cuh b/cpp/src/neighbors/mg/snmg.cuh index e0ad8509f8..9e66f58898 100644 --- a/cpp/src/neighbors/mg/snmg.cuh +++ b/cpp/src/neighbors/mg/snmg.cuh @@ -310,13 +310,13 @@ void sharded_search_with_direct_merge( ncclUint8, from_rank, raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); ncclRecv(in_distances.data_handle() + batch_offset, part_size * sizeof(float), ncclUint8, from_rank, raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); } ncclGroupEnd(); resource::sync_stream(dev_res); @@ -335,13 +335,13 @@ void sharded_search_with_direct_merge( ncclUint8, raft::resource::get_root_rank(clique), raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); ncclSend(d_distances.data_handle(), part_size * sizeof(float), ncclUint8, raft::resource::get_root_rank(clique), raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); ncclGroupEnd(); resource::sync_stream(dev_res); } @@ -464,7 +464,7 @@ void sharded_search_with_tree_merge( neighbors_view.data_handle(), translation_offset, part_size, - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); auto d_trans = raft::make_device_vector(dev_res, 2); raft::matrix::fill(dev_res, d_trans.view(), searchIdxT(0)); @@ -486,13 +486,13 @@ void sharded_search_with_tree_merge( ncclUint8, other_id, raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); ncclRecv(tmp_distances.data_handle() + part_size, part_size * sizeof(float), ncclUint8, other_id, raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); received_something = true; } } else if (rank % radix == offset) // This is one of the senders @@ -503,13 +503,13 @@ void sharded_search_with_tree_merge( ncclUint8, other_id, raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); ncclSend(tmp_distances.data_handle(), part_size * sizeof(float), ncclUint8, other_id, raft::resource::get_nccl_comm_for_rank(clique, rank), - raft::resource::get_cuda_stream(dev_res)); + raft::resource::get_cuda_stream(dev_res).get()); } ncclGroupEnd(); diff --git a/cpp/src/neighbors/scann/detail/scann_avq.cuh b/cpp/src/neighbors/scann/detail/scann_avq.cuh index 0e138ce1a3..d3b812b592 100644 --- a/cpp/src/neighbors/scann/detail/scann_avq.cuh +++ b/cpp/src/neighbors/scann/detail/scann_avq.cuh @@ -57,7 +57,7 @@ void compute_cluster_offsets(raft::resources const& dev_resources, raft::device_vector_view cluster_sizes, int64_t& max_cluster_size) { - cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources); + cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources).get(); rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(dev_resources); @@ -136,7 +136,7 @@ void sum_reduce_vector(raft::resources const& dev_resources, raft::device_vector_view v, raft::device_scalar_view s) { - cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources); + cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources).get(); rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(dev_resources); @@ -163,7 +163,7 @@ void cholesky_solver(raft::resources const& dev_resources, raft::device_vector_view b, raft::device_vector_view x) { - cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources); + cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources).get(); cusolverDnHandle_t cusolverH = raft::resource::get_cusolver_dn_handle(dev_resources); rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource_ref(dev_resources); @@ -280,7 +280,8 @@ void compute_avq_centroid(raft::resources const& dev_resources, raft::linalg::detail::cublas_device_pointer_mode pm(cublas_handle); - RAFT_CUBLAS_TRY(cublasSetStream(cublas_handle, raft::resource::get_cuda_stream(dev_resources))); + RAFT_CUBLAS_TRY( + cublasSetStream(cublas_handle, raft::resource::get_cuda_stream(dev_resources).get())); RAFT_CUBLAS_TRY(cublasSgemm(cublas_handle, cublasOperation_t::CUBLAS_OP_T, @@ -581,7 +582,7 @@ void apply_avq(raft::resources const& res, { // Compute clusters - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); auto cluster_offsets = raft::make_device_vector(res, centroids_view.extent(0)); auto clusters = raft::make_device_vector(res, dataset.extent(0)); int64_t max_cluster_size = 0; diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index c01e50bc83..141a265783 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -45,7 +45,7 @@ index build( const index_params& params, raft::mdspan, raft::row_major, Accessor> dataset) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); IdxT dim = dataset.extent(1); RAFT_LOG_DEBUG("Creating empty index"); @@ -138,7 +138,7 @@ index build( centroids_view, raft::make_const_mdspan(labels_view), params.partitioning_eta, - copy_stream); + copy_stream.get()); raft::device_vector_view soar_labels_view = idx.soar_labels(); diff --git a/cpp/src/neighbors/scann/detail/scann_quantize.cuh b/cpp/src/neighbors/scann/detail/scann_quantize.cuh index 95025f7f57..582e4c40cd 100644 --- a/cpp/src/neighbors/scann/detail/scann_quantize.cuh +++ b/cpp/src/neighbors/scann/detail/scann_quantize.cuh @@ -289,7 +289,7 @@ void quantize_bfloat16_noise_shaped(raft::resources const& res, raft::device_matrix_view bf16_dataset, float noise_shaping_threshold) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); IdxT n_rows = dataset.extent(0); auto norms = raft::make_device_vector(res, n_rows); @@ -375,7 +375,7 @@ auto sample_training_residuals( dataset, raft::make_const_mdspan(train_indices.view()), trainset.view(), - raft::resource::get_cuda_stream(res)); + raft::resource::get_cuda_stream(res).get()); // Considering labels as a single column matrix for use in gather auto labels_view = diff --git a/cpp/src/preprocessing/quantize/detail/binary.cuh b/cpp/src/preprocessing/quantize/detail/binary.cuh index 6695b8b126..310daaf3bb 100644 --- a/cpp/src/preprocessing/quantize/detail/binary.cuh +++ b/cpp/src/preprocessing/quantize/detail/binary.cuh @@ -193,7 +193,7 @@ auto train(raft::resources const& res, dataset_dim, dataset_dim, dataset_size, - raft::resource::get_cuda_stream(res)); + raft::resource::get_cuda_stream(res).get()); } } else if (params.threshold == cuvs::preprocessing::quantize::binary::bit_threshold::sampling_median) { @@ -352,7 +352,7 @@ void transform(raft::resources const& res, raft::device_matrix_view dataset, raft::device_matrix_view out) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); const uint32_t bits_per_pack = sizeof(QuantI) * 8; const uint32_t dataset_dim = dataset.extent(1); const uint32_t out_dim = out.extent(1); diff --git a/cpp/src/preprocessing/quantize/detail/pq.cuh b/cpp/src/preprocessing/quantize/detail/pq.cuh index 7fea89461a..8a2a6c3cc9 100644 --- a/cpp/src/preprocessing/quantize/detail/pq.cuh +++ b/cpp/src/preprocessing/quantize/detail/pq.cuh @@ -328,7 +328,7 @@ auto reconstruct_vectors( } }(pq_bits); dim3 blocks(raft::div_rounding_up_safe(n_rows, kBlockSize / threads_per_vec), 1, 1); - kernel<<>>( + kernel<<>>( codes, out_vectors, pq_centers, vq_centers, vq_labels, pq_bits, use_subspaces); RAFT_CUDA_TRY(cudaPeekAtLastError()); diff --git a/cpp/src/preprocessing/quantize/detail/scalar.cuh b/cpp/src/preprocessing/quantize/detail/scalar.cuh index 63a55aaf54..1fcb8f2f0f 100644 --- a/cpp/src/preprocessing/quantize/detail/scalar.cuh +++ b/cpp/src/preprocessing/quantize/detail/scalar.cuh @@ -73,7 +73,7 @@ std::tuple quantile_min_max( constexpr size_t max_num_samples = 1000000; constexpr int seed = 137; - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); // select subsample raft::random::RngState rng(seed); @@ -145,7 +145,7 @@ void transform(raft::resources const& res, raft::device_matrix_view dataset, raft::device_matrix_view out) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); raft::linalg::map(res, out, @@ -174,7 +174,7 @@ void inverse_transform(raft::resources const& res, raft::device_matrix_view dataset, raft::device_matrix_view out) { - cudaStream_t stream = raft::resource::get_cuda_stream(res); + cudaStream_t stream = raft::resource::get_cuda_stream(res).get(); raft::linalg::map(res, out, diff --git a/cpp/src/preprocessing/spectral/detail/spectral_embedding.cuh b/cpp/src/preprocessing/spectral/detail/spectral_embedding.cuh index 04c3e83b0a..bec3ba2d45 100644 --- a/cpp/src/preprocessing/spectral/detail/spectral_embedding.cuh +++ b/cpp/src/preprocessing/spectral/detail/spectral_embedding.cuh @@ -194,7 +194,7 @@ void create_connectivity_graph( sym_coo1_matrix.structure_view().get_rows().data(), sym_coo1_matrix.structure_view().get_cols().data(), sym_coo1_matrix.get_elements().data(), - stream); + stream.get()); raft::sparse::op::coo_remove_scalar<128, float, int, NNZType>( handle, diff --git a/cpp/src/sparse/neighbors/detail/cross_component_nn.cuh b/cpp/src/sparse/neighbors/detail/cross_component_nn.cuh index 9bfc593559..c9209281d9 100644 --- a/cpp/src/sparse/neighbors/detail/cross_component_nn.cuh +++ b/cpp/src/sparse/neighbors/detail/cross_component_nn.cuh @@ -312,7 +312,7 @@ void perform_1nn(raft::resources const& handle, raft::matrix::gather(handle, X_mutable_view, sort_plan_const_view, (value_idx)col_batch_size); // Get the number of unique components from the array of colors - value_idx n_components = get_n_components(colors, n_rows, stream); + value_idx n_components = get_n_components(colors, n_rows, stream.get()); // colors_group_idxs is an array containing the *end* indices of each color // component in colors. That is, the value of colors_group_idxs[j] indicates @@ -320,7 +320,7 @@ void perform_1nn(raft::resources const& handle, // the color components. auto colors_group_idxs = raft::make_device_vector(handle, n_components + 1); raft::sparse::convert::sorted_coo_to_csr( - colors, n_rows, colors_group_idxs.data_handle(), n_components + 1, stream); + colors, n_rows, colors_group_idxs.data_handle(), n_components + 1, stream.get()); auto group_idxs_view = raft::make_device_vector_view( colors_group_idxs.data_handle() + 1, n_components); @@ -562,7 +562,7 @@ void cross_component_nn( // Normalize colors so they are drawn from a monotonically increasing set constexpr bool zero_based = true; raft::label::make_monotonic( - colors.data(), const_cast(orig_colors), n_rows, stream, zero_based); + colors.data(), const_cast(orig_colors), n_rows, stream.get(), zero_based); /** * First compute 1-nn for all colors where the color of each data point @@ -597,7 +597,7 @@ void cross_component_nn( // Compute mask of duplicates rmm::device_uvector out_index(n_rows + 1, stream); raft::sparse::op::compute_duplicates_mask( - out_index.data(), colors.data(), nn_colors.data(), n_rows, stream); + out_index.data(), colors.data(), nn_colors.data(), n_rows, stream.get()); thrust::exclusive_scan(raft::resource::get_thrust_policy(handle), out_index.data(), @@ -614,11 +614,11 @@ void cross_component_nn( size++; - raft::sparse::COO min_edges(stream); - min_edges.allocate(size, n_rows, n_rows, true, stream); + raft::sparse::COO min_edges(stream.get()); + min_edges.allocate(size, n_rows, n_rows, true, stream.get()); min_components_by_color( - min_edges, out_index.data(), src_indices.data(), temp_inds_dists.data(), n_rows, stream); + min_edges, out_index.data(), src_indices.data(), temp_inds_dists.data(), n_rows, stream.get()); /** * Symmetrize resulting edge list diff --git a/cpp/src/stats/detail/batched/silhouette_score.cuh b/cpp/src/stats/detail/batched/silhouette_score.cuh index f41b843d5c..9b40eaf61c 100644 --- a/cpp/src/stats/detail/batched/silhouette_score.cuh +++ b/cpp/src/stats/detail/batched/silhouette_score.cuh @@ -116,7 +116,8 @@ rmm::device_uvector get_cluster_counts(raft::resources const& handle, rmm::device_uvector workspace(1, stream); - cuvs::stats::detail::countLabels(y, cluster_counts.data(), n_rows, n_labels, workspace, stream); + cuvs::stats::detail::countLabels( + y, cluster_counts.data(), n_rows, n_labels, workspace, stream.get()); return cluster_counts; } @@ -206,7 +207,7 @@ value_t silhouette_score( dim3 block_size(std::min(n_rows, 32), std::min(n_labels, 32)); dim3 grid_size(raft::ceildiv(n_rows, (value_idx)block_size.x), raft::ceildiv(n_labels, (label_idx)block_size.y)); - detail::fill_b_kernel<<>>( + detail::fill_b_kernel<<>>( b_ptr, y, n_rows, n_labels, cluster_counts.data()); raft::resource::wait_stream_pool_on_stream(handle); @@ -241,7 +242,7 @@ value_t silhouette_score( distances.data(), n_left_rows, n_right_rows, - chunk_stream); + chunk_stream.get()); } } diff --git a/cpp/src/stats/detail/trustworthiness_score.cuh b/cpp/src/stats/detail/trustworthiness_score.cuh index f44404fefd..21ee3a465d 100644 --- a/cpp/src/stats/detail/trustworthiness_score.cuh +++ b/cpp/src/stats/detail/trustworthiness_score.cuh @@ -124,7 +124,7 @@ double trustworthiness_score(const raft::resources& h, int n_neighbors, int batchSize = 512) { - cudaStream_t stream = raft::resource::get_cuda_stream(h); + cudaStream_t stream = raft::resource::get_cuda_stream(h).get(); const int KNN_ALLOC = n * (n_neighbors + 1); rmm::device_uvector emb_ind(KNN_ALLOC, stream); diff --git a/cpp/src/stats/silhouette_score.cu b/cpp/src/stats/silhouette_score.cu index 9ffe98a6af..bc5c765526 100644 --- a/cpp/src/stats/silhouette_score.cu +++ b/cpp/src/stats/silhouette_score.cu @@ -35,7 +35,7 @@ value_t _silhouette_score( labels.data_handle(), n_unique_labels, silhouette_score_per_sample_ptr, - raft::resource::get_cuda_stream(handle), + raft::resource::get_cuda_stream(handle).get(), metric); } diff --git a/cpp/tests/cluster/connect_knn.cu b/cpp/tests/cluster/connect_knn.cu index 57162d4e1e..4f24b50b27 100644 --- a/cpp/tests/cluster/connect_knn.cu +++ b/cpp/tests/cluster/connect_knn.cu @@ -43,7 +43,7 @@ template class ConnectKNNTest : public ::testing::TestWithParam { public: ConnectKNNTest() - : stream(handle.get_stream()), + : stream(handle.get_stream().get()), ps(::testing::TestWithParam::GetParam()), database(0, stream) { diff --git a/cpp/tests/cluster/gmm.cu b/cpp/tests/cluster/gmm.cu index d654dacfe4..8da610c6ea 100644 --- a/cpp/tests/cluster/gmm.cu +++ b/cpp/tests/cluster/gmm.cu @@ -51,7 +51,7 @@ inline int64_t cov_len(covariance_type ct, int d, int K) template class GMMTest : public ::testing::TestWithParam> { protected: - GMMTest() : stream(raft::resource::get_cuda_stream(handle)) {} + GMMTest() : stream(raft::resource::get_cuda_stream(handle).get()) {} void basicTest() { @@ -239,7 +239,7 @@ std::pair, raft::device_vector> make_g n, d, K, - raft::resource::get_cuda_stream(handle), + raft::resource::get_cuda_stream(handle).get(), true, nullptr, nullptr, @@ -305,7 +305,7 @@ TEST(GMMExtra, InitMethods) ASSERT_TRUE(std::isfinite(lb)) << "init " << (int)im; if (im == init_method::KMeans || im == init_method::KMeansPlusPlus) { double ari = - raft::stats::adjusted_rand_index(yref.data_handle(), labels.data_handle(), n, stream); + raft::stats::adjusted_rand_index(yref.data_handle(), labels.data_handle(), n, stream.get()); raft::resource::sync_stream(handle, stream); ASSERT_GT(ari, 0.95) << "init " << (int)im; } @@ -608,7 +608,7 @@ TEST(GMMExtra, IllDefinedCovarianceThrows) // All points identical -> any component covariance collapses to zero. auto X = raft::make_device_matrix(handle, n, d); - RAFT_CUDA_TRY(cudaMemsetAsync(X.data_handle(), 0, sizeof(float) * (size_t)n * d, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(X.data_handle(), 0, sizeof(float) * (size_t)n * d, stream.get())); int64_t cn = cov_len(covariance_type::FULL, d, K); auto weights = raft::make_device_vector(handle, K); diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 59051484f4..a0eaee1eed 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -269,15 +269,17 @@ class KmeansTest : public ::testing::TestWithParam> { raft::resource::sync_stream(handle, stream); - score = raft::stats::adjusted_rand_index( - d_labels_ref.data(), d_labels.data(), n_samples, raft::resource::get_cuda_stream(handle)); + score = raft::stats::adjusted_rand_index(d_labels_ref.data(), + d_labels.data(), + n_samples, + raft::resource::get_cuda_stream(handle).get()); if (score < 1.0) { std::stringstream ss; - ss << "Expected: " << raft::arr2Str(d_labels_ref.data(), 25, "d_labels_ref", stream); + ss << "Expected: " << raft::arr2Str(d_labels_ref.data(), 25, "d_labels_ref", stream.get()); std::cout << (ss.str().c_str()) << '\n'; ss.str(std::string()); - ss << "Actual: " << raft::arr2Str(d_labels.data(), 25, "d_labels", stream); + ss << "Actual: " << raft::arr2Str(d_labels.data(), 25, "d_labels", stream.get()); std::cout << (ss.str().c_str()) << '\n'; std::cout << "Score = " << score << '\n'; } @@ -443,7 +445,7 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(T(1e-2)), - stream); + stream.get()); T ref_pred_inertia = 0; cuvs::cluster::kmeans::predict(handle, @@ -470,14 +472,15 @@ class KmeansFitBatchedTest : public ::testing::TestWithParamdata_handle(), d_labels->data_handle(), n_samples, - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); if (score < 0.99) { std::stringstream ss; - ss << "Expected: " << raft::arr2Str(d_labels_ref->data_handle(), 25, "d_labels_ref", stream); + ss << "Expected: " + << raft::arr2Str(d_labels_ref->data_handle(), 25, "d_labels_ref", stream.get()); std::cout << (ss.str().c_str()) << '\n'; ss.str(std::string()); - ss << "Actual: " << raft::arr2Str(d_labels->data_handle(), 25, "d_labels", stream); + ss << "Actual: " << raft::arr2Str(d_labels->data_handle(), 25, "d_labels", stream.get()); std::cout << (ss.str().c_str()) << '\n'; std::cout << "Score = " << score << '\n'; } diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index 10d4ce9cbb..c9e109e24d 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -58,7 +58,7 @@ template > { protected: KmeansBalancedTest() - : stream(raft::resource::get_cuda_stream(handle)), + : stream(raft::resource::get_cuda_stream(handle).get()), d_labels(0, stream), d_labels_ref(0, stream), d_centroids(0, stream) @@ -128,8 +128,10 @@ class KmeansBalancedTest : public ::testing::TestWithParam class KmeansFindKTest : public ::testing::TestWithParam> { protected: KmeansFindKTest() - : stream(raft::resource::get_cuda_stream(handle)), best_k(raft::make_host_scalar(0)) + : stream(raft::resource::get_cuda_stream(handle).get()), best_k(raft::make_host_scalar(0)) { } diff --git a/cpp/tests/cluster/kmeans_mg.cu b/cpp/tests/cluster/kmeans_mg.cu index 85bf238d33..811da72e35 100644 --- a/cpp/tests/cluster/kmeans_mg.cu +++ b/cpp/tests/cluster/kmeans_mg.cu @@ -198,10 +198,10 @@ class KmeansSNMGTest : public ::testing::TestWithParam> { raft::make_host_scalar_view(&pred_inertia_sg)); ari_vs_ref_ = raft::stats::adjusted_rand_index( - d_labels_ref.data(), d_labels_snmg.data(), n_samples, sg_stream); + d_labels_ref.data(), d_labels_snmg.data(), n_samples, sg_stream.get()); ari_vs_sg_ = raft::stats::adjusted_rand_index( - d_labels_sg.data(), d_labels_snmg.data(), n_samples, sg_stream); + d_labels_sg.data(), d_labels_snmg.data(), n_samples, sg_stream.get()); snmg_inertia_ = snmg_inertia; sg_inertia_ = sg_inertia; diff --git a/cpp/tests/cluster/kmeans_test_blobs.cuh b/cpp/tests/cluster/kmeans_test_blobs.cuh index 45cdc41778..79091512a8 100644 --- a/cpp/tests/cluster/kmeans_test_blobs.cuh +++ b/cpp/tests/cluster/kmeans_test_blobs.cuh @@ -84,7 +84,7 @@ kmeans_blob_inputs make_kmeans_blob_inputs(raft::resources const& handle, n_samples, n_features, n_clusters, - stream, + stream.get(), /* row_major */ true, /* centers */ nullptr, /* cluster_std */ nullptr, diff --git a/cpp/tests/cluster/linkage.cu b/cpp/tests/cluster/linkage.cu index 43824b16ad..25347956c2 100644 --- a/cpp/tests/cluster/linkage.cu +++ b/cpp/tests/cluster/linkage.cu @@ -199,7 +199,7 @@ class LinkageTest : public ::testing::TestWithParam> { raft::resource::sync_stream(handle, stream); - score = compute_rand_index(labels.data(), labels_ref.data(), params.n_row, stream); + score = compute_rand_index(labels.data(), labels_ref.data(), params.n_row, stream.get()); } void SetUp() override { basicTest(); } diff --git a/cpp/tests/cluster/spectral.cu b/cpp/tests/cluster/spectral.cu index 9c55cef766..59c65b0b13 100644 --- a/cpp/tests/cluster/spectral.cu +++ b/cpp/tests/cluster/spectral.cu @@ -67,7 +67,7 @@ class SpectralClusteringTest : public ::testing::TestWithParam> { true, metric, 0.0, - stream); + stream.get()); } } @@ -162,7 +164,7 @@ class NNTest : public ::testing::TestWithParam> { private: raft::resources handle; - rmm::cuda_stream_view stream; + cuda::stream_ref stream; NNInputs params_; ComparisonSummary summary; IdxT m; diff --git a/cpp/tests/neighbors/epsilon_neighborhood.cu b/cpp/tests/neighbors/epsilon_neighborhood.cu index d5695c60b2..d39deb5e22 100644 --- a/cpp/tests/neighbors/epsilon_neighborhood.cu +++ b/cpp/tests/neighbors/epsilon_neighborhood.cu @@ -55,13 +55,13 @@ class EpsNeighTest : public ::testing::TestWithParam> { batchSize = param.n_row / param.n_batches; adj.resize(param.n_row * batchSize, stream); vd.resize(batchSize + 1, stream); - RAFT_CUDA_TRY(cudaMemsetAsync(vd.data(), 0, vd.size() * sizeof(IdxT), stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(vd.data(), 0, vd.size() * sizeof(IdxT), stream.get())); raft::random::make_blobs(data.data(), labels.data(), param.n_row, param.n_col, param.n_centers, - stream, + stream.get(), true, nullptr, nullptr, @@ -450,11 +450,11 @@ TEST(EpsNeighborhood, LargeNDimension) raft::matrix::fill(handle, adj_expected.view(), true); auto stream = raft::resource::get_cuda_stream(handle); ASSERT_TRUE(cuvs::devArrMatch( - adj_expected.data_handle(), adj.data_handle(), m * n, cuvs::Compare(), stream)); + adj_expected.data_handle(), adj.data_handle(), m * n, cuvs::Compare(), stream.get())); int64_t expected_vd0 = n; ASSERT_TRUE( - cuvs::devArrMatch(&expected_vd0, vd.data_handle(), 1, cuvs::Compare(), stream)); + cuvs::devArrMatch(&expected_vd0, vd.data_handle(), 1, cuvs::Compare(), stream.get())); } }; // namespace cuvs::neighbors::epsilon_neighborhood diff --git a/cpp/tests/neighbors/hnsw.cu b/cpp/tests/neighbors/hnsw.cu index ac86c19123..66a37b527d 100644 --- a/cpp/tests/neighbors/hnsw.cu +++ b/cpp/tests/neighbors/hnsw.cu @@ -159,7 +159,7 @@ class AnnHNSWTest : public ::testing::TestWithParam { private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnHNSWInputs ps; rmm::device_uvector database; rmm::device_uvector queries; diff --git a/cpp/tests/neighbors/knn_merge_parts.cu b/cpp/tests/neighbors/knn_merge_parts.cu index 65db1a08ce..332959ee4a 100644 --- a/cpp/tests/neighbors/knn_merge_parts.cu +++ b/cpp/tests/neighbors/knn_merge_parts.cu @@ -93,7 +93,7 @@ void run_merge(bool select_min, n_queries, k, 0.0f, - stream, + stream.get(), true)); } diff --git a/cpp/tests/neighbors/naive_knn.cuh b/cpp/tests/neighbors/naive_knn.cuh index d21e982902..ff822e52c9 100644 --- a/cpp/tests/neighbors/naive_knn.cuh +++ b/cpp/tests/neighbors/naive_knn.cuh @@ -121,7 +121,7 @@ void naive_knn(raft::resources const& handle, size_t batch_size = std::min(max_batch_size, n_inputs - offset); dim3 grid_dim(raft::ceildiv(batch_size, block_dim.x), grid_y, 1); - naive_distance_kernel<<>>( + naive_distance_kernel<<>>( dist.data(), x + offset * dim, y, batch_size, input_len, dim, type); raft::matrix::detail::select_k(handle, @@ -134,7 +134,7 @@ void naive_knn(raft::resources const& handle, indices_topk + offset * k, cuvs::distance::is_min_close(type)); } - RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream.get())); } } // namespace cuvs::neighbors diff --git a/cpp/tests/neighbors/refine.cu b/cpp/tests/neighbors/refine.cu index bf8a2bd98e..8804e6a0e0 100644 --- a/cpp/tests/neighbors/refine.cu +++ b/cpp/tests/neighbors/refine.cu @@ -86,7 +86,7 @@ class RefineTest : public ::testing::TestWithParam> { public: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; RefineHelper data; }; diff --git a/cpp/tests/neighbors/refine_helper.cuh b/cpp/tests/neighbors/refine_helper.cuh index 2d82021d34..a9dd7f3c58 100644 --- a/cpp/tests/neighbors/refine_helper.cuh +++ b/cpp/tests/neighbors/refine_helper.cuh @@ -127,7 +127,7 @@ class RefineHelper { public: RefineInputs p; const raft::resources& handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; raft::device_matrix dataset; raft::device_matrix queries; diff --git a/cpp/tests/neighbors/sparse_brute_force.cu b/cpp/tests/neighbors/sparse_brute_force.cu index 0fd6a26831..d08891a20c 100644 --- a/cpp/tests/neighbors/sparse_brute_force.cu +++ b/cpp/tests/neighbors/sparse_brute_force.cu @@ -86,7 +86,7 @@ class SparseKNNTest : public ::testing::TestWithParam(out_indices.data(), n_rows, k), raft::make_device_matrix_view(out_dists.data(), n_rows, k)); - RAFT_CUDA_TRY(cudaStreamSynchronize(resource::get_cuda_stream(handle))); + RAFT_CUDA_TRY(cudaStreamSynchronize(resource::get_cuda_stream(handle).get())); } void compare() diff --git a/cpp/tests/neighbors/tiered_index.cu b/cpp/tests/neighbors/tiered_index.cu index 4d1fd9c18d..11ea38d6ec 100644 --- a/cpp/tests/neighbors/tiered_index.cu +++ b/cpp/tests/neighbors/tiered_index.cu @@ -220,7 +220,7 @@ class ANNTieredIndexTest : public ::testing::TestWithParam private: raft::resources handle_; - rmm::cuda_stream_view stream_; + cuda::stream_ref stream_; AnnTieredIndexInputs ps; rmm::device_uvector database; rmm::device_uvector queries; diff --git a/cpp/tests/preprocessing/binary_quantization.cu b/cpp/tests/preprocessing/binary_quantization.cu index 5fa6e4b2a4..771312cf15 100644 --- a/cpp/tests/preprocessing/binary_quantization.cu +++ b/cpp/tests/preprocessing/binary_quantization.cu @@ -42,7 +42,7 @@ class BinaryQuantizationTest : public ::testing::TestWithParam>::GetParam()), - stream(raft::resource::get_cuda_stream(handle)), + stream(raft::resource::get_cuda_stream(handle).get()), input_(0, stream) { } diff --git a/cpp/tests/preprocessing/pca.cu b/cpp/tests/preprocessing/pca.cu index 6be6d58820..facb4d235b 100644 --- a/cpp/tests/preprocessing/pca.cu +++ b/cpp/tests/preprocessing/pca.cu @@ -101,7 +101,7 @@ class PcaTest : public ::testing::TestWithParam> { public: PcaTest() : params_(::testing::TestWithParam>::GetParam()), - stream(raft::resource::get_cuda_stream(handle)), + stream(raft::resource::get_cuda_stream(handle).get()), explained_vars(params_.n_col, stream), explained_vars_ref(params_.n_col, stream), components(params_.n_col * params_.n_col, stream), @@ -209,31 +209,31 @@ class PcaTest : public ::testing::TestWithParam> { explained_vars_ref.data(), params_.n_col, cuvs::CompareApprox(params_.tolerance), - s)); + s.get())); ASSERT_TRUE(devArrMatch(components.data(), components_ref.data(), (params_.n_col * params_.n_col), cuvs::CompareApprox(params_.tolerance), - s)); + s.get())); ASSERT_TRUE(devArrMatch(trans_data.data(), trans_data_ref.data(), (params_.n_row * params_.n_col), cuvs::CompareApprox(params_.tolerance), - s)); + s.get())); ASSERT_TRUE(devArrMatch(data.data(), data_back.data(), (params_.n_row * params_.n_col), cuvs::CompareApprox(params_.tolerance), - s)); + s.get())); ASSERT_TRUE(devArrMatch(data2.data(), data2_back.data(), (params_.n_row2 * params_.n_col2), cuvs::CompareApprox(params_.tolerance), - s)); + s.get())); EXPECT_GT(max_recon_err, T(1e-5)) << "Error should be non-zero when n_components < n_cols"; EXPECT_LT(max_recon_err, T(2.0)) << "Reconstruction error should be bounded"; @@ -269,7 +269,7 @@ class PcaRowMajorTest : public ::testing::TestWithParam> { public: PcaRowMajorTest() : params_(::testing::TestWithParam>::GetParam()), - stream(raft::resource::get_cuda_stream(handle)) + stream(raft::resource::get_cuda_stream(handle).get()) { } diff --git a/cpp/tests/preprocessing/product_quantization.cu b/cpp/tests/preprocessing/product_quantization.cu index a392e7e1db..ff44edde20 100644 --- a/cpp/tests/preprocessing/product_quantization.cu +++ b/cpp/tests/preprocessing/product_quantization.cu @@ -93,7 +93,7 @@ class ProductQuantizationTest : public ::testing::TestWithParam>::GetParam()), - stream(raft::resource::get_cuda_stream(handle)), + stream(raft::resource::get_cuda_stream(handle).get()), dataset_(raft::make_device_matrix( handle, params_.n_samples, params_.n_features)), dataset_host_( diff --git a/cpp/tests/preprocessing/scalar_quantization.cu b/cpp/tests/preprocessing/scalar_quantization.cu index 055e80d83a..da23395e05 100644 --- a/cpp/tests/preprocessing/scalar_quantization.cu +++ b/cpp/tests/preprocessing/scalar_quantization.cu @@ -37,7 +37,7 @@ class QuantizationTest : public ::testing::TestWithParam> public: QuantizationTest() : params_(::testing::TestWithParam>::GetParam()), - stream(raft::resource::get_cuda_stream(handle)), + stream(raft::resource::get_cuda_stream(handle).get()), input_(0, stream) { } diff --git a/cpp/tests/preprocessing/spectral_embedding.cu b/cpp/tests/preprocessing/spectral_embedding.cu index 513a91c6ec..3e457ace69 100644 --- a/cpp/tests/preprocessing/spectral_embedding.cu +++ b/cpp/tests/preprocessing/spectral_embedding.cu @@ -41,7 +41,7 @@ class SpectralEmbeddingTest : public ::testing::TestWithParam>::GetParam()), - stream(raft::resource::get_cuda_stream(handle)), + stream(raft::resource::get_cuda_stream(handle).get()), input_(raft::make_device_matrix( handle, params_.n_samples, params_.n_features)), labels_(raft::make_device_vector(handle, params_.n_samples)), diff --git a/cpp/tests/sparse/gram.cu b/cpp/tests/sparse/gram.cu index 91d8c75777..d28711972d 100644 --- a/cpp/tests/sparse/gram.cu +++ b/cpp/tests/sparse/gram.cu @@ -116,7 +116,7 @@ class GramMatrixTest : public ::testing::TestWithParam { protected: GramMatrixTest() : params(GetParam()), - stream(raft::resource::get_cuda_stream(handle)), + stream(raft::resource::get_cuda_stream(handle).get()), x1(0, stream), x2(0, stream), x1_csr_indptr(0, stream), diff --git a/cpp/tests/sparse/neighbors/cross_component_nn.cu b/cpp/tests/sparse/neighbors/cross_component_nn.cu index fdb601807d..3f141f5509 100644 --- a/cpp/tests/sparse/neighbors/cross_component_nn.cu +++ b/cpp/tests/sparse/neighbors/cross_component_nn.cu @@ -59,8 +59,8 @@ class ConnectComponentsTest params = ::testing::TestWithParam>::GetParam(); - raft::sparse::COO out_edges(stream); - raft::sparse::COO out_edges_batched(stream); + raft::sparse::COO out_edges(stream.get()); + raft::sparse::COO out_edges_batched(stream.get()); rmm::device_uvector data(params.n_row * params.n_col, stream); @@ -72,7 +72,7 @@ class ConnectComponentsTest /** * 1. Construct knn graph */ - raft::sparse::COO knn_graph_coo(stream); + raft::sparse::COO knn_graph_coo(stream.get()); auto data_view = raft::make_device_matrix_view( data.data(), params.n_row, params.n_col); @@ -80,7 +80,7 @@ class ConnectComponentsTest handle, data_view, cuvs::distance::DistanceType::L2SqrtExpanded, knn_graph_coo, params.c); raft::sparse::convert::sorted_coo_to_csr( - knn_graph_coo.rows(), knn_graph_coo.nnz, indptr.data(), params.n_row + 1, stream); + knn_graph_coo.rows(), knn_graph_coo.nnz, indptr.data(), params.n_row + 1, stream.get()); /** * 2. Construct MST, sorted by weights @@ -95,7 +95,7 @@ class ConnectComponentsTest params.n_row, knn_graph_coo.nnz, colors.data(), - stream, + stream.get(), false, true); @@ -140,7 +140,7 @@ class ConnectComponentsTest rmm::device_uvector indptr2(params.n_row + 1, stream); raft::sparse::convert::sorted_coo_to_csr( - out_edges.rows(), out_edges.nnz, indptr2.data(), params.n_row + 1, stream); + out_edges.rows(), out_edges.nnz, indptr2.data(), params.n_row + 1, stream.get()); auto output_mst = raft::sparse::solver::mst(handle, indptr2.data(), @@ -149,7 +149,7 @@ class ConnectComponentsTest params.n_row, out_edges.nnz, colors.data(), - stream, + stream.get(), false, false); @@ -394,9 +394,9 @@ class ConnectComponentsEdgesTest ConnectComponentsMutualReachabilityInputs>::GetParam(); raft::sparse::COO out_edges_unbatched( - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); raft::sparse::COO out_edges_batched( - raft::resource::get_cuda_stream(handle)); + raft::resource::get_cuda_stream(handle).get()); rmm::device_uvector data(params.n_row * params.n_col, raft::resource::get_cuda_stream(handle)); diff --git a/cpp/tests/stats/silhouette_score.cu b/cpp/tests/stats/silhouette_score.cu index 15a4c61e81..f8beefdfbb 100644 --- a/cpp/tests/stats/silhouette_score.cu +++ b/cpp/tests/stats/silhouette_score.cu @@ -64,8 +64,9 @@ class silhouetteScoreTest : public ::testing::TestWithParam(neighbors.extent(0), topk); auto distances_host = raft::make_host_matrix(distances.extent(0), topk); - cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources); + cudaStream_t stream = raft::resource::get_cuda_stream(dev_resources).get(); raft::copy(neighbors_host.data_handle(), neighbors.data_handle(), neighbors.size(), stream); raft::copy(distances_host.data_handle(), distances.data_handle(), distances.size(), stream); diff --git a/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md b/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md index eeffd82ae6..7349cf13d7 100644 --- a/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md +++ b/fern/pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md @@ -14,7 +14,7 @@ Asynchronously copies elements between compatible memory locations. ```cpp template void copy(OutputIterator dst, InputIterator src, SizeType n, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); ``` **Parameters** @@ -24,7 +24,7 @@ void copy(OutputIterator dst, InputIterator src, SizeType n, | `dst` | `OutputIterator` | Destination pointer or iterator. | | `src` | `InputIterator` | Source pointer or iterator. | | `n` | `SizeType` | Number of elements to copy. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** @@ -39,7 +39,7 @@ Copies a dense matrix between compatible matrix views. ```cpp template -void copy_matrix(OutputView dst, InputView src, rmm::cuda_stream_view stream); +void copy_matrix(OutputView dst, InputView src, cuda::stream_ref stream); ``` **Parameters** @@ -48,7 +48,7 @@ void copy_matrix(OutputView dst, InputView src, rmm::cuda_stream_view stream); | --- | --- | --- | | `dst` | `OutputView` | Destination matrix view. | | `src` | `InputView` | Source matrix view. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** @@ -64,7 +64,7 @@ Convenience helper for copying host data to device memory. ```cpp template void update_device(DevicePointer dst, HostPointer src, SizeType n, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); ``` **Parameters** @@ -74,7 +74,7 @@ void update_device(DevicePointer dst, HostPointer src, SizeType n, | `dst` | `DevicePointer` | Destination device pointer. | | `src` | `HostPointer` | Source host pointer. | | `n` | `SizeType` | Number of elements to copy. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** @@ -90,7 +90,7 @@ Convenience helper for copying device data to host memory. ```cpp template void update_host(HostPointer dst, DevicePointer src, SizeType n, - rmm::cuda_stream_view stream); + cuda::stream_ref stream); ``` **Parameters** @@ -100,7 +100,7 @@ void update_host(HostPointer dst, DevicePointer src, SizeType n, | `dst` | `HostPointer` | Destination host pointer. | | `src` | `DevicePointer` | Source device pointer. | | `n` | `SizeType` | Number of elements to copy. | -| `stream` | `rmm::cuda_stream_view` | CUDA stream used for the copy. | +| `stream` | `cuda::stream_ref` | CUDA stream used for the copy. | **Returns** diff --git a/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md b/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md index 6f02ddcdbf..ffd7cfe6d1 100644 --- a/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md +++ b/fern/pages/cpp_api/cpp-api-common-types-execution-resources.md @@ -25,7 +25,7 @@ _Source header: `raft/core/resource/cuda_stream.hpp`_ Returns the CUDA stream associated with a resources object. ```cpp -rmm::cuda_stream_view get_cuda_stream(raft::resources const& res); +cuda::stream_ref get_cuda_stream(raft::resources const& res); ``` **Parameters** @@ -36,7 +36,7 @@ rmm::cuda_stream_view get_cuda_stream(raft::resources const& res); **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::resource::sync_stream @@ -47,7 +47,7 @@ Synchronizes the CUDA stream associated with a resources object. ```cpp void sync_stream(raft::resources const& res); -void sync_stream(raft::resources const& res, rmm::cuda_stream_view stream); +void sync_stream(raft::resources const& res, cuda::stream_ref stream); ``` **Parameters** @@ -55,7 +55,7 @@ void sync_stream(raft::resources const& res, rmm::cuda_stream_view stream); | Name | Type | Description | | --- | --- | --- | | `res` | `raft::resources const&` | Resources object to synchronize. | -| `stream` | `rmm::cuda_stream_view` | Optional stream to synchronize instead of the main stream. | +| `stream` | `cuda::stream_ref` | Optional stream to synchronize instead of the main stream. | **Returns** @@ -92,7 +92,7 @@ _Source header: `raft/core/resource/cuda_stream_pool.hpp`_ Returns a stream from the configured stream pool. ```cpp -rmm::cuda_stream_view get_stream_from_stream_pool(raft::resources const& res); +cuda::stream_ref get_stream_from_stream_pool(raft::resources const& res); ``` **Parameters** @@ -103,7 +103,7 @@ rmm::cuda_stream_view get_stream_from_stream_pool(raft::resources const& res); **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::resource::sync_stream_pool @@ -210,7 +210,7 @@ Constructs a single-GPU resources object. ```cpp device_resources( - rmm::cuda_stream_view stream_view = rmm::cuda_stream_per_thread, + cuda::stream_ref stream_view = rmm::cuda_stream_per_thread, std::shared_ptr stream_pool = nullptr, std::shared_ptr workspace_resource = nullptr, std::optional allocation_limit = std::nullopt); @@ -220,7 +220,7 @@ device_resources( | Name | Type | Description | | --- | --- | --- | -| `stream_view` | `rmm::cuda_stream_view` | Default CUDA stream used by algorithms. | +| `stream_view` | `cuda::stream_ref` | Default CUDA stream used by algorithms. | | `stream_pool` | `std::shared_ptr` | Optional CUDA stream pool. | | `workspace_resource` | `std::shared_ptr` | Optional workspace memory resource. | | `allocation_limit` | `std::optional` | Optional temporary workspace allocation limit in bytes. | @@ -232,14 +232,14 @@ Synchronizes either the main stream or a specific CUDA stream. ```cpp void sync_stream() const; -void sync_stream(rmm::cuda_stream_view stream) const; +void sync_stream(cuda::stream_ref stream) const; ``` **Parameters** | Name | Type | Description | | --- | --- | --- | -| `stream` | `rmm::cuda_stream_view` | Stream to synchronize. Omit to synchronize the main stream. | +| `stream` | `cuda::stream_ref` | Stream to synchronize. Omit to synchronize the main stream. | **Returns** @@ -251,12 +251,12 @@ void sync_stream(rmm::cuda_stream_view stream) const; Returns the main CUDA stream associated with the resources object. ```cpp -rmm::cuda_stream_view get_stream() const; +cuda::stream_ref get_stream() const; ``` **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::device_resources::is_stream_pool_initialized @@ -290,8 +290,8 @@ rmm::cuda_stream_pool const& get_stream_pool() const; Returns a stream from the configured CUDA stream pool. ```cpp -rmm::cuda_stream_view get_stream_from_stream_pool() const; -rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const; +cuda::stream_ref get_stream_from_stream_pool() const; +cuda::stream_ref get_stream_from_stream_pool(std::size_t stream_idx) const; ``` **Parameters** @@ -302,7 +302,7 @@ rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const; **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::device_resources::get_next_usable_stream @@ -310,8 +310,8 @@ rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const; Returns a stream from the pool when one exists; otherwise returns the main stream. ```cpp -rmm::cuda_stream_view get_next_usable_stream() const; -rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const; +cuda::stream_ref get_next_usable_stream() const; +cuda::stream_ref get_next_usable_stream(std::size_t stream_idx) const; ``` **Parameters** @@ -322,7 +322,7 @@ rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const; **Returns** -`rmm::cuda_stream_view` +`cuda::stream_ref` #### raft::device_resources::sync_stream_pool diff --git a/fern/pages/other/resources.md b/fern/pages/other/resources.md index 9cc45196b5..883d3e1265 100644 --- a/fern/pages/other/resources.md +++ b/fern/pages/other/resources.md @@ -369,7 +369,7 @@ cudaStreamDestroy(stream); cudaStream_t stream; cudaStreamCreate(&stream); -raft::device_resources resources{rmm::cuda_stream_view{stream}}; +raft::device_resources resources{cuda::stream_ref{stream}}; // cuVS C++ calls using resources are enqueued on stream. diff --git a/fern/scripts/generate_api_reference.py b/fern/scripts/generate_api_reference.py index dd7f23442e..35a54da44e 100755 --- a/fern/scripts/generate_api_reference.py +++ b/fern/scripts/generate_api_reference.py @@ -800,9 +800,9 @@ def add_symbol( "raft-resource-get-cuda-stream", "raft::resource::get_cuda_stream", "Returns the CUDA stream associated with a resources object.", - "rmm::cuda_stream_view get_cuda_stream(raft::resources const& res);", + "cuda::stream_ref get_cuda_stream(raft::resources const& res);", [("res", "raft::resources const&", "Resources object to query.")], - "rmm::cuda_stream_view", + "cuda::stream_ref", nested=True, ) add_symbol( @@ -812,7 +812,7 @@ def add_symbol( "Synchronizes the CUDA stream associated with a resources object.", ( "void sync_stream(raft::resources const& res);\n" - "void sync_stream(raft::resources const& res, rmm::cuda_stream_view stream);" + "void sync_stream(raft::resources const& res, cuda::stream_ref stream);" ), [ ( @@ -822,7 +822,7 @@ def add_symbol( ), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "Optional stream to synchronize instead of the main stream.", ), ], @@ -861,9 +861,9 @@ def add_symbol( "raft-resource-get-stream-from-stream-pool", "raft::resource::get_stream_from_stream_pool", "Returns a stream from the configured stream pool.", - "rmm::cuda_stream_view get_stream_from_stream_pool(raft::resources const& res);", + "cuda::stream_ref get_stream_from_stream_pool(raft::resources const& res);", [("res", "raft::resources const&", "Resources object to query.")], - "rmm::cuda_stream_view", + "cuda::stream_ref", nested=True, ) add_symbol( @@ -947,7 +947,7 @@ def add_symbol( "Constructs a single-GPU resources object.", ( "device_resources(\n" - " rmm::cuda_stream_view stream_view = rmm::cuda_stream_per_thread,\n" + " cuda::stream_ref stream_view = rmm::cuda_stream_per_thread,\n" " std::shared_ptr stream_pool = nullptr,\n" " std::shared_ptr workspace_resource = nullptr,\n" " std::optional allocation_limit = std::nullopt);" @@ -955,7 +955,7 @@ def add_symbol( [ ( "stream_view", - "rmm::cuda_stream_view", + "cuda::stream_ref", "Default CUDA stream used by algorithms.", ), ( @@ -982,12 +982,12 @@ def add_symbol( "Synchronizes either the main stream or a specific CUDA stream.", ( "void sync_stream() const;\n" - "void sync_stream(rmm::cuda_stream_view stream) const;" + "void sync_stream(cuda::stream_ref stream) const;" ), [ ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "Stream to synchronize. Omit to synchronize the main stream.", ) ], @@ -998,8 +998,8 @@ def add_symbol( "raft-device-resources-get-stream", "raft::device_resources::get_stream", "Returns the main CUDA stream associated with the resources object.", - "rmm::cuda_stream_view get_stream() const;", - returns="rmm::cuda_stream_view", + "cuda::stream_ref get_stream() const;", + returns="cuda::stream_ref", ) add_symbol( lines, @@ -1023,8 +1023,8 @@ def add_symbol( "raft::device_resources::get_stream_from_stream_pool", "Returns a stream from the configured CUDA stream pool.", ( - "rmm::cuda_stream_view get_stream_from_stream_pool() const;\n" - "rmm::cuda_stream_view get_stream_from_stream_pool(std::size_t stream_idx) const;" + "cuda::stream_ref get_stream_from_stream_pool() const;\n" + "cuda::stream_ref get_stream_from_stream_pool(std::size_t stream_idx) const;" ), [ ( @@ -1033,7 +1033,7 @@ def add_symbol( "Optional index of the stream in the stream pool.", ) ], - "rmm::cuda_stream_view", + "cuda::stream_ref", ) add_symbol( lines, @@ -1044,8 +1044,8 @@ def add_symbol( "the main stream." ), ( - "rmm::cuda_stream_view get_next_usable_stream() const;\n" - "rmm::cuda_stream_view get_next_usable_stream(std::size_t stream_idx) const;" + "cuda::stream_ref get_next_usable_stream() const;\n" + "cuda::stream_ref get_next_usable_stream(std::size_t stream_idx) const;" ), [ ( @@ -1054,7 +1054,7 @@ def add_symbol( "Optional stream pool index to use when a stream pool is configured.", ) ], - "rmm::cuda_stream_view", + "cuda::stream_ref", ) add_symbol( lines, @@ -2085,7 +2085,7 @@ def add_symbol( ( "template \n" "void copy(OutputIterator dst, InputIterator src, SizeType n,\n" - " rmm::cuda_stream_view stream);" + " cuda::stream_ref stream);" ), [ ("dst", "OutputIterator", "Destination pointer or iterator."), @@ -2093,7 +2093,7 @@ def add_symbol( ("n", "SizeType", "Number of elements to copy."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ], @@ -2104,13 +2104,13 @@ def add_symbol( "raft-copy-matrix", "raft::copy_matrix", "Copies a dense matrix between compatible matrix views.", - "template \nvoid copy_matrix(OutputView dst, InputView src, rmm::cuda_stream_view stream);", + "template \nvoid copy_matrix(OutputView dst, InputView src, cuda::stream_ref stream);", [ ("dst", "OutputView", "Destination matrix view."), ("src", "InputView", "Source matrix view."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ], @@ -2124,7 +2124,7 @@ def add_symbol( ( "template \n" "void update_device(DevicePointer dst, HostPointer src, SizeType n,\n" - " rmm::cuda_stream_view stream);" + " cuda::stream_ref stream);" ), [ ("dst", "DevicePointer", "Destination device pointer."), @@ -2132,7 +2132,7 @@ def add_symbol( ("n", "SizeType", "Number of elements to copy."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ], @@ -2146,7 +2146,7 @@ def add_symbol( ( "template \n" "void update_host(HostPointer dst, DevicePointer src, SizeType n,\n" - " rmm::cuda_stream_view stream);" + " cuda::stream_ref stream);" ), [ ("dst", "HostPointer", "Destination host pointer."), @@ -2154,7 +2154,7 @@ def add_symbol( ("n", "SizeType", "Number of elements to copy."), ( "stream", - "rmm::cuda_stream_view", + "cuda::stream_ref", "CUDA stream used for the copy.", ), ],