statistics.hpp

 
#pragma once
#include <atomic>
#include <concepts>
#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/error.hpp>
#include <rapidsmpf/memory/memory_type.hpp>
#include <rapidsmpf/memory/pinned_memory_resource.hpp>
#include <rapidsmpf/memory/resource_types.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf {
 
class StreamOrderedTiming;
 
class Statistics : public std::enable_shared_from_this<Statistics> {
  public:
    enum class Formatter : std::uint8_t {
        Default = 0,
        Bytes,
        Duration,
        HitRate,
        MemoryThroughput,
        _Count,  
    };
 
    enum class Mode : std::uint8_t {
        Enabled,  
        Disabled,  
    };
 
    static std::shared_ptr<Statistics> create(Mode mode = Mode::Enabled);
 
    static std::shared_ptr<Statistics> disabled() {
        return create(Mode::Disabled);
    }
 
    static std::shared_ptr<Statistics> from_options(config::Options options);
 
    ~Statistics() noexcept;
 
    // `Statistics` is owned exclusively through `std::shared_ptr`. Use
    // `create()` or `from_options()` to construct instances.
    Statistics(Statistics const&) = delete;
    Statistics& operator=(Statistics const&) = delete;
    Statistics(Statistics&&) = delete;
    Statistics& operator=(Statistics&&) = delete;
 
    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);
    }
 
    struct ReportArgs {
        std::optional<any_device_resource> mr = std::nullopt;
        std::optional<any_host_device_resource> pinned_mr = std::nullopt;
        std::string_view header = "Statistics:";
    };
 
    std::string report(ReportArgs report_args) const;
 
    std::string report() const {
        return report(ReportArgs{});
    }
 
    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();
 
    // TODO: move MemoryRecord and MemoryRecorder to RmmResourceAdaptor?
 
    struct MemoryRecord {
        ScopedMemoryRecord scoped;  
        std::int64_t global_peak{0};  
        std::uint64_t num_calls{0};  
    };
 
    class MemoryRecorder {
      public:
        MemoryRecorder() = default;
 
        MemoryRecorder(
            std::shared_ptr<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:
        std::optional<RmmResourceAdaptor> mr_{std::nullopt};
        std::shared_ptr<Statistics> stats_{nullptr};
        std::string name_{};
    };
 
    MemoryRecorder create_memory_recorder(any_device_resource mr, 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;
    };
 
    explicit Statistics(bool enabled);
 
    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_;
};
 
template <typename T>
concept StatisticsProvider = requires(T const& t) {
    {
        t.statistics()
    } noexcept -> std::same_as<std::shared_ptr<Statistics>>;
};
 
// clang-format off
// Picks between _2 (stats, mr) and _3 (stats, mr, funcname) forms.
#define RAPIDSMPF_MEMORY_PROFILE_PICK_(_1, _2, _3, NAME, ...) NAME
#define RAPIDSMPF_MEMORY_PROFILE(...)                                                     \
    RAPIDSMPF_MEMORY_PROFILE_PICK_(                                                       \
        __VA_ARGS__, RAPIDSMPF_MEMORY_PROFILE_3, RAPIDSMPF_MEMORY_PROFILE_2, ~            \
    )(__VA_ARGS__)
// clang-format on
 
// Version with default function name (__func__)
#define RAPIDSMPF_MEMORY_PROFILE_2(stats, mr) \
    RAPIDSMPF_MEMORY_PROFILE_3(stats, mr, __func__)
 
// Version with custom function name
#define RAPIDSMPF_MEMORY_PROFILE_3(stats, mr, funcname)                                 \
    auto const& RAPIDSMPF_CONCAT(_rapidsmpf_memory_profile_stats_, __LINE__) = (stats); \
    RAPIDSMPF_EXPECTS(                                                                  \
        RAPIDSMPF_CONCAT(_rapidsmpf_memory_profile_stats_, __LINE__) != nullptr,        \
        "RAPIDSMPF_MEMORY_PROFILE: stats must not be null"                              \
    );                                                                                  \
    auto const RAPIDSMPF_CONCAT(_rapidsmpf_memory_recorder_, __LINE__) =                \
        RAPIDSMPF_CONCAT(_rapidsmpf_memory_profile_stats_, __LINE__)                    \
            -> create_memory_recorder(                                                  \
                (mr),                                                                   \
                std::string(__FILE__) + ":" + RAPIDSMPF_STRINGIFY(__LINE__) + "("       \
                    + std::string(funcname) + ")"                                       \
            )
 
}  // namespace rapidsmpf