text: #include "caffe2/core/prof_dag_counters.h"
  #include "caffe2/utils/string_utils.h"
  
  #include <ostream>
  #include <sstream>
  
  #include <c10/util/irange.h>
  
  namespace caffe2 {
  
  ProfDAGCounters::ProfDAGCounters(const std::shared_ptr<const NetDef>& net_def) {
    report_.net_name_ = net_def->name();
    report_.num_runs_ = 0;
    auto num_ops = net_def->op_size();
    report_.op_types_.reserve(num_ops);
    report_.op_extra_info_.reserve(num_ops);
  
    for (auto op_id = 0; op_id < num_ops; ++op_id) {
      const auto& op = net_def->op(op_id);
      if (op.engine() == "") {
        report_.op_types_.push_back(op.type());
      } else {
        report_.op_types_.push_back(op.type() + "(" + op.engine() + ")");
      }
      vector<std::string> op_extra_info;
      if (op.has_device_option() && op.device_option().extra_info_size() > 0) {
        for (auto i = 0; i < op.device_option().extra_info_size(); ++i) {
          std::string extra_info_str = op.device_option().extra_info(i);
          op_extra_info.push_back(extra_info_str);
        }
      }
      report_.op_extra_info_.push_back(op_extra_info);
    }
    report_.time_per_op_total_.resize(num_ops);
  }
  
  void ProfDAGCounters::ReportRunStart() {
    report_.num_runs_ += 1;
    timer_.Start();
    auto num_ops = report_.op_types_.size();
    op_start_times_run_.clear();
    op_start_times_run_.resize(num_ops, -1.0);
    op_end_times_run_.clear();
    op_end_times_run_.resize(num_ops, -1.0);
    op_async_end_times_run_.clear();
    op_async_end_times_run_.resize(num_ops, -1.0);
  }
  
  void ProfDAGCounters::AddPerOpStartTime(size_t op_id) {
    if (report_.num_runs_ <= 1) {
      return;
    }
  
    CAFFE_ENFORCE(op_id < op_start_times_run_.size());
    op_start_times_run_[op_id] = timer_.MilliSeconds();
  }
  
  void ProfDAGCounters::AddPerOpEndTime(size_t op_id) {
    if (report_.num_runs_ <= 1) {
      return;
    }
  
    CAFFE_ENFORCE(op_id < op_end_times_run_.size());
    op_end_times_run_[op_id] = timer_.MilliSeconds();
  }
  
  void ProfDAGCounters::AddPerOpAsyncEndTime(size_t op_id) {
    if (report_.num_runs_ <= 1) {
      return;
    }
  
    CAFFE_ENFORCE(op_id < op_async_end_times_run_.size());
    op_async_end_times_run_[op_id] = timer_.MilliSeconds();
  }
  
  void ProfDAGCounters::ReportRunEnd() {
    if (report_.num_runs_ <= 1) {
      return;
    }
  
    auto runtime = timer_.MilliSeconds();
  
    CaffeMap<std::string, float> cum_per_type_time_run;
    CaffeMap<std::string, float> cum_per_type_invocations_run;
    std::vector<float> per_op_time_run(report_.op_types_.size(), 0.0);
    for (const auto op_id : c10::irange(report_.op_types_.size())) {
      // check that we have valid times, otherwise return;
      // times might not be valid if network execution ended prematurely
      // because of operator errors
      if (op_start_times_run_[op_id] < 0.0) {
        return;
      }
  
      float op_time = 0.0;
      if (op_async_end_times_run_[op_id] > 0.0) {
        op_time = op_async_end_times_run_[op_id] - op_start_times_run_[op_id];
      } else {
        if (op_end_times_run_[op_id] < 0.0) {
          return;
        }
        op_time = op_end_times_run_[op_id] - op_start_times_run_[op_id];
      }
  
      per_op_time_run[op_id] = op_time;
  
      const string& op_type = report_.op_types_[op_id];
      cum_per_type_time_run[op_type] += op_time;
      cum_per_type_invocations_run[op_type] += 1;
    }
  
    // all operator times are valid, update report stats
    report_.runtime_stats_ += ProfDAGStats(runtime);
  
    for (const auto op_id : c10::irange(report_.op_types_.size())) {
      report_.time_per_op_total_[op_id] += ProfDAGStats(per_op_time_run[op_id]);
    }
  
    for (const auto& kv : cum_per_type_time_run) {
      report_.time_per_op_type_total_[kv.first] += ProfDAGStats(kv.second);
      report_.times_per_run_per_type_total_[kv.first] +=
          ProfDAGStats(cum_per_type_invocations_run[kv.first]);
    }
  }
  
  ProfDAGReport ProfDAGCounters::GetReport() const {
    return report_;
  }
  
  bool ProfDAGReport::hasStats() const {
    return runtime_stats_.cnt() > 0;
  }
  
  ProfDAGProto ProfDAGReport::statsProto(
      const std::string& name,
      const ProfDAGStats& stats,
      const std::vector<std::string>& op_extra_info) const {
    ProfDAGProto stats_proto;
    const auto& moments = stats.computeMoments();
    stats_proto.set_mean(moments.first);
    stats_proto.set_stddev(moments.second);
    stats_proto.set_name(name);
    for (auto& extra_info : op_extra_info) {
      stats_proto.add_extra_info(extra_info);
    }
    return stats_proto;
  }
  
  ProfDAGProtos ProfDAGReport::GetOperatorStats() const {
    ProfDAGProtos prof_dag_protos;
    prof_dag_protos.set_net_name(net_name_);
    if (hasStats()) {
      for (auto& item : time_per_op_type_total_) {
        auto buf = prof_dag_protos.add_stats();
        buf->CopyFrom(statsProto(item.first, item.second, vector<std::string>()));
      }
    }
    return prof_dag_protos;
  }
  
  ProfDAGProtos ProfDAGReport::GetPerOperatorCost() const {
    ProfDAGProtos prof_dag_protos;
    prof_dag_protos.set_net_name(net_name_);
    if (hasStats()) {
      for (const auto op_id : c10::irange(op_types_.size())) {
        const string& op_type = op_types_[op_id];
        auto buf = prof_dag_protos.add_stats();
        std::string op_output_name =
            net_name_ + "___" + to_string(op_id) + "___" + op_type;
        buf->CopyFrom(statsProto(
            op_output_name, time_per_op_total_[op_id], op_extra_info_[op_id]));
      }
    }
    return prof_dag_protos;
  }
  
  void ProfDAGReport::PrintStats() {
    if (!hasStats()) {
      LOG(INFO) << "Insufficient number of runs";
      return;
    }
  
    std::ostringstream debug_out;
    debug_out << "Measured operators over " << runtime_stats_.cnt()
              << " net runs (" << net_name_ << "), #ops: " << op_types_.size()
              << std::endl;
  
    debug_out << "Mean time in operator type per run (stddev):" << std::endl;
    for (const auto& item : time_per_op_type_total_) {
      const auto& moments = item.second.computeMoments();
      const auto& times_moments =
          times_per_run_per_type_total_[item.first].computeMoments();
      debug_out << std::setw(10) << std::setfill(' ') << moments.first
                << " ms/run (" << std::setw(10) << std::setfill(' ')
                << moments.second << " ms/run) "
                << " Op count per run: " << times_moments.first << "  "
                << item.first << std::endl;
    }
    const auto& runtime_moments = runtime_stats_.computeMoments();
    debug_out << net_name_ << " runtime: " << runtime_moments.first << " ms ("
              << runtime_moments.second << " ms)" << std::endl;
  
    LOG(INFO) << debug_out.str();
  }
  
  ProfDAGReport& ProfDAGReport::operator+=(const ProfDAGReport& rhs) {
    // Verify nets are compatible for addition
    CAFFE_ENFORCE_EQ(
        net_name_, rhs.net_name_, "Incompatible nets to add counters");
    CAFFE_ENFORCE_EQ(
        op_types_.size(),
        rhs.op_types_.size(),
        "Incompatible nets to add counters");
    for (const auto idx : c10::irange(op_types_.size())) {
      CAFFE_ENFORCE_EQ(
          op_types_[idx],
          rhs.op_types_[idx],
          "Incompatible nets to add counters");
    }
  
    if (!rhs.hasStats()) {
      // rhs does not have valid profiling results, do nothing
      return *this;
    } else if (!hasStats()) {
      // "this" does not have valid profiling results, but rhs does. copy rhs
      time_per_op_total_ = rhs.time_per_op_total_;
      time_per_op_type_total_ = rhs.time_per_op_type_total_;
      times_per_run_per_type_total_ = rhs.times_per_run_per_type_total_;
      runtime_stats_ = rhs.runtime_stats_;
      num_runs_ = rhs.num_runs_;
      return *this;
    }
  
    // Do the addition
    for (const auto idx : c10::irange(time_per_op_total_.size())) {
      time_per_op_total_[idx] += rhs.time_per_op_total_.at(idx);
    }
    for (auto& item : time_per_op_type_total_) {
      item.second += rhs.time_per_op_type_total_.at(item.first);
    }
    for (auto& item : times_per_run_per_type_total_) {
      item.second += rhs.times_per_run_per_type_total_.at(item.first);
    }
    runtime_stats_ += rhs.runtime_stats_;
    num_runs_ += rhs.num_runs_;
  
    return *this;
  }
  
  } // namespace caffe2
  