rmm_resource_adaptor.hpp

 
#pragma once
 
#include <cstddef>
#include <mutex>
#include <optional>
#include <stack>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <utility>
 
#include <rmm/error.hpp>
#include <rmm/mr/device_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/memory/scoped_memory_record.hpp>
 
namespace rapidsmpf {
 
class RmmResourceAdaptor final : public rmm::mr::device_memory_resource {
  public:
    RmmResourceAdaptor(
        rmm::device_async_resource_ref primary_mr,
        std::optional<rmm::device_async_resource_ref> fallback_mr = std::nullopt
    )
        : primary_mr_{std::move(primary_mr)}, fallback_mr_{std::move(fallback_mr)} {}
 
    RmmResourceAdaptor() = delete;
    ~RmmResourceAdaptor() override = default;
 
    [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept {
        return primary_mr_;
    }
 
    [[nodiscard]] std::optional<rmm::device_async_resource_ref>
    get_fallback_resource() const noexcept {
        return fallback_mr_;
    }
 
    [[nodiscard]] ScopedMemoryRecord get_main_record() const;
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept;
 
 
    void begin_scoped_memory_record();
 
    ScopedMemoryRecord end_scoped_memory_record();
 
  private:
    void* do_allocate(std::size_t bytes, rmm::cuda_stream_view stream) override;
 
    void do_deallocate(
        void* ptr, std::size_t bytes, rmm::cuda_stream_view stream
    ) noexcept override;
 
    [[nodiscard]] bool do_is_equal(
        rmm::mr::device_memory_resource const& other
    ) const noexcept override;
 
    mutable std::mutex mutex_;
    rmm::device_async_resource_ref primary_mr_;
    std::optional<rmm::device_async_resource_ref> 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