rmm_resource_adaptor_impl.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <mutex>
#include <optional>
#include <stack>
#include <thread>
#include <unordered_map>
#include <unordered_set>
 
#include <cuda/memory_resource>
 
#include <rmm/aligned.hpp>
#include <rmm/error.hpp>
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/memory/scoped_memory_record.hpp>
 
namespace rapidsmpf::detail {
 
class RmmResourceAdaptorImpl {
  public:
    RmmResourceAdaptorImpl(
        cuda::mr::any_resource<cuda::mr::device_accessible> primary_mr,
        std::optional<cuda::mr::any_resource<cuda::mr::device_accessible>> fallback_mr
    );
 
    ~RmmResourceAdaptorImpl() = default;
 
    RmmResourceAdaptorImpl(RmmResourceAdaptorImpl const&) = delete;
    RmmResourceAdaptorImpl(RmmResourceAdaptorImpl&&) = delete;
    RmmResourceAdaptorImpl& operator=(RmmResourceAdaptorImpl const&) = delete;
    RmmResourceAdaptorImpl& operator=(RmmResourceAdaptorImpl&&) = delete;
 
    [[nodiscard]] bool operator==(RmmResourceAdaptorImpl const& other) const noexcept;
 
    [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept;
 
    [[nodiscard]] std::optional<rmm::device_async_resource_ref>
    get_fallback_resource() const noexcept;
 
    [[nodiscard]] ScopedMemoryRecord get_main_record() const;
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept;
 
    void begin_scoped_memory_record();
 
    ScopedMemoryRecord end_scoped_memory_record();
 
    void* allocate(
        cuda::stream_ref stream,
        std::size_t bytes,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    );
 
    void deallocate(
        cuda::stream_ref stream,
        void* ptr,
        std::size_t bytes,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept;
 
    void* allocate_sync(
        std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    );
 
    void deallocate_sync(
        void* ptr,
        std::size_t bytes,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept;
 
    friend void get_property(
        RmmResourceAdaptorImpl const&, cuda::mr::device_accessible
    ) noexcept {}
 
  private:
    mutable std::mutex mutex_;
    cuda::mr::any_resource<cuda::mr::device_accessible> primary_mr_;
    std::optional<cuda::mr::any_resource<cuda::mr::device_accessible>> fallback_mr_;
    std::unordered_set<void*> fallback_allocations_;
 
    ScopedMemoryRecord main_record_;
    std::unordered_map<std::thread::id, std::stack<ScopedMemoryRecord>> record_stacks_;
    std::unordered_map<void*, std::thread::id> allocating_threads_;
};
 
}  // namespace rapidsmpf::detail