scoped_memory_record.hpp

 
#pragma once
 
#include <array>
#include <cstddef>
#include <type_traits>
 
#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"
);
 
}  // namespace rapidsmpf