pinned_memory_resource.hpp
1 /*
2  * Copyright (c) 2020-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/detail/aligned.hpp>
19 #include <rmm/detail/error.hpp>
20 #include <rmm/mr/host/host_memory_resource.hpp>
21 
22 #include <cstddef>
23 #include <utility>
24 
25 namespace rmm::mr {
26 
27 /*
28  * @brief A `host_memory_resource` that uses `cudaMallocHost` to allocate
29  * pinned/page-locked host memory.
30  *
31  * See https://devblogs.nvidia.com/how-optimize-data-transfers-cuda-cc/
32  */
34  public:
35  pinned_memory_resource() = default;
36  ~pinned_memory_resource() override = default;
39  pinned_memory_resource& operator=(pinned_memory_resource const&) = default;
40  pinned_memory_resource& operator=(pinned_memory_resource&&) = default;
41 
42  private:
55  void* do_allocate(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t)) override
56  {
57  // don't allocate anything if the user requested zero bytes
58  if (0 == bytes) { return nullptr; }
59 
60  // If the requested alignment isn't supported, use default
61  alignment = (rmm::detail::is_supported_alignment(alignment))
62  ? alignment
63  : rmm::detail::RMM_DEFAULT_HOST_ALIGNMENT;
64 
65  return rmm::detail::aligned_allocate(bytes, alignment, [](std::size_t size) {
66  void* ptr{nullptr};
67  auto status = cudaMallocHost(&ptr, size);
68  if (cudaSuccess != status) { throw std::bad_alloc{}; }
69  return ptr;
70  });
71  }
72 
88  void do_deallocate(void* ptr,
89  std::size_t bytes,
90  std::size_t alignment = alignof(std::max_align_t)) override
91  {
92  if (nullptr == ptr) { return; }
93  rmm::detail::aligned_deallocate(
94  ptr, bytes, alignment, [](void* ptr) { RMM_ASSERT_CUDA_SUCCESS(cudaFreeHost(ptr)); });
95  }
96 };
97 } // namespace rmm::mr
rmm::mr::host_memory_resource
Base class for host memory allocation.
Definition: host_memory_resource.hpp:47
rmm::mr::pinned_memory_resource
Definition: pinned_memory_resource.hpp:33