diff --git a/cpp/bench/sg/benchmark.cuh b/cpp/bench/sg/benchmark.cuh index 44237050ea..b9ebe95cdf 100644 --- a/cpp/bench/sg/benchmark.cuh +++ b/cpp/bench/sg/benchmark.cuh @@ -16,6 +16,7 @@ #include +#include #include #include @@ -33,7 +34,7 @@ class Fixture : public MLCommon::Bench::Fixture { { if (stream == 0) { RAFT_CUDA_TRY(cudaStreamCreate(&stream)); } auto stream_pool = std::make_shared(numStreams()); - handle.reset(new raft::handle_t{rmm::cuda_stream_view{stream}, stream_pool}); + handle.reset(new raft::handle_t{cuda::stream_ref{stream}, stream_pool}); MLCommon::Bench::Fixture::SetUp(state); } diff --git a/cpp/src/glm/qn/simple_mat/dense.hpp b/cpp/src/glm/qn/simple_mat/dense.hpp index a257b7290a..a55877f114 100644 --- a/cpp/src/glm/qn/simple_mat/dense.hpp +++ b/cpp/src/glm/qn/simple_mat/dense.hpp @@ -327,7 +327,7 @@ std::ostream& operator<<(std::ostream& os, const SimpleVec& v) { std::vector out(v.len); raft::update_host(&out[0], v.data, v.len, 0); - raft::interruptible::synchronize(rmm::cuda_stream_view()); + raft::interruptible::synchronize(cuda::stream_ref()); int it = 0; for (; it < v.len - 1;) { os << out[it] << " "; @@ -343,7 +343,7 @@ std::ostream& operator<<(std::ostream& os, const SimpleDenseMat& mat) os << "ord=" << (mat.ord == COL_MAJOR ? "CM" : "RM") << "\n"; std::vector out(mat.len); raft::update_host(&out[0], mat.data, mat.len, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}); - raft::interruptible::synchronize(rmm::cuda_stream_view()); + raft::interruptible::synchronize(cuda::stream_ref()); if (mat.ord == COL_MAJOR) { for (int r = 0; r < mat.m; r++) { int idx = r; diff --git a/cpp/src/glm/qn/simple_mat/sparse.hpp b/cpp/src/glm/qn/simple_mat/sparse.hpp index d5f8ef79b6..ce3372dae3 100644 --- a/cpp/src/glm/qn/simple_mat/sparse.hpp +++ b/cpp/src/glm/qn/simple_mat/sparse.hpp @@ -183,7 +183,7 @@ std::ostream& operator<<(std::ostream& os, const SimpleSparseMat& mat) raft::update_host(&cols[0], mat.cols, mat.nnz, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}); raft::update_host( &row_ids[0], mat.row_ids, mat.m + 1, cuda::stream_ref{cudaStream_t{cudaStreamDefault}}); - raft::interruptible::synchronize(rmm::cuda_stream_view()); + raft::interruptible::synchronize(cuda::stream_ref()); int i, row_end = 0; for (int row = 0; row < mat.m; row++) { diff --git a/cpp/src/umap/simpl_set_embed/algo.cuh b/cpp/src/umap/simpl_set_embed/algo.cuh index 127b621833..856d0be680 100644 --- a/cpp/src/umap/simpl_set_embed/algo.cuh +++ b/cpp/src/umap/simpl_set_embed/algo.cuh @@ -25,6 +25,7 @@ #include #include +#include #include #include #include @@ -127,7 +128,7 @@ T create_rounding_factor(T max_abs, int n) template T create_gradient_rounding_factor( - const int* head, nnz_t nnz, int n_samples, T alpha, rmm::cuda_stream_view stream) + const int* head, nnz_t nnz, int n_samples, T alpha, cuda::stream_ref stream) { rmm::device_uvector buffer(n_samples, stream); // calculate the maximum number of edges connected to 1 vertex. @@ -212,7 +213,7 @@ void optimize_layout(T* head_embedding, bool move_other = head_embedding == tail_embedding; T alpha = params->initial_alpha; - auto stream_view = rmm::cuda_stream_view(stream); + auto stream_view = cuda::stream_ref(stream); T rounding = create_gradient_rounding_factor(head, nnz, head_n, alpha, stream_view); diff --git a/cpp/src/umap/simpl_set_embed/optimize_batch_kernel.cuh b/cpp/src/umap/simpl_set_embed/optimize_batch_kernel.cuh index eac94c2660..ee58a0b9ae 100644 --- a/cpp/src/umap/simpl_set_embed/optimize_batch_kernel.cuh +++ b/cpp/src/umap/simpl_set_embed/optimize_batch_kernel.cuh @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -1045,8 +1046,6 @@ void call_optimize_batch_kernel(T* head_embedding, bool use_shared_mem = requiredSize < static_cast(raft::getSharedMemPerBlock()); T nsr_inv = T(1.0) / params->negative_sample_rate; - auto stream_view = rmm::cuda_stream_view(stream); - auto launch_kernel = [&](size_t offset = 0) { if (params->n_components == 2) { // multicore implementation with registers diff --git a/python/cuml/cuml/manifold/umap/umap.pyx b/python/cuml/cuml/manifold/umap/umap.pyx index fe6b611af4..9c4100331d 100644 --- a/python/cuml/cuml/manifold/umap/umap.pyx +++ b/python/cuml/cuml/manifold/umap/umap.pyx @@ -312,9 +312,10 @@ cdef class RaftCOO: cdef RaftCOO self = RaftCOO.__new__(RaftCOO) cdef handle_t* handle_ = handle.getHandle() - cdef lib.COO* coo = new lib.COO(handle_.get_stream()) + cdef cudaStream_t stream = handle_.get_stream() + cdef lib.COO* coo = new lib.COO(stream) self.ptr.reset(coo) - coo.allocate(arr.nnz, arr.shape[0], False, handle_.get_stream()) + coo.allocate(arr.nnz, arr.shape[0], False, stream) handle_.sync_stream() copy_from_cupy(coo.vals(), arr.data, np.float32) @@ -1345,7 +1346,7 @@ class UMAP( init.data.ptr if isinstance(init, cp.ndarray) else init.ctypes.data ), init.nbytes, - handle_.get_stream(), + handle_.get_stream(), any_resource[device_accessible]( get_current_device_resource().get_mr() )