unsupported/Eigen/CXX11/src/Tensor/TensorStats.h

#ifndef EIGEN_CXX11_TENSOR_TENSOR_STATS_H
#define EIGEN_CXX11_TENSOR_TENSOR_STATS_H

#include <mutex>
#include <map>
using ::std::map;
using ::std::multimap;
#include <csignal>

namespace Eigen {
namespace internal {

struct MatmulOp {
  enum class Algorithm { GEBP, XSMM, GEMV };
  Algorithm alg;
  size_t m, k, n;
  bool transposeA, transposeB;
  size_t num_threads;
  typedef std::tuple<Algorithm, size_t, size_t, size_t, bool, bool, size_t> tpl;
  tpl to_tuple() const {
    return tpl(alg, m, k, n, transposeA, transposeB, num_threads);
  }
};

#ifdef EIGEN_TENSOR_STATS

EIGEN_ALWAYS_INLINE bool operator<(const MatmulOp& m1,
                                   const MatmulOp& m2) {
  return m1.to_tuple() < m2.to_tuple();
}

EIGEN_ALWAYS_INLINE std::ostream& operator<<(std::ostream& os,
                                             const MatmulOp::Algorithm& a) {
  if (a == MatmulOp::Algorithm::GEBP) {
    os << "eigen-generic";
  } else if (a == MatmulOp::Algorithm::XSMM) {
    os << "libxsmm";
  } else if (a == MatmulOp::Algorithm::GEMV) {
    os << "matrix-vector";
  }
  return os;
}

EIGEN_ALWAYS_INLINE std::ostream& operator<<(std::ostream& os,
                                             const MatmulOp& m) {
  os << "MATMUL ";
  os << (m.transposeA ? "A'" : "A") << "*";
  os << (m.transposeB ? "B'" : "B") << "\t";
  os << "(" << m.m << ", " << m.k << ", " << m.n << ")" << "\t";
  os << m.num_threads << (m.num_threads != 1 ? " threads" : " thread") << "\t";
  os << "using " << m.alg;
  os << std::endl;
  return os;
}

class EigenStats {
  std::map<MatmulOp, unsigned> counts;
  std::mutex lock;
 public:
  void add(MatmulOp m) {
    std::lock_guard<std::mutex> l(lock);
    counts[m] += 1;
  }
  void print() {
    std::cerr << *this;
  }
  friend std::ostream& operator<<(std::ostream& os, EigenStats& s);
};

EIGEN_ALWAYS_INLINE std::ostream& operator<<(std::ostream& os, EigenStats& s) {
  std::lock_guard<std::mutex> l(s.lock);
  os << "Eigen STATS:" << std::endl;
  for (auto& kv : s.counts) {
    os << kv.second << '\t' << kv.first;
  }
  return os;
}

#else
// Mock class so that everything about stats gets optimized-out.
struct EigenStats {
  void add(MatmulOp m) {}
};
#endif


struct EigenStatsWrapper {
  static EigenStats * get() {
    static EigenStats stats;
#ifdef EIGEN_TENSOR_STATS
    static bool first_time = true;
    if (first_time) {
      first_time = false;
      atexit([]() { EigenStatsWrapper::get()->print(); });
      signal(SIGINT, [](int s) { exit(1); });
    }
#endif
    return &stats;
  }
};

}  // namespace internal
}  // namespace Eigen

#endif  // EIGEN_CXX11_TENSOR_TENSOR_STATS_H