Some profiler fixes and cleanup.

PiperOrigin-RevId: 168069346

8 files changed, 104 insertions, 43 deletions

diff --git a/tensorflow/core/profiler/README.md b/tensorflow/core/profiler/README.md
index 2ddac0a79b..7855799417 100644
--- a/tensorflow/core/profiler/README.md
+++ b/tensorflow/core/profiler/README.md
@@ -216,7 +216,7 @@ seq2seq_attention_model.py:363:build_graph:self._add_train_o..., cpu: 1.28sec, a
 
 ```shell
 # The following example generates a timeline.
-tfprof> graph -step 0 -max_depth 100000 -output timeline:outfile=<filename>
+tfprof> graph -step -1 -max_depth 100000 -output timeline:outfile=<filename>
 
 generating trace file.
 
diff --git a/tensorflow/core/profiler/g3doc/command_line.md b/tensorflow/core/profiler/g3doc/command_line.md
index fb4207c784..d41ac7290d 100644
--- a/tensorflow/core/profiler/g3doc/command_line.md
+++ b/tensorflow/core/profiler/g3doc/command_line.md
@@ -14,7 +14,12 @@
 
 ### Command Line Inputs
 
-tfprof command line tool uses the following inputs:
+tfprof command line tool uses the following input:
+
+<b>--profile_path:</b> A ProfileProto binary proto file.
+See QuickStart on generating the file.
+
+<b>THE OLD WAY BELOW IS DEPRECATED:</b>
 
 <b>--graph_path:</b> GraphDef proto file (required). Used to build in-memory
 data structure of the model. For example, graph.pbtxt written by tf.Supervisor
diff --git a/tensorflow/core/profiler/internal/print_model_analysis.cc b/tensorflow/core/profiler/internal/print_model_analysis.cc
index ddf3c7f1f2..575ae182ee 100644
--- a/tensorflow/core/profiler/internal/print_model_analysis.cc
+++ b/tensorflow/core/profiler/internal/print_model_analysis.cc
@@ -84,7 +84,6 @@ string RunProfile(const string& command, const string& options,
 }  // namespace
 
 bool NewProfiler(const string* graph, const string* op_log) {
-  CHECK(!tf_stat) << "Currently only 1 living tfprof profiler is allowed";
   CHECK(graph) << "graph mustn't be null";
   std::unique_ptr<GraphDef> graph_ptr(new GraphDef());
   if (!graph_ptr->ParseFromString(*graph)) {
diff --git a/tensorflow/core/profiler/internal/tfprof_node.h b/tensorflow/core/profiler/internal/tfprof_node.h
index 55d53f3923..95d199e5b9 100644
--- a/tensorflow/core/profiler/internal/tfprof_node.h
+++ b/tensorflow/core/profiler/internal/tfprof_node.h
@@ -175,22 +175,22 @@ class ExecStep {
   std::map<int32, std::pair<int64, uint64>> output_memory_;
 };
 
-#define GRAPH_NODE_BYTES(type)                                \
-  do {                                                        \
-    if (execs_.empty()) {                                     \
-      return 0;                                               \
-    }                                                         \
-    if (step >= 0) {                                          \
-      auto exec = execs_.find(step);                          \
-      CHECK(exec != execs_.end()) << "unknown step " << step; \
-      return exec->second.type##_bytes();                     \
-    }                                                         \
-                                                              \
-    int64 bytes = 0;                                          \
-    for (const auto& exec : execs_) {                         \
-      bytes += exec.second.type##_bytes();                    \
-    }                                                         \
-    return bytes / execs_.size();                             \
+#define GRAPH_NODE_BYTES(type)             \
+  do {                                     \
+    if (execs_.empty()) {                  \
+      return 0;                            \
+    }                                      \
+    if (step >= 0) {                       \
+      auto exec = execs_.find(step);       \
+      if (exec == execs_.end()) return 0;  \
+      return exec->second.type##_bytes();  \
+    }                                      \
+                                           \
+    int64 bytes = 0;                       \
+    for (const auto& exec : execs_) {      \
+      bytes += exec.second.type##_bytes(); \
+    }                                      \
+    return bytes / execs_.size();          \
   } while (0)
 
 class TFGraphNode {
@@ -372,7 +372,9 @@ class TFGraphNode {
     }
     if (step >= 0) {
       auto exec = execs_.find(step);
-      CHECK(exec != execs_.end());
+      if (exec == execs_.end()) {
+        return 0;
+      }
       return exec->second.run_count();
     }
     int64 total_run_count = 0;
@@ -390,7 +392,9 @@ class TFGraphNode {
     }
     if (step >= 0) {
       auto exec = execs_.find(step);
-      CHECK(exec != execs_.end());
+      if (exec == execs_.end()) {
+        return 0;
+      }
       return exec->second.exec_micros();
     }
 
@@ -410,7 +414,9 @@ class TFGraphNode {
     }
     if (step >= 0) {
       auto exec = execs_.find(step);
-      CHECK(exec != execs_.end());
+      if (exec == execs_.end()) {
+        return 0;
+      }
       return exec->second.accelerator_exec_micros();
     }
 
@@ -430,7 +436,9 @@ class TFGraphNode {
     }
     if (step >= 0) {
       auto exec = execs_.find(step);
-      CHECK(exec != execs_.end());
+      if (exec == execs_.end()) {
+        return 0;
+      }
       return exec->second.cpu_exec_micros();
     }
 
@@ -448,20 +456,26 @@ class TFGraphNode {
 
   int64 all_start_micros(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.all_start_micros();
   }
 
   int64 latest_end_micros(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.latest_end_micros();
   }
 
   const std::map<string, std::vector<std::pair<int64, int64>>>& op_execs(
       int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return empty_op_execs_;
+    }
     return exec->second.op_execs();
   }
 
@@ -469,33 +483,45 @@ class TFGraphNode {
 
   int64 accelerator_temp_bytes(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.accelerator_temp_bytes();
   }
   int64 host_temp_bytes(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.host_temp_bytes();
   }
   int64 accelerator_persistent_bytes(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.accelerator_persistent_bytes();
   }
   int64 host_persistent_bytes(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.host_persistent_bytes();
   }
   const std::map<int32, std::pair<int64, uint64>>& output_memory(
       int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return empty_output_memory_;
+    }
     return exec->second.output_memory();
   }
   int64 allocator_bytes_in_use(int64 step) const {
     auto exec = execs_.find(step);
-    CHECK(exec != execs_.end()) << "unknown step " << step;
+    if (exec == execs_.end()) {
+      return 0;
+    }
     return exec->second.allocator_bytes_in_use();
   }
 
@@ -566,6 +592,9 @@ class TFGraphNode {
   std::set<string> op_types_;
 
   std::map<int64, ExecStep> execs_;
+
+  std::map<int32, std::pair<int64, uint64>> empty_output_memory_;
+  std::map<string, std::vector<std::pair<int64, int64>>> empty_op_execs_;
 };
 
 class TFMultiGraphNode {
diff --git a/tensorflow/core/profiler/internal/tfprof_stats.cc b/tensorflow/core/profiler/internal/tfprof_stats.cc
index 5b54958344..eb84bada13 100644
--- a/tensorflow/core/profiler/internal/tfprof_stats.cc
+++ b/tensorflow/core/profiler/internal/tfprof_stats.cc
@@ -88,6 +88,9 @@ TFStats::TFStats(const string& filename,
         node_pb.second.name(), std::move(node)));
   }
   has_code_traces_ = profile.has_trace();
+  for (int64 s : profile.steps()) {
+    steps_.insert(s);
+  }
 }
 
 void TFStats::BuildView(const string& cmd) {
@@ -136,6 +139,14 @@ const GraphNodeProto& TFStats::ShowGraphNode(const string& cmd,
   if (cmd == kCmds[0]) {
     return scope_view_->Show(opts);
   } else if (cmd == kCmds[1]) {
+    if (opts.step < 0 && opts.output_type == kOutput[0]) {
+      for (int64 step : steps_) {
+        Options nopts = opts;
+        nopts.step = step;
+        graph_view_->Show(nopts);
+      }
+      return empty_graph_node_;
+    }
     return graph_view_->Show(opts);
   } else {
     fprintf(stderr, "Unknown command: %s\n", cmd.c_str());
@@ -148,7 +159,11 @@ const MultiGraphNodeProto& TFStats::ShowMultiGraphNode(
   if (!Validate(opts)) {
     return empty_multi_graph_node_;
   }
-  if (cmd == kCmds[2] && has_code_traces()) {
+  if (cmd == kCmds[2]) {
+    if (!has_code_traces()) {
+      fprintf(stderr, "No code trace information\n");
+      return empty_multi_graph_node_;
+    }
     return code_view_->Show(opts);
   } else if (cmd == kCmds[3]) {
     return op_view_->Show(opts);
@@ -212,7 +227,9 @@ void TFStats::AddOpLogProto(std::unique_ptr<OpLogProto> op_log) {
     }
     if (entry.has_code_def()) {
       has_code_traces_ = true;
-      node->second->AddCode(entry.code_def());
+      if (node->second->code().traces_size() == 0) {
+        node->second->AddCode(entry.code_def());
+      }
     }
   }
 }
@@ -258,9 +275,11 @@ void TFStats::WriteProfile(const string& filename) {
     }
     (*profile.mutable_nodes())[it->second->id()].MergeFrom(
         it->second->ToProto(nodes_map_));
-    if (it->second->code().traces_size() > 0) {
-      profile.set_has_trace(true);
-    }
+  }
+
+  profile.set_has_trace(has_code_traces_);
+  for (int64 s : steps_) {
+    profile.add_steps(s);
   }
   Status s =
       WriteStringToFile(Env::Default(), filename, profile.SerializeAsString());
@@ -271,7 +290,12 @@ void TFStats::WriteProfile(const string& filename) {
 
 bool TFStats::Validate(const Options& opts) const {
   if (opts.step >= 0 && steps_.find(opts.step) == steps_.end()) {
-    fprintf(stderr, "Options -step=%lld not found\n", opts.step);
+    fprintf(stderr,
+            "Options -step=%lld not found.\nAvailable steps: ", opts.step);
+    for (int64 s : steps_) {
+      fprintf(stderr, "%lld ", s);
+    }
+    fprintf(stderr, "\n");
     return false;
   }
   return true;
diff --git a/tensorflow/core/profiler/internal/tfprof_timeline.cc b/tensorflow/core/profiler/internal/tfprof_timeline.cc
index f3934860d9..1732574cc4 100644
--- a/tensorflow/core/profiler/internal/tfprof_timeline.cc
+++ b/tensorflow/core/profiler/internal/tfprof_timeline.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 
 namespace tensorflow {
@@ -303,11 +304,12 @@ void Timeline::GenerateCodeTimeline(const CodeNode* node) {
 }
 
 void Timeline::OutputTimeline() {
+  string outfile = strings::Printf("%s_%lld", outfile_.c_str(), step());
   Status s =
-      WriteStringToFile(Env::Default(), outfile_, chrome_formatter_.Format());
+      WriteStringToFile(Env::Default(), outfile, chrome_formatter_.Format());
   if (!s.ok()) {
     fprintf(stderr, "Failed to write timeline file: %s\nError: %s\n",
-            outfile_.c_str(), s.ToString().c_str());
+            outfile.c_str(), s.ToString().c_str());
     return;
   }
   fprintf(stdout, "\n******************************************************\n");
@@ -315,7 +317,7 @@ void Timeline::OutputTimeline() {
           "Timeline file is written to %s.\n"
           "Open a Chrome browser, enter URL chrome://tracing and "
           "load the timeline file.",
-          outfile_.c_str());
+          outfile.c_str());
   fprintf(stdout, "\n******************************************************\n");
   fflush(stdout);
 }
diff --git a/tensorflow/core/profiler/internal/tfprof_timeline_test.cc b/tensorflow/core/profiler/internal/tfprof_timeline_test.cc
index 5dd440e9a2..babae395ba 100644
--- a/tensorflow/core/profiler/internal/tfprof_timeline_test.cc
+++ b/tensorflow/core/profiler/internal/tfprof_timeline_test.cc
@@ -70,7 +70,7 @@ TEST_F(TFProfTimelineTest, GraphView) {
   tf_stats_->ShowGraphNode("graph", opts);
 
   string dump_str;
-  TF_CHECK_OK(ReadFileToString(Env::Default(), dump_file, &dump_str));
+  TF_CHECK_OK(ReadFileToString(Env::Default(), dump_file + "_0", &dump_str));
   EXPECT_EQ(1754536562981488144ull, Hash64(dump_str));
 }
 
@@ -84,7 +84,7 @@ TEST_F(TFProfTimelineTest, ScopeView) {
   tf_stats_->ShowGraphNode("scope", opts);
 
   string dump_str;
-  TF_CHECK_OK(ReadFileToString(Env::Default(), dump_file, &dump_str));
+  TF_CHECK_OK(ReadFileToString(Env::Default(), dump_file + "_0", &dump_str));
   EXPECT_EQ(17545174915963890413ull, Hash64(dump_str));
 }
 
diff --git a/tensorflow/core/profiler/tfprof_log.proto b/tensorflow/core/profiler/tfprof_log.proto
index 1ce5a5eecf..ae571e2540 100644
--- a/tensorflow/core/profiler/tfprof_log.proto
+++ b/tensorflow/core/profiler/tfprof_log.proto
@@ -42,6 +42,8 @@ message ProfileProto {
   map<int64, ProfileNode> nodes = 1;
   // Whether or not has code traces.
   bool has_trace = 2;
+  // Traced steps.
+  repeated int64 steps = 3;
 }
 
 message ProfileNode {