callback_memory_resource.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2022-2024, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
 
#include <cstddef>
#include <functional>
#include <utility>
 
namespace RMM_NAMESPACE {
namespace mr {
using allocate_callback_t = std::function<void*(std::size_t, cuda_stream_view, void*)>;
 
using deallocate_callback_t = std::function<void(void*, std::size_t, cuda_stream_view, void*)>;
 
class callback_memory_resource final : public device_memory_resource {
 public:
  callback_memory_resource(
    allocate_callback_t allocate_callback,
    deallocate_callback_t deallocate_callback,
    void* allocate_callback_arg   = nullptr,  // NOLINT(bugprone-easily-swappable-parameters)
    void* deallocate_callback_arg = nullptr) noexcept
    : allocate_callback_(std::move(allocate_callback)),
      deallocate_callback_(std::move(deallocate_callback)),
      allocate_callback_arg_(allocate_callback_arg),
      deallocate_callback_arg_(deallocate_callback_arg)
  {
  }
 
  callback_memory_resource()                                           = delete;
  ~callback_memory_resource() override                                 = default;
  callback_memory_resource(callback_memory_resource const&)            = delete;
  callback_memory_resource& operator=(callback_memory_resource const&) = delete;
  callback_memory_resource(callback_memory_resource&&) noexcept =
    default;  
  callback_memory_resource& operator=(callback_memory_resource&&) noexcept =
    default;  
 
 private:
  void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
  {
    return allocate_callback_(bytes, stream, allocate_callback_arg_);
  }
 
  void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override
  {
    deallocate_callback_(ptr, bytes, stream, deallocate_callback_arg_);
  }
 
  allocate_callback_t allocate_callback_;
  deallocate_callback_t deallocate_callback_;
  void* allocate_callback_arg_;
  void* deallocate_callback_arg_;
};
  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE