callback_memory_resource.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#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) noexcept 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