statistics.hpp

 
#pragma once
#include <atomic>
#include <cstddef>
#include <filesystem>
#include <functional>
#include <limits>
#include <map>
#include <mutex>
#include <ostream>
#include <string>
#include <utility>
#include <vector>
 
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/memory/memory_type.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf {
 
class StreamOrderedTiming;
 
class Statistics {
  public:
    Statistics(bool enabled = true);
 
    Statistics(RmmResourceAdaptor* mr);
 
    static std::shared_ptr<Statistics> from_options(
        RmmResourceAdaptor* mr, config::Options options
    );
 
    ~Statistics() noexcept;
    Statistics(Statistics const&) = delete;
    Statistics& operator=(Statistics const&) = delete;
 
    static std::shared_ptr<Statistics> disabled();
 
    Statistics(Statistics&& o) noexcept
        : enabled_(o.enabled()),
          stats_{std::move(o.stats_)},
          formatters_{std::move(o.formatters_)} {}
 
    Statistics& operator=(Statistics&& o) noexcept {
        enabled_ = o.enabled();
        stats_ = std::move(o.stats_);
        formatters_ = std::move(o.formatters_);
        return *this;
    }
 
    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;
 
    class Stat {
      public:
        Stat() = default;
 
        auto operator<=>(Stat const&) const noexcept = default;
 
        void add(double value) {
            ++count_;
            value_ += value;
            max_ = std::max(max_, value);
        }
 
        [[nodiscard]] std::size_t count() const noexcept {
            return count_;
        }
 
        [[nodiscard]] double value() const noexcept {
            return value_;
        }
 
        [[nodiscard]] double max() const noexcept {
            return max_;
        }
 
      private:
        std::size_t count_{0};
        double value_{0};
        double max_{-std::numeric_limits<double>::infinity()};
    };
 
    using Formatter = std::function<void(std::ostream&, std::vector<Stat> const&)>;
 
    Stat get_stat(std::string const& name) const;
 
    void add_stat(std::string const& name, double value);
 
    bool exist_report_entry_name(std::string const& name) const;
 
    void register_formatter(std::string const& name, Formatter formatter);
 
    void register_formatter(
        std::string const& report_entry_name,
        std::vector<std::string> const& 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};
        RmmResourceAdaptor* mr_{nullptr};
        std::string name_;
        ScopedMemoryRecord main_record_;
    };
 
    MemoryRecorder create_memory_recorder(std::string name);
 
    std::unordered_map<std::string, MemoryRecord> const& get_memory_records() const;
 
  private:
    struct FormatterEntry {
        std::vector<std::string> stat_names;  
        Formatter fn;
    };
 
    mutable std::mutex mutex_;
    std::atomic<bool> enabled_;
    std::map<std::string, Stat> stats_;
    std::map<std::string, FormatterEntry> formatters_;
    std::unordered_map<std::string, MemoryRecord> memory_records_;
    RmmResourceAdaptor* 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