statistics.hpp

 
#pragma once
#include <cstddef>
#include <functional>
#include <map>
#include <mutex>
#include <string>
#include <utility>
 
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/utils.hpp>
 
namespace rapidsmpf {
 
 
class Statistics {
  public:
    Statistics(bool enabled = true);
 
    Statistics(RmmResourceAdaptor* mr);
 
    ~Statistics() noexcept = default;
    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_)} {}
 
    Statistics& operator=(Statistics&& o) noexcept {
        enabled_ = o.enabled_;
        stats_ = std::move(o.stats_);
        return *this;
    }
 
    bool enabled() const noexcept {
        return enabled_;
    }
 
    std::string report(std::string const& header = "Statistics:") const;
 
    using Formatter = std::function<void(std::ostream&, std::size_t, double)>;
 
    static void FormatterDefault(std::ostream& os, std::size_t count, double val);
 
    class Stat {
      public:
        Stat(Formatter formatter) : formatter_{std::move(formatter)} {}
 
        bool operator==(Stat const& o) const noexcept {
            return count_ == o.count() && value_ == o.value();
        }
 
        double add(double value) {
            ++count_;
            return value_ += value;
        }
 
        [[nodiscard]] std::size_t count() const noexcept {
            return count_;
        }
 
        [[nodiscard]] double value() const noexcept {
            return value_;
        }
 
        [[nodiscard]] Formatter const& formatter() const noexcept {
            return formatter_;
        }
 
      private:
        std::size_t count_{0};
        double value_{0};
        Formatter formatter_;
    };
 
    Stat get_stat(std::string const& name) const;
 
    double add_stat(
        std::string const& name,
        double value,
        Formatter const& formatter = FormatterDefault
    );
 
    std::size_t add_bytes_stat(std::string const& name, std::size_t nbytes);
 
    Duration add_duration_stat(std::string const& name, Duration seconds);
 
    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:
    mutable std::mutex mutex_;
    bool enabled_;
    std::map<std::string, Stat> stats_;
    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 const RAPIDSMPF_CONCAT(_rapidsmpf_memory_recorder_, __LINE__) =             \
        ((rapidsmpf::detail::to_pointer(stats)                                       \
          && rapidsmpf::detail::to_pointer(stats) -> is_memory_profiling_enabled())  \
             ? rapidsmpf::detail::to_pointer(stats)->create_memory_recorder(         \
                   std::string(__FILE__) + ":" + RAPIDSMPF_STRINGIFY(__LINE__) + "(" \
                   + std::string(funcname) + ")"                                     \
               )                                                                     \
             : rapidsmpf::Statistics::MemoryRecorder{})
 
}  // namespace rapidsmpf