profiler C++ API.

PiperOrigin-RevId: 181397308

tensorflow/BUILD

@@ -782,6 +782,7 @@ tf_cc_shared_object(
         "//tensorflow/cc:cc_ops",
         "//tensorflow/cc:client_session",
         "//tensorflow/cc:scope",
+        "//tensorflow/cc/profiler",
         "//tensorflow/core:tensorflow",
     ],
 )

tensorflow/cc/profiler/BUILD

@@ -0,0 +1,36 @@
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])  # Apache 2.0
+
+load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
+
+tf_cuda_cc_test(
+    name = "profiler_test",
+    srcs = ["profiler_test.cc"],
+    deps = [
+        ":profiler",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+cc_library(
+    name = "profiler",
+    srcs = ["profiler.cc"],
+    hdrs = ["profiler.h"],
+    deps = [
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
+        "//tensorflow/core/profiler/internal:tfprof_stats",
+    ],
+)

tensorflow/cc/profiler/profiler.cc

@@ -0,0 +1,57 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/cc/profiler/profiler.h"
+
+namespace tensorflow {
+namespace tfprof {
+
+Profiler::Profiler(const GraphDef& graph) {
+  std::unique_ptr<GraphDef> graph_ptr(new GraphDef());
+  *graph_ptr = graph;
+  stats_.reset(new TFStats(std::move(graph_ptr), nullptr, nullptr, nullptr));
+}
+
+void Profiler::AddStep(int64 step, const RunMetadata& run_meta) {
+  std::unique_ptr<RunMetadata> run_meta_ptr(new RunMetadata());
+  *run_meta_ptr = run_meta;
+  stats_->AddRunMeta(step, std::move(run_meta_ptr));
+}
+
+GraphNodeProto Profiler::ProfileGraph(const Options& options) {
+  stats_->BuildView(kCmds[1]);
+  return stats_->ShowGraphNode(kCmds[1], options);
+}
+
+GraphNodeProto Profiler::ProfileNameScope(const Options& options) {
+  stats_->BuildView(kCmds[0]);
+  return stats_->ShowGraphNode(kCmds[0], options);
+}
+
+MultiGraphNodeProto Profiler::ProfileOperations(const Options& options) {
+  stats_->BuildView(kCmds[3]);
+  return stats_->ShowMultiGraphNode(kCmds[3], options);
+}
+
+Status Profiler::SerializeToString(string* content) {
+  if (!content) {
+    return Status(error::Code::INVALID_ARGUMENT,
+                  "Cannot use null string pointer for SerializeToString.");
+  }
+  stats_->SerializeToString(content);
+  return Status::OK();
+}
+
+}  // namespace tfprof
+}  // namespace tensorflow

tensorflow/cc/profiler/profiler.h

@@ -0,0 +1,97 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef THIRD_PARTY_TENSORFLOW_CC_PROFILER_PROFILER_H_
+#define THIRD_PARTY_TENSORFLOW_CC_PROFILER_PROFILER_H_
+
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/profiler/internal/tfprof_stats.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
+#include "tensorflow/core/profiler/tfprof_output.pb.h"
+
+namespace tensorflow {
+namespace tfprof {
+
+/// @addtogroup core
+/// @{
+
+/// A `Profiler` object lets the caller profile the execution of a graph.
+///
+/// Example:
+///     // First build a graph and run tracing.
+///     Scope root = Scope::NewRootScope();
+///     auto a = Placeholder(root, DT_INT32);
+///     auto c = Add(root, a, {41});
+///
+///     ClientSession session(root);
+///     std::vector<Tensor> outputs;
+///     RunOptions run_options;
+///     run_options.set_trace_level(RunOptions::FULL_TRACE);
+///     RunMetadata run_meta;
+///     Status s = session.Run(run_options, { {a, {1}} }, {c}, &outputs,
+///                            &run_meta);
+///     if (!s.ok()) { ... }
+///
+///     // Then create profiler to do profiling.
+///     GraphDef graph;
+///     root.ToGraphDef(&graph);
+///     Profiler profiler(graph);
+///     profiler.AddStep(0, run_meta);
+///     Options opts = ...  // TODO(xpan): Support option building API.
+///     MultiGraphNodeProto r = profiler.ProfileOperations(opts);
+///
+class Profiler {
+ public:
+  /// `graph` is the model's GraphDef.
+  Profiler(const GraphDef& graph);
+
+  /// Adds tracing information `run_meta` to profiler. A `run_meta` is
+  /// generated by a TensorFlow session run call. `step` is the key
+  /// to the `run_meta`. When calling ProfileXXX methods, caller can specify
+  /// `step` in `options` to seletively profile the corresponding `run_meta`.
+  /// Multiple different `run_meta` can be keyed by the same `step` in order
+  /// to group them together.
+  void AddStep(int64 step, const RunMetadata& run_meta);
+
+  /// Profiles the model by organizing nodes in graph structure.
+  /// Each node is an op and the nodes are contected by the op inputs/outputs.
+  GraphNodeProto ProfileGraph(const Options& options);
+
+  /// Profiles the model by organizing nodes in name scope structure.
+  /// Each node is an op, and nodes are organized by the ops' name
+  /// scope, similar to a filesystem tree.
+  /// E.g. /foo is the root of operation /foo/matmul_1 and foo/conv_2.
+  GraphNodeProto ProfileNameScope(const Options& options);
+
+  /// Profiles the model by organizing nodes by operation types.
+  /// Each node is an operation type (e.g. Conv2D or MatMul), containing all
+  /// ops belonging to that type in the model.
+  MultiGraphNodeProto ProfileOperations(const Options& options);
+
+  /// Serialize the profile content (ProfileProto) into a binary string,
+  /// User can write the string to file for offline analysis by
+  /// tfprof command-line tools or graphical user interface.
+  Status SerializeToString(string* content);
+
+ private:
+  std::unique_ptr<TFStats> stats_;
+};
+/// @}
+
+}  // namespace tfprof
+}  // namespace tensorflow
+
+#endif  // THIRD_PARTY_TENSORFLOW_CC_PROFILER_PROFILER_H_

tensorflow/cc/profiler/profiler_test.cc

@@ -0,0 +1,177 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/platform/test.h"
+
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/cc/profiler/profiler.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/graph/default_device.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/public/session.h"
+
+namespace tensorflow {
+namespace tfprof {
+
+class ProfilerTest : public ::testing::Test {
+ protected:
+  ProfilerTest() {}
+};
+
+GraphDef CreateGraphDef() {
+  Scope root = Scope::NewRootScope();
+
+  auto a = ops::Const<float>(root, {{3, 2}, {-1, 0}});
+
+  auto x = ops::Const(root.WithOpName("x"), {{1.f}, {1.f}});
+
+  auto y = ops::MatMul(root.WithOpName("y"), a, x);
+
+  auto y2 = ops::Square(root, y);
+
+  auto y2_sum = ops::Sum(root, y2, 0);
+
+  auto y_norm = ops::Sqrt(root, y2_sum);
+
+  auto y_div = ops::Div(root.WithOpName("y_normalized"), y, y_norm);
+
+  GraphDef def;
+  TF_CHECK_OK(root.ToGraphDef(&def));
+
+  return def;
+}
+
+Options Default() {
+  Options opts(1000,       /* max_depth */
+               0,          /* min_bytes */
+               0,          /* min_peak_bytes */
+               0,          /* min_residual_bytes */
+               0,          /* min_output_bytes */
+               0,          /* min_micros */
+               0,          /* min_accelerator_micros */
+               0,          /* min_cpu_micros */
+               0,          /* min_params */
+               0,          /* min_float_ops */
+               0,          /* min_occurrence */
+               0,          /* step */
+               "name",     /* order_by */
+               {".*"},     /* account_type_regexes */
+               {".*"},     /* start_name_regexes */
+               {},         /* trim_name_regexes */
+               {".*"}, {}, /* hide_name_regexes */
+               false,      /* account_displayed_op_only */
+               {"micros"}, /* select */
+               {"none"},   /* output_type */
+               {});
+  return opts;
+}
+
+template <typename T>
+const T* ExtractNode(const T& pb, const string& name) {
+  if (pb.name() == name) {
+    return &pb;
+  }
+  for (const T& c : pb.children()) {
+    const T* ret = ExtractNode(c, name);
+    if (ret) return ret;
+  }
+  return nullptr;
+}
+
+TEST_F(ProfilerTest, Basics) {
+  SessionOptions options;
+  options.config.set_allow_soft_placement(true);
+  std::unique_ptr<Session> session(NewSession(options));
+  GraphDef def = CreateGraphDef();
+  if (options.target.empty()) {
+    graph::SetDefaultDevice("/gpu:0", &def);
+  }
+
+  TF_CHECK_OK(session->Create(def));
+
+  Tensor x(DT_FLOAT, TensorShape({2, 1}));
+  auto x_flat = x.flat<float>();
+  x_flat.setRandom();
+  Eigen::Tensor<float, 0, Eigen::RowMajor> inv_norm =
+      x_flat.square().sum().sqrt().inverse();
+  x_flat = x_flat * inv_norm();
+
+  std::vector<Tensor> outputs;
+  RunOptions run_options;
+  run_options.set_trace_level(RunOptions::FULL_TRACE);
+  RunMetadata run_metadata;
+  outputs.clear();
+
+  Profiler profiler(def);
+  for (int i = 0; i < 2; ++i) {
+    TF_CHECK_OK(session->Run(run_options, {{"x", x}}, {"y:0", "y_normalized:0"},
+                             {}, &outputs, &run_metadata));
+    profiler.AddStep(i, run_metadata);
+    CHECK_EQ(size_t{2}, outputs.size());
+  }
+
+  std::vector<DeviceAttributes> resp;
+  TF_CHECK_OK(session->ListDevices(&resp));
+  bool has_gpu = false;
+  for (const auto& dev : resp) {
+    if (dev.device_type() == "GPU") {
+      has_gpu = true;
+    }
+  }
+
+  GraphNodeProto ret = profiler.ProfileNameScope(Default());
+  const GraphNodeProto* matmul = ExtractNode(ret, "y");
+  EXPECT_TRUE(matmul);
+  EXPECT_GT(matmul->exec_micros(), 0);
+  if (has_gpu) {
+    EXPECT_GT(matmul->accelerator_exec_micros(), 0);
+  } else {
+    EXPECT_EQ(matmul->accelerator_exec_micros(), 0);
+  }
+  const GraphNodeProto* square = ExtractNode(ret, "Square");
+  EXPECT_TRUE(square);
+  EXPECT_GT(square->exec_micros(), 0);
+  if (has_gpu) {
+    EXPECT_GT(square->accelerator_exec_micros(), 0);
+  } else {
+    EXPECT_EQ(square->accelerator_exec_micros(), 0);
+  }
+
+  Options opts2 = Default();
+  opts2.output_type = "timeline";
+  string timeline_file = io::JoinPath(testing::TmpDir(), "timeline");
+  opts2.output_options["outfile"] = timeline_file;
+  GraphNodeProto ret2 = profiler.ProfileGraph(opts2);
+  string s;
+  TF_CHECK_OK(ReadFileToString(Env::Default(), timeline_file + "_0", &s));
+  EXPECT_TRUE(s.find("Square") != s.npos);
+
+  MultiGraphNodeProto ret3 = profiler.ProfileOperations(Default());
+  const MultiGraphNodeProto* matmul2 = ExtractNode(ret3, "MatMul");
+  EXPECT_TRUE(matmul2);
+  EXPECT_GT(matmul2->exec_micros(), 0);
+  if (has_gpu) {
+    EXPECT_GT(matmul2->accelerator_exec_micros(), 0);
+  } else {
+    EXPECT_EQ(matmul2->accelerator_exec_micros(), 0);
+  }
+
+  TF_CHECK_OK(session->Close());
+}
+
+}  // namespace tfprof
+}  // namespace tensorflow

tensorflow/contrib/cmake/tf_core_profiler.cmake

@@ -17,6 +17,8 @@
 ########################################################
 file(GLOB_RECURSE tf_core_profiler_srcs
     "${tensorflow_source_dir}/tensorflow/core/profiler/*.proto"
+    "${tensorflow_source_dir}/tensorflow/core/profiler/tfprof_options.h"
+    "${tensorflow_source_dir}/tensorflow/core/profiler/tfprof_options.cc"
     "${tensorflow_source_dir}/tensorflow/core/profiler/internal/*.h"
     "${tensorflow_source_dir}/tensorflow/core/profiler/internal/*.cc"
     "${tensorflow_source_dir}/tensorflow/core/profiler/internal/advisor/*.h"

tensorflow/core/profiler/BUILD

@@ -34,7 +34,7 @@ tf_cc_binary(
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
-        "//tensorflow/core/profiler/internal:tfprof_options",
+        "//tensorflow/core/profiler:tfprof_options",
         "//tensorflow/core/profiler/internal:tfprof_stats",
         "//tensorflow/core/profiler/internal:tfprof_utils",
         "//tensorflow/core/profiler/internal/advisor:tfprof_advisor",
@@ -42,6 +42,17 @@ tf_cc_binary(
     ],
 )
 
+cc_library(
+    name = "tfprof_options",
+    srcs = ["tfprof_options.cc"],
+    hdrs = ["tfprof_options.h"],
+    deps = [
+        "//tensorflow/core:framework_headers_lib",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/profiler:protos_all_cc",
+    ],
+)
+
 tf_proto_library(
     name = "protos_all",
     srcs = glob(["**/*.proto"]),

tensorflow/core/profiler/internal/BUILD

@@ -16,7 +16,6 @@ cc_library(
         ":tfprof_graph",
         ":tfprof_node",
         ":tfprof_op",
-        ":tfprof_options",
         ":tfprof_scope",
         ":tfprof_show",
         ":tfprof_timeline",
@@ -26,6 +25,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -47,12 +47,12 @@ cc_library(
     srcs = ["tfprof_node.cc"],
     hdrs = ["tfprof_node.h"],
     deps = [
-        ":tfprof_options",
         ":tfprof_utils",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -63,7 +63,6 @@ cc_library(
     deps = [
         ":tfprof_constants",
         ":tfprof_node",
-        ":tfprof_options",
         ":tfprof_show",
         ":tfprof_tensor",
         ":tfprof_utils",
@@ -74,6 +73,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -84,7 +84,6 @@ cc_library(
     deps = [
         ":tfprof_constants",
         ":tfprof_node",
-        ":tfprof_options",
         ":tfprof_show_multi",
         ":tfprof_tensor",
         ":tfprof_utils",
@@ -94,6 +93,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -104,7 +104,6 @@ cc_library(
     deps = [
         ":tfprof_constants",
         ":tfprof_node",
-        ":tfprof_options",
         ":tfprof_show_multi",
         ":tfprof_timeline",
         ":tfprof_utils",
@@ -116,6 +115,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -126,7 +126,6 @@ cc_library(
     deps = [
         ":tfprof_constants",
         ":tfprof_node",
-        ":tfprof_options",
         ":tfprof_show",
         ":tfprof_tensor",
         ":tfprof_utils",
@@ -135,6 +134,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -145,11 +145,11 @@ cc_library(
     deps = [
         ":tfprof_constants",
         ":tfprof_node",
-        ":tfprof_options",
         ":tfprof_utils",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -161,7 +161,6 @@ cc_library(
         ":tfprof_constants",
         ":tfprof_node",
         ":tfprof_node_show",
-        ":tfprof_options",
         ":tfprof_tensor",
         ":tfprof_timeline",
         ":tfprof_utils",
@@ -170,6 +169,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -181,7 +181,6 @@ cc_library(
         ":tfprof_constants",
         ":tfprof_node",
         ":tfprof_node_show",
-        ":tfprof_options",
         ":tfprof_scope",
         ":tfprof_show",
         ":tfprof_tensor",
@@ -192,6 +191,7 @@ cc_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -209,7 +209,6 @@ tf_cc_test(
     ],
     deps = [
         ":tfprof_constants",
-        ":tfprof_options",
         ":tfprof_stats",
         ":tfprof_tf_testlib",
         ":tfprof_utils",
@@ -218,6 +217,7 @@ tf_cc_test(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -231,7 +231,6 @@ tf_cc_test(
     ],
     deps = [
         ":tfprof_constants",
-        ":tfprof_options",
         ":tfprof_stats",
         ":tfprof_tf_testlib",
         ":tfprof_utils",
@@ -241,6 +240,7 @@ tf_cc_test(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -250,21 +250,10 @@ cc_library(
     hdrs = ["tfprof_utils.h"],
     copts = if_not_windows(["-Wno-sign-compare"]),
     deps = [
-        ":tfprof_options",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:regexp_internal",
-    ],
-)
-
-cc_library(
-    name = "tfprof_options",
-    srcs = ["tfprof_options.cc"],
-    hdrs = ["tfprof_options.h"],
-    deps = [
-        "//tensorflow/core:framework_headers_lib",
-        "//tensorflow/core:lib",
-        "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -279,13 +268,13 @@ cc_library(
     srcs = ["print_model_analysis.cc"],
     hdrs = ["print_model_analysis.h"],
     deps = [
-        ":tfprof_options",
         ":tfprof_stats",
         "//tensorflow/c:checkpoint_reader",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
         "//tensorflow/core/profiler/internal/advisor:tfprof_advisor",
     ],
     alwayslink = 1,
@@ -305,7 +294,6 @@ tf_cc_test(
     ],
     deps = [
         ":tfprof_constants",
-        ":tfprof_options",
         ":tfprof_stats",
         ":tfprof_tf_testlib",
         ":tfprof_utils",
@@ -314,6 +302,7 @@ tf_cc_test(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 
@@ -340,7 +329,6 @@ tf_cc_test(
         "testdata/graph.pbtxt",
     ],
     deps = [
-        ":tfprof_options",
         ":tfprof_stats",
         ":tfprof_tf_testlib",
         ":tfprof_utils",
@@ -349,6 +337,7 @@ tf_cc_test(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core/profiler:protos_all_cc",
+        "//tensorflow/core/profiler:tfprof_options",
     ],
 )
 

tensorflow/core/profiler/internal/print_model_analysis.cc

@@ -22,13 +22,13 @@ limitations under the License.
 #include "tensorflow/c/checkpoint_reader.h"
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
-#include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/profiler/internal/advisor/tfprof_advisor.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_stats.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_options.pb.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
 
 namespace tensorflow {
 namespace tfprof {

tensorflow/core/profiler/internal/tfprof_code.h

@@ -28,12 +28,12 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_show_multi.h"
 #include "tensorflow/core/profiler/internal/tfprof_timeline.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/profile.pb.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_graph.h

@@ -30,9 +30,9 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_show.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_node.h

@@ -31,8 +31,8 @@ limitations under the License.
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/regexp.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 
 namespace tensorflow {
 namespace tfprof {

tensorflow/core/profiler/internal/tfprof_node_show.h

@@ -32,8 +32,8 @@ limitations under the License.
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_op.h

@@ -29,9 +29,9 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_show_multi.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_scope.h

@@ -29,9 +29,9 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_show.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_show.h

@@ -29,10 +29,10 @@ limitations under the License.
 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
 #include "tensorflow/core/profiler/internal/tfprof_node_show.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_tensor.h"
 #include "tensorflow/core/profiler/internal/tfprof_timeline.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_show_multi.h

@@ -29,11 +29,11 @@ limitations under the License.
 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
 #include "tensorflow/core/profiler/internal/tfprof_node_show.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_show.h"
 #include "tensorflow/core/profiler/internal/tfprof_tensor.h"
 #include "tensorflow/core/profiler/internal/tfprof_timeline.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 
 namespace tensorflow {

tensorflow/core/profiler/internal/tfprof_show_test.cc

@@ -23,9 +23,9 @@ limitations under the License.
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 

tensorflow/core/profiler/internal/tfprof_stats.cc

@@ -282,7 +282,7 @@ void TFStats::AddRunMeta(int64 step, std::unique_ptr<RunMetadata> run_meta) {
   }
 }
 
-void TFStats::WriteProfile(const string& filename) {
+void TFStats::SerializeToString(string* content) {
   ProfileProto profile;
   for (const auto& entry : id_to_string_) {
     (*profile.mutable_id_to_string())[entry.first] = entry.second;
@@ -299,8 +299,13 @@ void TFStats::WriteProfile(const string& filename) {
   for (int64 s : steps_) {
     profile.add_steps(s);
   }
-  Status s =
-      WriteStringToFile(Env::Default(), filename, profile.SerializeAsString());
+  *content = profile.SerializeAsString();
+}
+
+void TFStats::WriteProfile(const string& filename) {
+  string content;
+  SerializeToString(&content);
+  Status s = WriteStringToFile(Env::Default(), filename, content);
   if (!s.ok()) {
     fprintf(stderr, "%s\n", s.ToString().c_str());
   }

tensorflow/core/profiler/internal/tfprof_stats.h

@@ -34,17 +34,17 @@ limitations under the License.
 #include "tensorflow/core/framework/step_stats.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
-#include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/profiler/internal/tfprof_code.h"
 #include "tensorflow/core/profiler/internal/tfprof_graph.h"
 #include "tensorflow/core/profiler/internal/tfprof_node.h"
 #include "tensorflow/core/profiler/internal/tfprof_op.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_scope.h"
 #include "tensorflow/core/profiler/internal/tfprof_show.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
 
 namespace tensorflow {
 namespace tfprof {
@@ -92,6 +92,7 @@ class TFStats {
   // and code traces.
   void AddOpLogProto(std::unique_ptr<OpLogProto> op_log);
 
+  void SerializeToString(string* content);
   void WriteProfile(const string& filename);
 
   // For test purpose only.

tensorflow/core/profiler/internal/tfprof_stats_test.cc

@@ -24,9 +24,9 @@ limitations under the License.
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 

tensorflow/core/profiler/internal/tfprof_tensor_test.cc

@@ -18,10 +18,10 @@ limitations under the License.
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_stats.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 

tensorflow/core/profiler/internal/tfprof_timeline_test.cc

@@ -23,12 +23,12 @@ limitations under the License.
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/test.h"
-#include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/profiler/internal/tfprof_constants.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/profiler/tfprof_output.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
 
 namespace tensorflow {
 namespace tfprof {

tensorflow/core/profiler/internal/tfprof_utils.h

@@ -22,8 +22,8 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 #include "tensorflow/core/protobuf/config.pb.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 
 namespace tensorflow {
 namespace tfprof {

tensorflow/core/profiler/profiler.cc

@@ -31,13 +31,13 @@ limitations under the License.
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/init_main.h"
 #include "tensorflow/core/platform/protobuf.h"
-#include "tensorflow/core/protobuf/config.pb.h"
-#include "tensorflow/core/util/command_line_flags.h"
 #include "tensorflow/core/profiler/internal/advisor/tfprof_advisor.h"
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
 #include "tensorflow/core/profiler/internal/tfprof_stats.h"
 #include "tensorflow/core/profiler/internal/tfprof_utils.h"
 #include "tensorflow/core/profiler/tfprof_log.pb.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/util/command_line_flags.h"
 
 namespace tensorflow {
 namespace tfprof {

tensorflow/core/profiler/tfprof_options.cc

@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/core/profiler/internal/tfprof_options.h"
+#include "tensorflow/core/profiler/tfprof_options.h"
 
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"

tensorflow/python/profiler/model_analyzer.py

@@ -177,8 +177,9 @@ class Profiler(object):
     """Add statistics of a step.
 
     Args:
-      step: int, A step used to identify the RunMetadata. Must be different
-         across different AddStep() calls.
+      step: int, An id used to group one or more different `run_meta` together.
+          When profiling with the profile_xxx APIs, user can use the `step`
+          id in the `options` to profile these `run_meta` together.
       run_meta: RunMetadata proto that contains statistics of a session run.
     """
     # pylint: disable=protected-access