Skip to content

Adopt CUDA stream compatibility accessors #2557

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bdice wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Adopt CUDA stream compatibility accessors #2557

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8cc00f8
Use stream compatibility accessors
bdice Sep 4, 2026
e278c97
Complete stream compatibility accessor preparation
bdice Sep 5, 2026
62aaa35
Extract native handles in CUDA checks
bdice Sep 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion c/src/core/c_api.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<raft::resources*>(res);
*stream = raft::resource::get_cuda_stream(*res_ptr);
*stream = raft::resource::get_cuda_stream(*res_ptr).get();
});
}

Expand Down
2 changes: 1 addition & 1 deletion c/src/neighbors/nn_descent.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion c/tests/neighbors/ann_ivf_sq_c.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion cpp/include/cuvs/neighbors/common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
4 changes: 2 additions & 2 deletions cpp/src/cluster/detail/agglomerative.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -280,7 +280,7 @@ void extract_flattened_clusters(raft::resources const& handle,
rmm::device_uvector<value_idx> levels(n_vertices, stream);

value_idx n_blocks = raft::ceildiv(n_vertices, (value_idx)tpb);
write_levels_kernel<<<n_blocks, tpb, 0, stream>>>(children, levels.data(), n_vertices);
write_levels_kernel<<<n_blocks, tpb, 0, stream.get()>>>(children, levels.data(), n_vertices);
/**
* Step 1: Find label roots:
*
Expand Down Expand Up @@ -323,7 +323,7 @@ void extract_flattened_clusters(raft::resources const& handle,
*/
value_idx cut_level = (n_edges / 2) - (n_clusters - 1);

inherit_labels<<<n_blocks, tpb, 0, stream>>>(
inherit_labels<<<n_blocks, tpb, 0, stream.get()>>>(
children, levels.data(), n_leaves, tmp_labels.data(), cut_level, n_vertices);

// copy tmp labels to actual labels
Expand Down
6 changes: 3 additions & 3 deletions cpp/src/cluster/detail/connectivities.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ struct distance_graph_impl<Linkage::KNN_GRAPH, value_idx, value_t> {
auto thrust_policy = raft::resource::get_thrust_policy(handle);

// Need to symmetrize knn into undirected graph
raft::sparse::COO<value_t, value_idx> knn_graph_coo(stream);
raft::sparse::COO<value_t, value_idx> knn_graph_coo(stream.get());

auto X_view = raft::make_device_matrix_view<const value_t, value_idx, raft::row_major>(X, m, n);
cuvs::neighbors::detail::knn_graph<value_idx, value_t, size_t>(
Expand Down Expand Up @@ -92,7 +92,7 @@ struct distance_graph_impl<Linkage::KNN_GRAPH, value_idx, value_t> {
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
Expand Down Expand Up @@ -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<value_idx><<<blocks, 256, 0, stream>>>(indices, m, nnz);
fill_indices2<value_idx><<<blocks, 256, 0, stream.get()>>>(indices, m, nnz);

raft::linalg::map_offset(handle,
raft::make_device_vector_view<value_idx, value_idx>(indptr, m),
Expand Down
10 changes: 5 additions & 5 deletions cpp/src/cluster/detail/kmeans.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ void kmeansPlusPlus(raft::resources const& handle,
rmm::device_uvector<char>& workspace)
{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> 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;
Expand Down Expand Up @@ -309,7 +309,7 @@ void initScalableKMeansPlusPlus(raft::resources const& handle,
{
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> 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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -1035,7 +1035,7 @@ void kmeans_predict(raft::resources const& handle,
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> 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,
Expand Down Expand Up @@ -1186,7 +1186,7 @@ void kmeans_transform(raft::resources const& handle,
"kmeans only supports L2Expanded or L2SqrtExpanded distance metrics.");
raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> 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;
Expand Down
12 changes: 6 additions & 6 deletions cpp/src/cluster/detail/kmeans_balanced.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ inline std::enable_if_t<std::is_floating_point_v<MathT>> predict_core(
&beta,
distances.data(),
n_clusters,
stream);
stream.get());

auto distances_const_view = raft::make_device_matrix_view<const MathT, IdxT, raft::row_major>(
distances.data(), n_rows, n_clusters);
Expand Down Expand Up @@ -287,12 +287,12 @@ 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<T, MathT>) {
raft::linalg::reduce_rows_by_key(
dataset, dim, labels, nullptr, n_rows, dim, n_clusters, centers, stream, reset_counters);
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<MappingOpT, const T*> 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);
mapping_itr, dim, labels, nullptr, n_rows, dim, n_clusters, centers, stream.get(), reset_counters);
}

// Compute weight of each cluster
Expand Down Expand Up @@ -689,7 +689,7 @@ 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<IdxT>(kBlockDimY)), 1, 1);
adjust_centers_random_donor_kernel<kBlockDimY>
<<<grid_dim, block_dim, 0, stream>>>(centers,
<<<grid_dim, block_dim, 0, stream.get()>>>(centers,
n_clusters,
dim,
dataset,
Expand All @@ -709,7 +709,7 @@ auto adjust_centers(const raft::resources& handle,
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<IdxT>(kBlockDimY)), 1, 1);
adjust_centers_kernel<kBlockDimY><<<grid_dim, block_dim, 0, stream>>>(centers,
adjust_centers_kernel<kBlockDimY><<<grid_dim, block_dim, 0, stream.get()>>>(centers,
n_pairs,
dim,
dataset,
Expand Down Expand Up @@ -1068,7 +1068,7 @@ auto build_fine_clusters(const raft::resources& handle,
}

thrust::transform_iterator<MappingOpT, const T*> 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) {
Expand Down
20 changes: 10 additions & 10 deletions cpp/src/cluster/detail/kmeans_common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ void countLabels(raft::resources const& handle,
IndexT n_clusters,
rmm::device_uvector<char>& 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<IndexT> CubIndexT;
Expand Down Expand Up @@ -177,7 +177,7 @@ void weightSum(

if constexpr (raft::is_device_mdspan_v<decltype(weight)>) {
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);
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -254,7 +254,7 @@ void sampleCentroids(raft::resources const& handle,
rmm::device_uvector<DataT>& inRankCp,
rmm::device_uvector<char>& 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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -455,7 +455,7 @@ void countSamplesInCluster(raft::resources const& handle,
rmm::device_uvector<char>& workspace,
raft::device_vector_view<DataT, IndexT> 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);
Expand Down Expand Up @@ -534,7 +534,7 @@ void compute_centroid_adjustments(
rmm::device_uvector<char>& 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);
Expand Down Expand Up @@ -582,7 +582,7 @@ void finalize_centroids(raft::resources const& handle,
raft::device_matrix_view<const DataT, IndexT> old_centroids,
raft::device_matrix_view<DataT, IndexT> 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<raft::Apply::ALONG_COLUMNS>(handle,
raft::make_const_mdspan(centroid_sums),
Expand Down Expand Up @@ -617,7 +617,7 @@ void compute_centroid_shift(raft::resources const& handle,
raft::device_matrix_view<const DataT, IndexT> new_centroids,
raft::device_scalar_view<DataT> 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{},
Expand Down Expand Up @@ -702,7 +702,7 @@ void process_batch(
raft::device_scalar_view<DataT> clustering_cost,
rmm::device_uvector<char>& batch_workspace)
{
cudaStream_t stream = raft::resource::get_cuda_stream(handle);
cudaStream_t stream = raft::resource::get_cuda_stream(handle).get();

minClusterAndDistanceCompute<DataT, IndexT>(handle,
batch_data,
Expand Down
4 changes: 2 additions & 2 deletions cpp/src/cluster/detail/minClusterDistanceCompute.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ void minClusterAndDistanceCompute(
int batch_centroids,
rmm::device_uvector<char>& 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);
Expand Down Expand Up @@ -242,7 +242,7 @@ void minClusterDistanceCompute(raft::resources const& handle,
int batch_centroids,
rmm::device_uvector<char>& 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);
Expand Down
32 changes: 16 additions & 16 deletions cpp/src/cluster/detail/mst.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ void connect_knn_graph(
{
auto stream = raft::resource::get_cuda_stream(handle);

raft::sparse::COO<value_t, value_idx> connected_edges(stream);
raft::sparse::COO<value_t, value_idx> connected_edges(stream.get());

// default row and column batch sizes are chosen for computing cross component nearest neighbors.
// Reference: PR #1445
Expand All @@ -95,7 +95,7 @@ void connect_knn_graph(

rmm::device_uvector<value_idx> 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
Expand All @@ -107,11 +107,11 @@ void connect_knn_graph(
m,
connected_edges.nnz,
color,
stream,
stream.get(),
false,
false);

merge_msts<value_idx, value_t>(msf, new_mst, stream);
merge_msts<value_idx, value_t>(msf, new_mst, stream.get());
}

/**
Expand Down Expand Up @@ -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<value_idx> 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<value_idx> h_color(m);
raft::copy(handle,
Expand Down Expand Up @@ -251,7 +251,7 @@ void connect_knn_graph(

rmm::device_uvector<value_idx> 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
Expand All @@ -263,11 +263,11 @@ void connect_knn_graph(
m,
new_nnz,
color,
stream,
stream.get(),
false,
false);

merge_msts<value_idx, value_t>(msf, new_mst, stream);
merge_msts<value_idx, value_t>(msf, new_mst, stream.get());
}

/**
Expand Down Expand Up @@ -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<value_idx, value_idx, value_t, double>(
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;

Expand Down Expand Up @@ -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());
}

/**
Expand All @@ -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
Loading
Loading