All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
cuda_stream.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2020-2024, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <rmm/cuda_stream_view.hpp>
20 #include <rmm/detail/error.hpp>
21 #include <rmm/detail/export.hpp>
22 #include <rmm/detail/logging_assert.hpp>
23 
24 #include <cuda_runtime_api.h>
25 
26 #include <functional>
27 #include <memory>
28 
29 namespace RMM_NAMESPACE {
41 class cuda_stream {
42  public:
49  cuda_stream(cuda_stream&&) = default;
59  ~cuda_stream() = default;
60  cuda_stream(cuda_stream const&) = delete; // Copying disallowed: one stream one owner
61  cuda_stream& operator=(cuda_stream&) = delete;
62 
69  : stream_{[]() {
70  auto* stream = new cudaStream_t; // NOLINT(cppcoreguidelines-owning-memory)
71  RMM_CUDA_TRY(cudaStreamCreate(stream));
72  return stream;
73  }(),
74  [](cudaStream_t* stream) {
75  RMM_ASSERT_CUDA_SUCCESS(cudaStreamDestroy(*stream));
76  delete stream; // NOLINT(cppcoreguidelines-owning-memory)
77  }}
78  {
79  }
80 
87  [[nodiscard]] bool is_valid() const { return stream_ != nullptr; }
88 
94  [[nodiscard]] cudaStream_t value() const
95  {
96  RMM_LOGGING_ASSERT(is_valid());
97  return *stream_;
98  }
99 
103  explicit operator cudaStream_t() const noexcept { return value(); }
104 
110  [[nodiscard]] cuda_stream_view view() const { return cuda_stream_view{value()}; }
111 
117  operator cuda_stream_view() const { return view(); }
118 
126  void synchronize() const { RMM_CUDA_TRY(cudaStreamSynchronize(value())); }
127 
133  void synchronize_no_throw() const noexcept
134  {
135  RMM_ASSERT_CUDA_SUCCESS(cudaStreamSynchronize(value()));
136  }
137 
138  private:
139  std::unique_ptr<cudaStream_t, std::function<void(cudaStream_t*)>> stream_;
140 };
141  // end of group
143 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:41
Owning wrapper for a CUDA stream.
Definition: cuda_stream.hpp:41
void synchronize_no_throw() const noexcept
Synchronize the owned CUDA stream. Does not throw if there is an error.
Definition: cuda_stream.hpp:133
bool is_valid() const
Returns true if the owned stream is non-null.
Definition: cuda_stream.hpp:87
void synchronize() const
Synchronize the owned CUDA stream.
Definition: cuda_stream.hpp:126
cuda_stream()
Construct a new cuda stream object.
Definition: cuda_stream.hpp:68
cuda_stream & operator=(cuda_stream &&)=default
Move copy assignment operator (default)
cudaStream_t value() const
Get the value of the wrapped CUDA stream.
Definition: cuda_stream.hpp:94
cuda_stream_view view() const
Creates an immutable, non-owning view of the wrapped CUDA stream.
Definition: cuda_stream.hpp:110
cuda_stream(cuda_stream &&)=default
Move constructor (default)