Skip to content
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

Accounting for rmm::cuda_stream_pool not having a constructor for 0 streams #329

Merged
merged 4 commits into from
Sep 17, 2021
Merged

Accounting for rmm::cuda_stream_pool not having a constructor for 0 streams #329

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
94e28a3
making rmm::cuda_stream_pool on raft handle a unique pointer
divyegala Sep 17, 2021
3bf26b9
minor update
divyegala Sep 17, 2021
f525414
fixing compilation errors
divyegala Sep 17, 2021
d073c12
tests finally passing
divyegala Sep 17, 2021
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
22 changes: 17 additions & 5 deletions cpp/include/raft/handle.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,13 @@ class handle_t {
CUDA_CHECK(cudaGetDevice(&cur_dev));
return cur_dev;
}()),
streams_(n_streams) {
streams_{[n_streams]() {
if (n_streams == 0) {
return std::nullptr_t;
} else {
return std::make_unique<rmm::cuda_stream_pool>(n_streams);
}
}()} {
create_resources();
thrust_policy_ = std::make_unique<rmm::exec_policy>(user_stream_);
}
Expand Down Expand Up @@ -140,14 +146,20 @@ class handle_t {

// legacy compatibility for cuML
cudaStream_t get_internal_stream(int sid) const {
return streams_.get_stream(sid).value();
RAFT_EXPECTS(
streams_.get() != nullptr,
"ERROR: rmm::cuda_stream_pool was not initialized with a non-zero value");
return streams_->get_stream(sid).value();
}
// new accessor return rmm::cuda_stream_view
rmm::cuda_stream_view get_internal_stream_view(int sid) const {
return streams_.get_stream(sid);
RAFT_EXPECTS(
streams_.get() != nullptr,
"ERROR: rmm::cuda_stream_pool was not initialized with a non-zero value");
return streams_->get_stream(sid);
}

int get_num_internal_streams() const { return streams_.get_pool_size(); }
int get_num_internal_streams() const { return streams_->get_pool_size(); }
std::vector<cudaStream_t> get_internal_streams() const {
std::vector<cudaStream_t> int_streams_vec;
for (int i = 0; i < get_num_internal_streams(); i++) {
Expand Down Expand Up @@ -212,7 +224,7 @@ class handle_t {
std::unordered_map<std::string, std::shared_ptr<comms::comms_t>> subcomms_;

const int dev_id_;
rmm::cuda_stream_pool streams_{0};
std::unique_ptr<rmm::cuda_stream_pool> streams_;
mutable cublasHandle_t cublas_handle_;
mutable bool cublas_initialized_{false};
mutable cusolverDnHandle_t cusolver_dn_handle_;
Expand Down