Skip to content

Adopt CUDA stream compatibility accessors #3136

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 8 commits into
base: main
Choose a base branch
from
Open

Adopt CUDA stream compatibility accessors #3136

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
5fbe443
Use stream compatibility accessors
bdice Sep 4, 2026
60231c7
Complete stream compatibility accessor preparation
bdice Sep 5, 2026
c48c50c
Cover new raw stream boundaries
bdice Sep 5, 2026
662f9f5
Cover new solver raw stream boundaries
bdice Sep 5, 2026
d730bd5
Add missing benchmark stream accessors
bdice Sep 5, 2026
f356e94
Add benchmark raw stream accessors
bdice Sep 5, 2026
51a2375
Extract native handles at CUDA stream boundaries
bdice Sep 5, 2026
e40580b
Use stream compatibility synchronization in benchmark
bdice Sep 5, 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
6 changes: 3 additions & 3 deletions cpp/bench/prims/common/benchmark.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ struct cuda_event_timer {
RAFT_CUDA_TRY(cudaEventCreate(&start_));
RAFT_CUDA_TRY(cudaEventCreate(&stop_));
raft::interruptible::synchronize(stream_);
RAFT_CUDA_TRY(cudaEventRecord(start_, stream_));
RAFT_CUDA_TRY(cudaEventRecord(start_, stream_.get()));
}
cuda_event_timer() = delete;

