cuda_async_view_memory_resource.hpp
1 /*
2  * Copyright (c) 2021, 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 #pragma once
17 
18 #include <rmm/cuda_device.hpp>
19 #include <rmm/cuda_stream_view.hpp>
20 #include <rmm/detail/cuda_util.hpp>
21 #include <rmm/detail/dynamic_load_runtime.hpp>
22 #include <rmm/detail/error.hpp>
23 #include <rmm/mr/device/device_memory_resource.hpp>
24 
25 #include <thrust/optional.h>
26 
27 #include <cuda_runtime_api.h>
28 
29 #include <cstddef>
30 #include <limits>
31 
32 #if CUDART_VERSION >= 11020 // 11.2 introduced cudaMallocAsync
33 #define RMM_CUDA_MALLOC_ASYNC_SUPPORT
34 #endif
35 
36 namespace rmm::mr {
37 
43  public:
44 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
45 
55  cuda_async_view_memory_resource(cudaMemPool_t valid_pool_handle)
56  : cuda_pool_handle_{[valid_pool_handle]() {
57  RMM_EXPECTS(nullptr != valid_pool_handle, "Unexpected null pool handle.");
58  return valid_pool_handle;
59  }()}
60  {
61  // Check if cudaMallocAsync Memory pool supported
62  auto const device = rmm::detail::current_device();
63  int cuda_pool_supported{};
64  auto result =
65  cudaDeviceGetAttribute(&cuda_pool_supported, cudaDevAttrMemoryPoolsSupported, device.value());
66  RMM_EXPECTS(result == cudaSuccess && cuda_pool_supported,
67  "cudaMallocAsync not supported with this CUDA driver/runtime version");
68  }
69 #endif
70 
71 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
72 
76  [[nodiscard]] cudaMemPool_t pool_handle() const noexcept { return cuda_pool_handle_; }
77 #endif
78 
84 
91  [[nodiscard]] bool supports_streams() const noexcept override { return true; }
92 
98  [[nodiscard]] bool supports_get_mem_info() const noexcept override { return false; }
99 
100  private:
101 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
102  cudaMemPool_t cuda_pool_handle_{};
103 #endif
104 
115  void* do_allocate(std::size_t bytes, rmm::cuda_stream_view stream) override
116  {
117  void* ptr{nullptr};
118 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
119  if (bytes > 0) {
120  RMM_CUDA_TRY_ALLOC(rmm::detail::async_alloc::cudaMallocFromPoolAsync(
121  &ptr, bytes, pool_handle(), stream.value()));
122  }
123 #else
124  (void)bytes;
125  (void)stream;
126 #endif
127  return ptr;
128  }
129 
137  void do_deallocate(void* ptr, std::size_t, rmm::cuda_stream_view stream) override
138  {
139 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
140  if (ptr != nullptr) {
141  RMM_ASSERT_CUDA_SUCCESS(rmm::detail::async_alloc::cudaFreeAsync(ptr, stream.value()));
142  }
143 #else
144  (void)ptr;
145  (void)stream;
146 #endif
147  }
148 
158  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
159  {
160  return dynamic_cast<cuda_async_view_memory_resource const*>(&other) != nullptr;
161  }
162 
170  [[nodiscard]] std::pair<std::size_t, std::size_t> do_get_mem_info(
171  rmm::cuda_stream_view) const override
172  {
173  return std::make_pair(0, 0);
174  }
175 };
176 
177 } // namespace rmm::mr
rmm::mr::cuda_async_view_memory_resource::supports_streams
bool supports_streams() const noexcept override
Query whether the resource supports use of non-null CUDA streams for allocation/deallocation....
Definition: cuda_async_view_memory_resource.hpp:91
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::cuda_async_view_memory_resource
device_memory_resource derived class that uses cudaMallocAsync/cudaFreeAsync for allocation/deallocat...
Definition: cuda_async_view_memory_resource.hpp:42
rmm::mr::cuda_async_view_memory_resource::supports_get_mem_info
bool supports_get_mem_info() const noexcept override
Query whether the resource supports the get_mem_info API.
Definition: cuda_async_view_memory_resource.hpp:98
rmm::cuda_stream_view::value
constexpr cudaStream_t value() const noexcept
Get the wrapped stream.
Definition: cuda_stream_view.hpp:57
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82