device_memory_resource.hpp
1 /*
2  * Copyright (c) 2019-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_stream_view.hpp>
19 #include <rmm/detail/aligned.hpp>
20 
21 #include <cstddef>
22 #include <utility>
23 
24 namespace rmm::mr {
25 
83  public:
84  device_memory_resource() = default;
85  virtual ~device_memory_resource() = default;
87  device_memory_resource& operator=(device_memory_resource const&) = default;
88  device_memory_resource(device_memory_resource&&) noexcept = default;
89  device_memory_resource& operator=(device_memory_resource&&) noexcept = default;
90 
106  void* allocate(std::size_t bytes, cuda_stream_view stream = cuda_stream_view{})
107  {
108  return do_allocate(rmm::detail::align_up(bytes, allocation_size_alignment), stream);
109  }
110 
129  void deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream = cuda_stream_view{})
130  {
131  do_deallocate(ptr, rmm::detail::align_up(bytes, allocation_size_alignment), stream);
132  }
133 
147  [[nodiscard]] bool is_equal(device_memory_resource const& other) const noexcept
148  {
149  return do_is_equal(other);
150  }
151 
158  [[nodiscard]] virtual bool supports_streams() const noexcept = 0;
159 
165  [[nodiscard]] virtual bool supports_get_mem_info() const noexcept = 0;
166 
175  [[nodiscard]] std::pair<std::size_t, std::size_t> get_mem_info(cuda_stream_view stream) const
176  {
177  return do_get_mem_info(stream);
178  }
179 
180  private:
181  // All allocations are padded to a multiple of allocation_size_alignment bytes.
182  static constexpr auto allocation_size_alignment = std::size_t{8};
183 
196  virtual void* do_allocate(std::size_t bytes, cuda_stream_view stream) = 0;
197 
209  virtual void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) = 0;
210 
225  [[nodiscard]] virtual bool do_is_equal(device_memory_resource const& other) const noexcept
226  {
227  return this == &other;
228  }
229 
238  [[nodiscard]] virtual std::pair<std::size_t, std::size_t> do_get_mem_info(
239  cuda_stream_view stream) const = 0;
240 };
241 } // namespace rmm::mr
rmm::mr::device_memory_resource::supports_streams
virtual bool supports_streams() const noexcept=0
Query whether the resource supports use of non-null CUDA streams for allocation/deallocation.
rmm::mr::device_memory_resource::allocate
void * allocate(std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:106
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::device_memory_resource::deallocate
void deallocate(void *ptr, std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:129
rmm::mr::device_memory_resource::is_equal
bool is_equal(device_memory_resource const &other) const noexcept
Compare this resource to another.
Definition: device_memory_resource.hpp:147
rmm::mr::device_memory_resource::supports_get_mem_info
virtual bool supports_get_mem_info() const noexcept=0
Query whether the resource supports the get_mem_info API.
rmm::mr::device_memory_resource::get_mem_info
std::pair< std::size_t, std::size_t > get_mem_info(cuda_stream_view stream) const
Queries the amount of free and total memory for the resource.
Definition: device_memory_resource.hpp:175
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82