stats.h

#pragma once

#include <atomic>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
#include "caffe2/core/logging.h"
#include "caffe2/core/static_tracepoint.h"

namespace caffe2 {

class CAFFE2_API StatValue {
  std::atomic<int64_t> v_{0};

 public:
  int64_t increment(int64_t inc) {
    return v_ += inc;
  }

  int64_t reset(int64_t value = 0) {
    return v_.exchange(value);
  }

  int64_t get() const {
    return v_.load();
  }
};

struct CAFFE2_API ExportedStatValue {
  std::string key;
  int64_t value;
  std::chrono::time_point<std::chrono::high_resolution_clock> ts;
};

using ExportedStatList = std::vector<ExportedStatValue>;
using ExportedStatMap = std::unordered_map<std::string, int64_t>;

CAFFE2_API ExportedStatMap toMap(const ExportedStatList& stats);

class CAFFE2_API StatRegistry {
  std::mutex mutex_;
  std::unordered_map<std::string, std::unique_ptr<StatValue>> stats_;

 public:
  static StatRegistry& get();

  StatValue* add(const std::string& name);

  void publish(ExportedStatList& exported, bool reset = false);

  ExportedStatList publish(bool reset = false) {
    ExportedStatList stats;
    publish(stats, reset);
    return stats;
  }

  void update(const ExportedStatList& data);

  ~StatRegistry();
};

struct CAFFE2_API Stat {
  std::string groupName;
  std::string name;
  Stat(const std::string& gn, const std::string& n) : groupName(gn), name(n) {}

  template <typename... Unused>
  int64_t increment(Unused...) {
    return -1;
  }
};

class CAFFE2_API ExportedStat : public Stat {
  StatValue* value_;

 public:
  ExportedStat(const std::string& gn, const std::string& n)
      : Stat(gn, n), value_(StatRegistry::get().add(gn + "/" + n)) {}

  int64_t increment(int64_t value = 1) {
    return value_->increment(value);
  }

  template <typename T, typename Unused1, typename... Unused>
  int64_t increment(T value, Unused1, Unused...) {
    return increment(value);
  }
};

class CAFFE2_API AvgExportedStat : public ExportedStat {
 private:
  ExportedStat count_;

 public:
  AvgExportedStat(const std::string& gn, const std::string& n)
      : ExportedStat(gn, n + "/sum"), count_(gn, n + "/count") {}

  int64_t increment(int64_t value = 1) {
    count_.increment();
    return ExportedStat::increment(value);
  }

  template <typename T, typename Unused1, typename... Unused>
  int64_t increment(T value, Unused1, Unused...) {
    return increment(value);
  }
};

class CAFFE2_API StdDevExportedStat : public ExportedStat {
  // Uses an offset (first_) to remove issue of cancellation
  // Variance is then (sumsqoffset_ - (sumoffset_^2) / count_) / (count_ - 1)
 private:
  ExportedStat count_;
  ExportedStat sumsqoffset_;
  ExportedStat sumoffset_;
  std::atomic<int64_t> first_{std::numeric_limits<int64_t>::min()};
  int64_t const_min_{std::numeric_limits<int64_t>::min()};

 public:
  StdDevExportedStat(const std::string& gn, const std::string& n)
      : ExportedStat(gn, n + "/sum"),
        count_(gn, n + "/count"),
        sumsqoffset_(gn, n + "/sumsqoffset"),
        sumoffset_(gn, n + "/sumoffset") {}

  int64_t increment(int64_t value = 1) {
    first_.compare_exchange_strong(const_min_, value);
    int64_t offset_value = first_.load();
    int64_t orig_value = value;
    value -= offset_value;
    count_.increment();
    sumsqoffset_.increment(value * value);
    sumoffset_.increment(value);
    return ExportedStat::increment(orig_value);
  }

  template <typename T, typename Unused1, typename... Unused>
  int64_t increment(T value, Unused1, Unused...) {
    return increment(value);
  }
};

class CAFFE2_API DetailedExportedStat : public ExportedStat {
 private:
  std::vector<ExportedStat> details_;

 public:
  DetailedExportedStat(const std::string& gn, const std::string& n)
      : ExportedStat(gn, n) {}

  void setDetails(const std::vector<std::string>& detailNames) {
    details_.clear();
    for (const auto& detailName : detailNames) {
      details_.emplace_back(groupName, name + "/" + detailName);
    }
  }

  template <typename T, typename... Unused>
  int64_t increment(T value, size_t detailIndex, Unused...) {
    if (detailIndex < details_.size()) {
      details_[detailIndex].increment(value);
    }
    return ExportedStat::increment(value);
  }
};

class CAFFE2_API StaticStat : public Stat {
 private:
  StatValue* value_;

 public:
  StaticStat(const std::string& groupName, const std::string& name)
      : Stat(groupName, name),
        value_(StatRegistry::get().add(groupName + "/" + name)) {}

  int64_t increment(int64_t value = 1) {
    return value_->reset(value);
  }

  template <typename T, typename Unused1, typename... Unused>
  int64_t increment(T value, Unused1, Unused...) {
    return increment(value);
  }
};

namespace detail {

template <class T>
struct _ScopeGuard {
  T f_;
  std::chrono::high_resolution_clock::time_point start_;

  explicit _ScopeGuard(T f)
      : f_(f), start_(std::chrono::high_resolution_clock::now()) {}
  ~_ScopeGuard() {
    using namespace std::chrono;
    auto duration = high_resolution_clock::now() - start_;
    int64_t nanos = duration_cast<nanoseconds>(duration).count();
    f_(nanos);
  }

  // Using implicit cast to bool so that it can be used in an 'if' condition
  // within CAFFE_DURATION macro below.
  /* implicit */ operator bool() {
    return true;
  }
};

template <class T>
_ScopeGuard<T> ScopeGuard(T f) {
  return _ScopeGuard<T>(f);
}
} // namespace detail

#define CAFFE_STAT_CTOR(ClassName)                 \
  ClassName(std::string name) : groupName(name) {} \
  std::string groupName

#define CAFFE_EXPORTED_STAT(name) \
  ExportedStat name {             \
    groupName, #name              \
  }

#define CAFFE_AVG_EXPORTED_STAT(name) \
  AvgExportedStat name {              \
    groupName, #name                  \
  }

#define CAFFE_STDDEV_EXPORTED_STAT(name) \
  StdDevExportedStat name {              \
    groupName, #name                     \
  }

#define CAFFE_DETAILED_EXPORTED_STAT(name) \
  DetailedExportedStat name {              \
    groupName, #name                       \
  }

#define CAFFE_STAT(name) \
  Stat name {            \
    groupName, #name     \
  }

#define CAFFE_STATIC_STAT(name) \
  StaticStat name {             \
    groupName, #name            \
  }

#define CAFFE_EVENT(stats, field, ...)                              \
  {                                                                 \
    auto __caffe_event_value_ = stats.field.increment(__VA_ARGS__); \
    CAFFE_SDT(                                                      \
        field,                                                      \
        stats.field.groupName.c_str(),                              \
        __caffe_event_value_,                                       \
        ##__VA_ARGS__);                                             \
  }

#define CAFFE_DURATION(stats, field, ...)                        \
  if (auto g = ::caffe2::detail::ScopeGuard([&](int64_t nanos) { \
        CAFFE_EVENT(stats, field, nanos, ##__VA_ARGS__);         \
      }))
} // namespace caffe2