rmm_resource_adaptor.hpp

 
#pragma once
 
#include <array>
#include <cstddef>
#include <mutex>
#include <optional>
#include <stack>
#include <thread>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
 
#include <rmm/error.hpp>
#include <rmm/mr/device_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
 
namespace rapidsmpf {
 
struct ScopedMemoryRecord {
    enum class AllocType : std::size_t {
        PRIMARY = 0,  
        FALLBACK = 1,  
        ALL = 2  
    };
 
    using AllocTypeArray = std::array<std::int64_t, 2>;
 
    [[nodiscard]] std::int64_t num_total_allocs(
        AllocType alloc_type = AllocType::ALL
    ) const noexcept;
 
    [[nodiscard]] std::int64_t num_current_allocs(
        AllocType alloc_type = AllocType::ALL
    ) const noexcept;
 
    [[nodiscard]] std::int64_t current(
        AllocType alloc_type = AllocType::ALL
    ) const noexcept;
 
    [[nodiscard]] std::int64_t total(
        AllocType alloc_type = AllocType::ALL
    ) const noexcept;
 
    [[nodiscard]] std::int64_t peak(AllocType alloc_type = AllocType::ALL) const noexcept;
 
    void record_allocation(AllocType alloc_type, std::int64_t nbytes);
 
    void record_deallocation(AllocType alloc_type, std::int64_t nbytes);
 
    ScopedMemoryRecord& add_subscope(ScopedMemoryRecord const& subscope);
 
    ScopedMemoryRecord& add_scope(ScopedMemoryRecord const& scope);
 
  private:
    AllocTypeArray num_current_allocs_{{0, 0}};
    AllocTypeArray num_total_allocs_{{0, 0}};
    AllocTypeArray current_{{0, 0}};
    AllocTypeArray total_{{0, 0}};
    AllocTypeArray peak_{{0, 0}};
    std::int64_t highest_peak_{0};
};
 
static_assert(
    std::is_trivially_copyable_v<ScopedMemoryRecord>,
    "ScopedMemoryRecord must be trivially copyable"
);
 
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_{primary_mr}, fallback_mr_{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