statistics.hpp

 
#pragma once
#include <atomic>
#include <cstddef>
#include <filesystem>
#include <initializer_list>
#include <limits>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <ostream>
#include <span>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
 
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/memory/memory_type.hpp>
#include <rapidsmpf/memory/pinned_memory_resource.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf {
 
class StreamOrderedTiming;
 
class Statistics {
  public:
    enum class Formatter : std::uint8_t {
        Default = 0,
        Bytes,
        Duration,
        HitRate,
        MemoryThroughput,
        _Count,  
    };
 
    Statistics(bool enabled = true);
 
    Statistics(
        RmmResourceAdaptor mr,
        std::shared_ptr<PinnedMemoryResource> pinned_mr = PinnedMemoryResource::Disabled
    );
 
    static std::shared_ptr<Statistics> from_options(
        RmmResourceAdaptor mr,
        config::Options options,
        std::shared_ptr<PinnedMemoryResource> pinned_mr = PinnedMemoryResource::Disabled
    );
 
    ~Statistics() noexcept;
 
    // `Statistics` is owned exclusively through `std::shared_ptr` (see `disabled()` and
    // `from_options()`).
    Statistics(Statistics const&) = delete;
    Statistics& operator=(Statistics const&) = delete;
    Statistics(Statistics&&) = delete;
    Statistics& operator=(Statistics&&) = delete;
 
    static std::shared_ptr<Statistics> disabled();
 
    bool enabled() const noexcept {
        return enabled_.load(std::memory_order_acquire);
    }
 
    void enable() noexcept {
        enabled_.store(true, std::memory_order_release);
    }
 
    void disable() noexcept {
        enabled_.store(false, std::memory_order_release);
    }
 
    std::string report(std::string const& header = "Statistics:") const;
 
    void write_json(std::ostream& os) const;
 
    void write_json(std::filesystem::path const& filepath) const;
 
    [[nodiscard]] std::shared_ptr<Statistics> copy() const;
 
    [[nodiscard]] std::vector<std::uint8_t> serialize() const;
 
    [[nodiscard]] static std::shared_ptr<Statistics> deserialize(
        std::span<std::uint8_t const> data
    );
 
    [[nodiscard]] static std::shared_ptr<Statistics> merge(
        std::span<std::shared_ptr<Statistics> const> stats
    );
 
    class Stat {
      public:
        Stat() = default;
 
        Stat(std::size_t count, double value, double max);
 
        auto operator<=>(Stat const&) const noexcept = default;
 
        void add(double value);
 
        [[nodiscard]] std::size_t count() const noexcept;
 
        [[nodiscard]] double value() const noexcept;
 
        [[nodiscard]] double max() const noexcept;
 
        [[nodiscard]] static constexpr std::size_t serialized_size() noexcept {
            return sizeof(std::uint64_t) + sizeof(double) + sizeof(double);
        }
 
        std::uint8_t* serialize(std::uint8_t* out) const;
 
        [[nodiscard]] static std::pair<Stat, std::span<std::uint8_t const>> deserialize(
            std::span<std::uint8_t const> data
        );
 
        [[nodiscard]] Stat merge(Stat const& other) const;
 
      private:
        std::size_t count_{0};
        double value_{0};
        double max_{-std::numeric_limits<double>::infinity()};
    };
 
    Stat get_stat(std::string const& name) const;
 
    void add_stat(std::string const& name, double value);
 
    void add_report_entry(
        std::string const& report_entry_name,
        std::initializer_list<std::string_view> stat_names,
        Formatter formatter
    );
 
    // clang-format off
    // clang-format on
    void add_report_entry(
        std::string const& report_entry_name,
        std::vector<std::string> stat_names,
        Formatter formatter
    );
 
    void add_bytes_stat(std::string const& name, std::size_t nbytes);
 
    void add_duration_stat(std::string const& name, Duration seconds);
 
    void record_copy(
        MemoryType src, MemoryType dst, std::size_t nbytes, StreamOrderedTiming&& timing
    );
 
    void record_alloc(
        MemoryType mem_type, std::size_t nbytes, StreamOrderedTiming&& timing
    );
 
    std::vector<std::string> list_stat_names() const;
 
    void clear();
 
    bool is_memory_profiling_enabled() const;
 
    struct MemoryRecord {
        ScopedMemoryRecord scoped;  
        std::int64_t global_peak{0};  
        std::uint64_t num_calls{0};  
    };
 
    class MemoryRecorder {
      public:
        MemoryRecorder() = default;
 
        MemoryRecorder(Statistics* stats, RmmResourceAdaptor mr, std::string name);
 
        ~MemoryRecorder();
 
        MemoryRecorder(MemoryRecorder const&) = delete;
        MemoryRecorder& operator=(MemoryRecorder const&) = delete;
        MemoryRecorder(MemoryRecorder&&) = delete;
        MemoryRecorder& operator=(MemoryRecorder&&) = delete;
 
      private:
        Statistics* stats_{nullptr};
        std::optional<RmmResourceAdaptor> mr_;
        std::string name_;
    };
 
    MemoryRecorder create_memory_recorder(std::string name);
 
    std::unordered_map<std::string, MemoryRecord> const& get_memory_records() const;
 
  private:
    struct ReportEntry {
        std::vector<std::string> stat_names;
        Formatter formatter;
    };
 
    mutable std::mutex mutex_;
    std::atomic<bool> enabled_;
    std::map<std::string, Stat> stats_;
    std::map<std::string, ReportEntry> report_entries_;
    std::unordered_map<std::string, MemoryRecord> memory_records_;
    std::optional<RmmResourceAdaptor> mr_;
    std::shared_ptr<PinnedMemoryResource>
        pinned_mr_;  
};
 
#define RAPIDSMPF_MEMORY_PROFILE(...)                                       \
    RAPIDSMPF_OVERLOAD_BY_ARG_COUNT(                                        \
        __VA_ARGS__, RAPIDSMPF_MEMORY_PROFILE_2, RAPIDSMPF_MEMORY_PROFILE_1 \
    )                                                                       \
    (__VA_ARGS__)
 
// Version with default function name (__func__)
#define RAPIDSMPF_MEMORY_PROFILE_1(stats) RAPIDSMPF_MEMORY_PROFILE_2(stats, __func__)
 
// Version with custom function name
#define RAPIDSMPF_MEMORY_PROFILE_2(stats, funcname)                                      \
    auto&& RAPIDSMPF_CONCAT(_rapidsmpf_stats_, __LINE__) = (stats);                      \
    auto const RAPIDSMPF_CONCAT(_rapidsmpf_memory_recorder_, __LINE__) =                 \
        ((rapidsmpf::detail::to_pointer(RAPIDSMPF_CONCAT(_rapidsmpf_stats_, __LINE__))   \
          && rapidsmpf::detail::to_pointer(                                              \
                 RAPIDSMPF_CONCAT(_rapidsmpf_stats_, __LINE__)                           \
          ) -> is_memory_profiling_enabled())                                            \
             ? rapidsmpf::detail::to_pointer(                                            \
                   RAPIDSMPF_CONCAT(_rapidsmpf_stats_, __LINE__)                         \
               )                                                                         \
                   ->create_memory_recorder(                                             \
                       std::string(__FILE__) + ":" + RAPIDSMPF_STRINGIFY(__LINE__) + "(" \
                       + std::string(funcname) + ")"                                     \
                   )                                                                     \
             : rapidsmpf::Statistics::MemoryRecorder{})
 
}  // namespace rapidsmpf