Expand All @@ -85,7 +85,7 @@ struct cuda_event_timer {
*/
~cuda_event_timer()
{
RAFT_CUDA_TRY_NO_THROW(cudaEventRecord(stop_, stream_));
RAFT_CUDA_TRY_NO_THROW(cudaEventRecord(stop_, stream_.get()));
raft::interruptible::synchronize(stop_);
float milliseconds = 0.0f;
RAFT_CUDA_TRY_NO_THROW(cudaEventElapsedTime(&milliseconds, start_, stop_));
Expand Down Expand Up @@ -137,7 +137,7 @@ class fixture {
/** The helper that writes zeroes to some buffer in GPU memory to flush the L2 cache. */
void flush_L2_cache()
{
RAFT_CUDA_TRY(cudaMemsetAsync(scratch_buf_.data(), 0, scratch_buf_.size(), stream));
RAFT_CUDA_TRY(cudaMemsetAsync(scratch_buf_.data(), 0, scratch_buf_.size(), stream.get()));
}

/**
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/linalg/add.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ struct add : public fixture {
void run_benchmark(::benchmark::State& state) override
{
loop_on_state(state, [this]() {
raft::linalg::add(ptr0.data(), ptr0.data(), ptr1.data(), params.len, stream);
raft::linalg::add(ptr0.data(), ptr0.data(), ptr1.data(), params.len, stream.get());
});
}

Expand Down
3 changes: 2 additions & 1 deletion cpp/bench/prims/linalg/map_then_reduce.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ struct map_then_reduce : public fixture {
void run_benchmark(::benchmark::State& state) override
{
loop_on_state(state, [this]() {
raft::linalg::mapThenSumReduce(out.data(), params.len, Identity<T>(), stream, in.data());
raft::linalg::mapThenSumReduce(
out.data(), params.len, Identity<T>(), stream.get(), in.data());
});
}

Expand Down
16 changes: 8 additions & 8 deletions cpp/bench/prims/linalg/matrix_vector_op.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ struct mat_vec_op : public fixture {
params.cols,
params.rows,
OpT{},
stream);
stream.get());
} else {
raft::linalg::matrixVectorOp<true, false>(out.data() + params.outAlignOffset,
in.data() + params.inAlignOffset,
Expand All @@ -63,7 +63,7 @@ struct mat_vec_op : public fixture {
params.cols,
params.rows,
OpT{},
stream);
stream.get());
}
} else {
if (params.bcastAlongRows) {
Expand All @@ -74,7 +74,7 @@ struct mat_vec_op : public fixture {
params.cols,
params.rows,
OpT{},
stream);
stream.get());
} else {
raft::linalg::matrixVectorOp<false, false>(out.data() + params.outAlignOffset,
in.data() + params.inAlignOffset,
Expand All @@ -83,7 +83,7 @@ struct mat_vec_op : public fixture {
params.cols,
params.rows,
OpT{},
stream);
stream.get());
}
}
} else {
Expand All @@ -95,15 +95,15 @@ struct mat_vec_op : public fixture {
params.cols,
params.rows,
OpT{},
stream);
stream.get());
} else {
raft::linalg::matrixVectorOp<true, false>(out.data() + params.outAlignOffset,
in.data() + params.inAlignOffset,
vec1.data(),
params.cols,
params.rows,
OpT{},
stream);
stream.get());
}
} else {
if (params.bcastAlongRows) {
Expand All @@ -113,15 +113,15 @@ struct mat_vec_op : public fixture {
params.cols,
params.rows,
OpT{},
stream);
stream.get());
} else {
raft::linalg::matrixVectorOp<false, false>(out.data() + params.outAlignOffset,
in.data() + params.inAlignOffset,
vec1.data(),
params.cols,
params.rows,
OpT{},
stream);
stream.get());
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions cpp/bench/prims/linalg/reduce.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@ struct reduce : public fixture {
loop_on_state(state, [this]() {
if (along_rows) {
raft::linalg::reduce<true, true>(
out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream);
out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream.get());
} else {
raft::linalg::reduce<true, false>(
out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream);
out.data(), in.data(), input_size.cols, input_size.rows, T(0.f), stream.get());
}
});
}
Expand Down
10 changes: 8 additions & 2 deletions cpp/bench/prims/linalg/reduce_cols_by_key.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,14 @@ struct reduce_cols_by_key : public fixture {
state.SetLabel(label_stream.str());

loop_on_state(state, [this]() {
raft::linalg::reduce_cols_by_key(
in.data(), keys.data(), out.data(), params.rows, params.cols, params.keys, stream, false);
raft::linalg::reduce_cols_by_key(in.data(),
keys.data(),
out.data(),
params.rows,
params.cols,
params.keys,
stream.get(),
false);
});
}

Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/linalg/reduce_rows_by_key.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ struct reduce_rows_by_key : public fixture {
params.cols,
params.keys,
out.data(),
stream,
stream.get(),
false);
});
}
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/linalg/transpose.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ struct TransposeBench : public fixture {
output_view.data_handle(),
params.rows,
params.cols,
handle.get_stream());
handle.get_stream().get());
});
}

Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/matrix/select_k.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ struct selection : public fixture {
out_dists_(p.batch_size * p.k, stream),
out_ids_(p.batch_size * p.k, stream)
{
raft::sparse::iota_fill(in_ids_.data(), IdxT(p.batch_size), IdxT(p.len), stream);
raft::sparse::iota_fill(in_ids_.data(), IdxT(p.batch_size), IdxT(p.len), stream.get());
raft::random::RngState state{42};

KeyT min_value = -1.0;
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/random/make_blobs.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ struct make_blobs : public fixture {
params.rows,
params.cols,
params.clusters,
this->stream,
this->stream.get(),
params.row_major);
});
}
Expand Down
4 changes: 2 additions & 2 deletions cpp/bench/prims/random/permute.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ struct permute : public fixture {
params.cols,
params.rows,
params.rowMajor,
stream,
stream.get(),
123456ULL);
});
}
Expand Down Expand Up @@ -97,7 +97,7 @@ struct permute_perms_only : public fixture {
IntType(0),
IntType(n_rows),
true,
stream,
stream.get(),
123456ULL);
bytes_processed += size_t(n_rows) * sizeof(IntType);
});
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/sparse/convert_csr.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ struct bench_base : public fixture {
// col_ind is over-dimensioned because nnz is unknown at this point
col_ind(p.num_rows * p.num_cols, stream)
{
init_adj(adj.data(), p.num_rows, p.num_cols, p.divisor, stream);
init_adj(adj.data(), p.num_rows, p.num_cols, p.divisor, stream.get());

std::vector<index_t> row_ind_host(p.num_rows);
for (size_t i = 0; i < row_ind_host.size(); ++i) {
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/prims/sparse/select_k_csr.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ struct SelectKCsrTest : public fixture {
1,
nnz,
1,
stream,
stream.get(),
false,
nullptr,
nullptr,
Expand Down
4 changes: 2 additions & 2 deletions cpp/bench/prims/sparse/svds.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,7 @@ class svds_bench_base : public fixture {
CUBLAS_OP_N,
one,
zero,
stream);
stream.get());

raft::linalg::gemm(handle,
Vt.data_handle(),
Expand All @@ -268,7 +268,7 @@ class svds_bench_base : public fixture {
CUBLAS_OP_T,
one,
zero,
stream);
stream.get());

std::vector<value_t> h_utu(static_cast<std::size_t>(params.k) * params.k);
std::vector<value_t> h_vvt(static_cast<std::size_t>(params.k) * params.k);
Expand Down
8 changes: 4 additions & 4 deletions cpp/bench/prims/util/fast_int_div.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,20 +75,20 @@ struct fast_int_div_bench : public fixture {
h_numerators.data(),
h_numerators.size() * sizeof(IntT),
cudaMemcpyHostToDevice,
stream));
stream.get()));
RAFT_CUDA_TRY(cudaMemcpyAsync(d_divisors.data(),
h_divisors.data(),
h_divisors.size() * sizeof(divisor_t),
cudaMemcpyHostToDevice,
stream));
stream.synchronize();
stream.get()));
stream.sync();
}

void run_benchmark(::benchmark::State& state) override
{
const auto* divisors = static_cast<const divisor_t*>(d_divisors.data());
loop_on_state(state, [this, divisors]() {
divmod_kernel<IntT, divisor_t><<<kBlocks, kThreads, 0, stream>>>(
divmod_kernel<IntT, divisor_t><<<kBlocks, kThreads, 0, stream.get()>>>(
d_numerators.data(), kNumNumerators, divisors, kNumDivisors, out_d.data());
RAFT_CUDA_TRY(cudaPeekAtLastError());
});
Expand Down
4 changes: 2 additions & 2 deletions cpp/include/raft/comms/detail/mpi_comms.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ class mpi_comms : public comms_iface {
rank_(1),
status_(stream),
next_request_id_(0),
stream_(stream)
stream_(stream.get())
{
int mpi_is_initialized = 0;
RAFT_MPI_TRY(MPI_Initialized(&mpi_is_initialized));
Expand Down Expand Up @@ -134,7 +134,7 @@ class mpi_comms : public comms_iface {
rank_(1),
status_(stream),
next_request_id_(0),
stream_(stream)
stream_(stream.get())
{
int mpi_is_initialized = 0;
RAFT_MPI_TRY(MPI_Initialized(&mpi_is_initialized));
Expand Down
4 changes: 2 additions & 2 deletions cpp/include/raft/comms/detail/std_comms.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ class std_comms : public comms_iface {
rmm::cuda_stream_view stream,
bool subcomms_ucp = true)
: nccl_comm_(nccl_comm),
stream_(stream),
stream_(stream.get()),
status_(stream),
num_ranks_(num_ranks),
rank_(rank),
Expand All @@ -97,7 +97,7 @@ class std_comms : public comms_iface {
rmm::cuda_stream_view stream,
bool own_nccl_comm = false)
: nccl_comm_(nccl_comm),
stream_(stream),
stream_(stream.get()),
status_(stream),
num_ranks_(num_ranks),
rank_(rank),
Expand Down
20 changes: 10 additions & 10 deletions cpp/include/raft/comms/detail/test.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ bool test_collective_allreduce(raft::resources const& handle, int root)

int const send = 1;

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_scalar<int> temp_d(stream);
RAFT_CUDA_TRY(cudaMemcpyAsync(temp_d.data(), &send, 1, cudaMemcpyHostToDevice, stream));
Expand Down Expand Up @@ -65,7 +65,7 @@ bool test_collective_broadcast(raft::resources const& handle, int root)

int const send = root;

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_scalar<int> temp_d(stream);

Expand Down Expand Up @@ -100,7 +100,7 @@ bool test_collective_reduce(raft::resources const& handle, int root)

int const send = root;

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_scalar<int> temp_d(stream);

Expand Down Expand Up @@ -136,7 +136,7 @@ bool test_collective_allgather(raft::resources const& handle, int root)

int const send = communicator.get_rank();

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_scalar<int> temp_d(stream);
rmm::device_uvector<int> recv_d(communicator.get_size(), stream);
Expand Down Expand Up @@ -173,7 +173,7 @@ bool test_collective_gather(raft::resources const& handle, int root)

int const send = communicator.get_rank();

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_scalar<int> temp_d(stream);
rmm::device_uvector<int> recv_d(communicator.get_rank() == root ? communicator.get_size() : 0,
Expand Down Expand Up @@ -217,7 +217,7 @@ bool test_collective_gatherv(raft::resources const& handle, int root)
displacements[communicator.get_rank() + 1] - displacements[communicator.get_rank()],
communicator.get_rank());

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_uvector<int> temp_d(sends.size(), stream);
rmm::device_uvector<int> recv_d(communicator.get_rank() == root ? displacements.back() : 0,
Expand Down Expand Up @@ -269,7 +269,7 @@ bool test_collective_reducescatter(raft::resources const& handle, int root)

std::vector<int> sends(communicator.get_size(), 1);

cudaStream_t stream = resource::get_cuda_stream(handle);
cudaStream_t stream = resource::get_cuda_stream(handle).get();

rmm::device_uvector<int> temp_d(sends.size(), stream);
rmm::device_scalar<int> recv_d(stream);
Expand Down Expand Up @@ -367,7 +367,7 @@ bool test_pointToPoint_device_send_or_recv(raft::resources const& h, int numTria
{
comms_t const& communicator = resource::get_comms(h);
int const rank = communicator.get_rank();
cudaStream_t stream = resource::get_cuda_stream(h);
cudaStream_t stream = resource::get_cuda_stream(h).get();

bool ret = true;
for (int i = 0; i < numTrials; i++) {
Expand Down Expand Up @@ -410,7 +410,7 @@ bool test_pointToPoint_device_sendrecv(raft::resources const& h, int numTrials)
{
comms_t const& communicator = resource::get_comms(h);
int const rank = communicator.get_rank();
cudaStream_t stream = resource::get_cuda_stream(h);
cudaStream_t stream = resource::get_cuda_stream(h).get();

bool ret = true;
for (int i = 0; i < numTrials; i++) {
Expand Down Expand Up @@ -457,7 +457,7 @@ bool test_pointToPoint_device_multicast_sendrecv(raft::resources const& h, int n
{
comms_t const& communicator = resource::get_comms(h);
int const rank = communicator.get_rank();
cudaStream_t stream = resource::get_cuda_stream(h);
cudaStream_t stream = resource::get_cuda_stream(h).get();

bool ret = true;
for (int i = 0; i < numTrials; i++) {
Expand Down
Loading
Loading