Refactor StatSummarizer extract common functionality without proto dependencies.

PiperOrigin-RevId: 197816405

1 files changed, 36 insertions, 152 deletions

diff --git a/tensorflow/core/util/stat_summarizer.h b/tensorflow/core/util/stat_summarizer.h
index 79fa63723e..39cd948525 100644
--- a/tensorflow/core/util/stat_summarizer.h
+++ b/tensorflow/core/util/stat_summarizer.h
@@ -13,20 +13,23 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_UTIL_STAT_SUMMARIZER_H_
-#define TENSORFLOW_UTIL_STAT_SUMMARIZER_H_
+#ifndef TENSORFLOW_CORE_UTIL_STAT_SUMMARIZER_H_
+#define TENSORFLOW_CORE_UTIL_STAT_SUMMARIZER_H_
 
 #include <stdlib.h>
 
 #include <cmath>
 #include <limits>
 #include <map>
+#include <memory>
 #include <sstream>
 #include <string>
 
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/stat_summarizer_options.h"
+#include "tensorflow/core/util/stats_calculator.h"
 
 namespace tensorflow {
 
@@ -34,103 +37,6 @@ class GraphDef;
 class StepStats;
 class NodeExecStats;
 
-template <typename ValueType, typename HighPrecisionValueType = double>
-class Stat {
- public:
-  void UpdateStat(ValueType v) {
-    if (count_ == 0) {
-      first_ = v;
-    }
-
-    newest_ = v;
-    max_ = std::max(v, max_);
-    min_ = std::min(v, min_);
-    ++count_;
-    sum_ += v;
-    squared_sum_ += static_cast<HighPrecisionValueType>(v) * v;
-  }
-
-  void Reset() { new (this) Stat<ValueType, HighPrecisionValueType>(); }
-
-  bool empty() const { return count_ == 0; }
-
-  ValueType first() const { return first_; }
-
-  ValueType newest() const { return newest_; }
-
-  ValueType max() const { return max_; }
-
-  ValueType min() const { return min_; }
-
-  int64 count() const { return count_; }
-
-  ValueType sum() const { return sum_; }
-
-  HighPrecisionValueType squared_sum() const { return squared_sum_; }
-
-  bool all_same() const { return (count_ == 0 || min_ == max_); }
-
-  HighPrecisionValueType avg() const {
-    return empty() ? std::numeric_limits<ValueType>::quiet_NaN()
-                   : static_cast<HighPrecisionValueType>(sum_) / count_;
-  }
-
-  ValueType std_deviation() const {
-    return all_same() ? 0 : sqrt(squared_sum_ / count_ - avg() * avg());
-  }
-
-  void OutputToStream(std::ostream* stream) const {
-    if (empty()) {
-      *stream << "count=0";
-    } else if (all_same()) {
-      *stream << "count=" << count_ << " curr=" << newest_;
-      if (count_ > 1) *stream << "(all same)";
-    } else {
-      *stream << "count=" << count_ << " first=" << first_
-              << " curr=" << newest_ << " min=" << min_ << " max=" << max_
-              << " avg=" << avg() << " std=" << std_deviation();
-    }
-  }
-
-  friend std::ostream& operator<<(std::ostream& stream,
-                                  const Stat<ValueType>& stat) {
-    stat.OutputToStream(&stream);
-    return stream;
-  }
-
- private:
-  ValueType first_ = 0;
-  ValueType newest_ = 0;
-  ValueType max_ = std::numeric_limits<ValueType>::min();
-  ValueType min_ = std::numeric_limits<ValueType>::max();
-  int64 count_ = 0;
-  ValueType sum_ = 0;
-  HighPrecisionValueType squared_sum_ = 0;
-};
-
-// Used to control the output of the statistics summarizer;
-class StatSummarizerOptions {
- public:
-  StatSummarizerOptions()
-      : show_run_order(true),
-        run_order_limit(0),
-        show_time(true),
-        time_limit(10),
-        show_memory(true),
-        memory_limit(10),
-        show_type(true),
-        show_summary(true) {}
-
-  bool show_run_order;
-  int run_order_limit;
-  bool show_time;
-  int time_limit;
-  bool show_memory;
-  int memory_limit;
-  bool show_type;
-  bool show_summary;
-};
-
 // A StatSummarizer assists in performance analysis of Graph executions.
 //
 // It summarizes time spent executing (on GPU/CPU), memory used etc. across
@@ -140,14 +46,6 @@ class StatSummarizerOptions {
 // See tensorflow/tools/benchmark/benchmark_model.cc for an example usage.
 class StatSummarizer {
  public:
-  enum SortingMetric {
-    BY_NAME,
-    BY_RUN_ORDER,
-    BY_TIME,
-    BY_MEMORY,
-    BY_TYPE,
-  };
-
   explicit StatSummarizer(const StatSummarizerOptions& options);
 
   // Deprecated: Use StatSummarizer(const StatSummarizerOptions&) instead. The
@@ -161,65 +59,51 @@ class StatSummarizer {
 
   // Returns a string detailing the accumulated runtime stats in a tab-separated
   // format which can be pasted into a spreadsheet for further analysis.
-  std::string GetOutputString() const;
+  std::string GetOutputString() const {
+    return stats_calculator_->GetOutputString();
+  }
 
-  std::string ShortSummary() const;
+  std::string ShortSummary() const {
+    return stats_calculator_->GetShortSummary();
+  }
 
   // Prints the string returned by GetOutputString().
-  void PrintStepStats() const;
+  void PrintStepStats() const { stats_calculator_->PrintStepStats(); }
 
   // Prints the output tensor sizes and types for each node.
   void PrintOutputs() const;
 
-  void ComputeStatsByType(std::map<string, int64>* node_type_map_count,
-                          std::map<string, int64>* node_type_map_time,
-                          std::map<string, int64>* node_type_map_memory,
-                          std::map<string, int64>* node_type_map_times_called,
-                          int64* accumulated_us) const;
+  void ComputeStatsByType(
+      std::map<std::string, int64_t>* node_type_map_count,
+      std::map<std::string, int64_t>* node_type_map_time,
+      std::map<std::string, int64_t>* node_type_map_memory,
+      std::map<std::string, int64_t>* node_type_map_times_called,
+      int64_t* accumulated_us) const {
+    stats_calculator_->ComputeStatsByType(
+        node_type_map_count, node_type_map_time, node_type_map_memory,
+        node_type_map_times_called, accumulated_us);
+  }
 
-  std::string GetStatsByNodeType() const;
+  std::string GetStatsByNodeType() const {
+    return stats_calculator_->GetStatsByNodeType();
+  }
 
   std::string GetStatsByMetric(const string& title,
-                               SortingMetric sorting_metric,
-                               int num_stats) const;
-
-  void Reset();
+                               StatsCalculator::SortingMetric sorting_metric,
+                               int num_stats) const {
+    return stats_calculator_->GetStatsByMetric(title, sorting_metric,
+                                               num_stats);
+  }
 
-  // Returns number of runs.
-  int num_runs() const { return static_cast<int>(run_total_us_.count()); }
+ private:
+  void Validate(const std::vector<TensorDescription>* outputs,
+                const NodeExecStats& ns) const;
 
-  // Returns stats of total microseconds spent by all nodes in each run.
-  const Stat<int64>& run_total_us() const { return run_total_us_; }
+  std::map<std::string, std::vector<TensorDescription> > outputs_;
 
- private:
-  struct Detail {
-    string name;
-    string type;
-    int64 run_order;
-    Stat<int64> start_us;
-    Stat<int64> rel_end_us;
-    Stat<int64> mem_used;
-    std::vector<TensorDescription> outputs;
-    int64 times_called;
-  };
-
-  void Validate(const Detail* detail, const NodeExecStats& ns) const;
-
-  void OrderNodesByMetric(SortingMetric sorting_metric,
-                          std::vector<const Detail*>* details) const;
-
-  std::string HeaderString(const string& title) const;
-  std::string ColumnString(const Detail& detail,
-                           const int64 cumulative_stat_on_node,
-                           const Stat<int64>& stat) const;
-
-  Stat<int64> run_total_us_;
-  Stat<int64> memory_;
-
-  std::map<std::string, Detail> details_;
-  StatSummarizerOptions options_;
+  std::unique_ptr<StatsCalculator> stats_calculator_;
 };
 
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_UTIL_STAT_SUMMARIZER_H_
+#endif  // TENSORFLOW_CORE_UTIL_STAT_SUMMARIZER_H_