Output per-node-type timings to proto in benchmark_model

Change: 151366175

1 files changed, 26 insertions, 25 deletions

diff --git a/tensorflow/core/util/stat_summarizer.cc b/tensorflow/core/util/stat_summarizer.cc
index a90ebab271..f8f071e0fa 100644
--- a/tensorflow/core/util/stat_summarizer.cc
+++ b/tensorflow/core/util/stat_summarizer.cc
@@ -266,41 +266,44 @@ void StatSummarizer::OrderNodesByMetric(
   }
 }
 
-std::string StatSummarizer::GetStatsByNodeType() const {
-  std::stringstream stream;
-
-  stream << "============================== Summary by node type "
-            "=============================="
-         << std::endl;
-
-  int64 accumulated_us = 0;
-  int64 accumulated_bytes = 0;
-  std::map<string, int64> node_type_map_count;
-  std::map<string, int64> node_type_map_time;
-  std::map<string, int64> node_type_map_memory;
-
-  int64 num_processed = 0;
-
-  LOG(INFO) << "Number of nodes executed: " << details_.size();
+void StatSummarizer::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,
+    int64* accumulated_us) const {
   for (const auto& det : details_) {
     const string node_name = det.first;
     const Detail& detail = det.second;
 
     int64 curr_time_val = detail.rel_end_us.avg();
-    accumulated_us += curr_time_val;
+    *accumulated_us += curr_time_val;
 
-    ++num_processed;
     int64 curr_memory_val = detail.mem_used.newest();
-    accumulated_bytes += curr_memory_val;
 
     const string& node_type = detail.type;
 
-    node_type_map_count[node_type] += 1;
-    node_type_map_time[node_type] += curr_time_val;
-    node_type_map_memory[node_type] += curr_memory_val;
+    (*node_type_map_count)[node_type] += 1;
+    (*node_type_map_time)[node_type] += curr_time_val;
+    (*node_type_map_memory)[node_type] += curr_memory_val;
   }
+}
+
+std::string StatSummarizer::GetStatsByNodeType() const {
+  std::stringstream stream;
+
+  stream << "============================== Summary by node type "
+            "=============================="
+         << std::endl;
+
+  LOG(INFO) << "Number of nodes executed: " << details_.size();
+
+  std::map<string, int64> node_type_map_count;
+  std::map<string, int64> node_type_map_time;
+  std::map<string, int64> node_type_map_memory;
+  int64 accumulated_us = 0;
 
-  LOG(INFO) << "Processed " << num_processed << " nodes";
+  ComputeStatsByType(&node_type_map_count, &node_type_map_time,
+                     &node_type_map_memory, &accumulated_us);
 
   // Sort them.
   std::priority_queue<std::pair<int64, std::pair<string, int64>>> timings;
@@ -318,7 +321,6 @@ std::string StatSummarizer::GetStatsByNodeType() const {
   InitField(stream, 10) << "[mem KB]";
   stream << std::endl;
 
-  float avg_total_time_ms = 0.0f;
   float cdf = 0.0f;
   while (!timings.empty()) {
     auto entry = timings.top();
@@ -330,7 +332,6 @@ std::string StatSummarizer::GetStatsByNodeType() const {
     const int64 node_type_total_us = entry.first;
     const float time_per_run_ms = node_type_total_us / 1000.0f;
 
-    avg_total_time_ms += time_per_run_ms;
     const float percentage =
         ((entry.first / static_cast<float>(accumulated_us)) * 100.0f);
     cdf += percentage;