All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
callback_memory_resource.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2022-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 #pragma once
17 
18 #include <rmm/detail/export.hpp>
20 
21 #include <cstddef>
22 #include <functional>
23 #include <utility>
24 
25 namespace RMM_NAMESPACE {
26 namespace mr {
48 using allocate_callback_t = std::function<void*(std::size_t, cuda_stream_view, void*)>;
49 
66 using deallocate_callback_t = std::function<void(void*, std::size_t, cuda_stream_view, void*)>;
67 
73  public:
90  allocate_callback_t allocate_callback,
91  deallocate_callback_t deallocate_callback,
92  void* allocate_callback_arg = nullptr, // NOLINT(bugprone-easily-swappable-parameters)
93  void* deallocate_callback_arg = nullptr) noexcept
94  : allocate_callback_(std::move(allocate_callback)),
95  deallocate_callback_(std::move(deallocate_callback)),
96  allocate_callback_arg_(allocate_callback_arg),
97  deallocate_callback_arg_(deallocate_callback_arg)
98  {
99  }
100 
101  callback_memory_resource() = delete;
102  ~callback_memory_resource() override = default;
104  callback_memory_resource& operator=(callback_memory_resource const&) = delete;
106  default;
108  default;
109 
110  private:
123  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
124  {
125  return allocate_callback_(bytes, stream, allocate_callback_arg_);
126  }
127 
139  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
140  {
141  deallocate_callback_(ptr, bytes, stream, deallocate_callback_arg_);
142  }
143 
144  allocate_callback_t allocate_callback_;
145  deallocate_callback_t deallocate_callback_;
146  void* allocate_callback_arg_;
147  void* deallocate_callback_arg_;
148 };
149  // end of group
151 } // namespace mr
152 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
A device memory resource that uses the provided callbacks for memory allocation and deallocation.
Definition: callback_memory_resource.hpp:72
callback_memory_resource(callback_memory_resource &&) noexcept=default
Default move constructor.
callback_memory_resource(allocate_callback_t allocate_callback, deallocate_callback_t deallocate_callback, void *allocate_callback_arg=nullptr, void *deallocate_callback_arg=nullptr) noexcept
Construct a new callback memory resource.
Definition: callback_memory_resource.hpp:89
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
std::function< void *(std::size_t, cuda_stream_view, void *)> allocate_callback_t
Callback function type used by callback memory resource for allocation.
Definition: callback_memory_resource.hpp:48
std::function< void(void *, std::size_t, cuda_stream_view, void *)> deallocate_callback_t
Callback function type used by callback_memory_resource for deallocation.
Definition: callback_memory_resource.hpp:66